ollie/toolsrv/spec9p.go

// spec9p.go - fsedsl-based namespace specification for toolsrv
package toolsrv

import (
	"context"
	"fmt"
	"os"
	"strconv"
	"strings"

	"ollie/fsedsl"
)

// ToolsrvCtx is the context passed through the fsedsl tree.
type ToolsrvCtx struct {
	Server *Server9P
	Agent  *AgentInfo // set when traversing into agent-specific paths
	Proc   *Proc9P    // set when traversing into /proc/{pid}
}

// Type aliases from fsedsl, specialized to ToolsrvCtx.
type (
	FsNodeDecl = fsedsl.FsNodeDecl[ToolsrvCtx]
	Binding    = fsedsl.Binding[ToolsrvCtx]
)

// Specialized fsedsl constructors for ToolsrvCtx.
var (
	DirNode   = fsedsl.DirNode[ToolsrvCtx]
	FileNode  = fsedsl.FileNode[ToolsrvCtx]
	Each      = fsedsl.Each[ToolsrvCtx]
	Doc       = fsedsl.Doc[ToolsrvCtx]
	Read      = fsedsl.Read[ToolsrvCtx]
	Write     = fsedsl.Write[ToolsrvCtx]
	BlockOnce = fsedsl.BlockOnce[ToolsrvCtx]
	Request   = fsedsl.Request[ToolsrvCtx]
)

// ToolsrvSpec returns the fsedsl specification for the toolsrv namespace.
func ToolsrvSpec() FsNodeDecl {
	return DirNode("/",
		FileNode("register", 0666,
			Doc("Register agent: write aname=<id>\\ncwd=<path>\\nsig=<sig>, read token"),
			Request(handleRegister),
		),
		FileNode("ctl", 0222,
			Doc("Control: write 'load <tool>' or 'unload <tool>'"),
			Write(handleCtlWrite),
		),
		FileNode("tools", 0666,
			Doc("Tools: read to list, write tool name to load"),
			Read(handleToolsRead),
			Write(handleToolsWrite),
		),
		FileNode("info", 0444,
			Doc("Host info: platform, arch"),
			Read(handleInfoRead),
		),
		DirNode("proc",
			FileNode("new", 0666,
				Doc("Execute tool: write token + tool + args, read result (blocking)"),
				Request(handleProcNew),
			),
			FileNode("new.bg", 0222,
				Doc("Execute tool in background: write token + tool + args, returns pid"),
				Write(handleProcNewBg),
			),
			Each("{pid}", procBindings,
				FileNode("out", 0444,
					Doc("Process output"),
					Read(handleProcOutRead),
				),
				FileNode("wait", 0444,
					Doc("Block until process exits, returns exit code"),
					BlockOnce(handleProcWait),
				),
				FileNode("stat", 0444,
					Doc("Process status: running/exited, runtime, tool"),
					Read(handleProcStatRead),
				),
				FileNode("ctl", 0222,
					Doc("Process control: write 'signal <N>' or 'dismiss'"),
					Write(handleProcCtlWrite),
				),
			),
		),
	)
}

// --- Handlers ---

func handleRegister(ctx ToolsrvCtx, data []byte) ([]byte, error) {
	args := parseKV(string(data))
	id := args["aname"]
	cwd := args["cwd"]
	sig := args["sig"]

	if id == "" || cwd == "" || sig == "" {
		return nil, fmt.Errorf("missing aname, cwd, or sig")
	}

	token, err := ctx.Server.RegisterAgent(id, cwd, sig)
	if err != nil {
		return nil, err
	}
	return []byte(token + "\n"), nil
}

func handleCtlWrite(ctx ToolsrvCtx, data []byte) error {
	return ctx.Server.fs.HandleCtl(strings.TrimSpace(string(data)))
}

func handleToolsRead(ctx ToolsrvCtx) ([]byte, error) {
	return []byte(ctx.Server.fs.HandleToolsRead()), nil
}

func handleToolsWrite(ctx ToolsrvCtx, data []byte) error {
	return ctx.Server.fs.HandleToolsWrite(strings.TrimSpace(string(data)))
}

func handleInfoRead(ctx ToolsrvCtx) ([]byte, error) {
	return []byte(ctx.Server.HostInfo()), nil
}

func handleProcNew(ctx ToolsrvCtx, data []byte) ([]byte, error) {
	args := parseKV(string(data))
	token := args["token"]

	agent, ok := ctx.Server.GetAgent(token)
	if !ok {
		return nil, fmt.Errorf("invalid token")
	}

	// Set cwd for this execution
	ctx.Server.fs.SetCWD(agent.CWD)

	// Build payload without token
	var payload strings.Builder
	for k, v := range args {
		if k == "token" {
			continue
		}
		payload.WriteString(k)
		payload.WriteString("=")
		payload.WriteString(v)
		payload.WriteString("\n")
	}

	result, err := ctx.Server.fs.HandleProcNew(context.Background(), payload.String())
	if err != nil {
		return nil, err
	}
	return []byte(result), nil
}

func handleProcNewBg(ctx ToolsrvCtx, data []byte) error {
	args := parseKV(string(data))
	token := args["token"]

	agent, ok := ctx.Server.GetAgent(token)
	if !ok {
		return fmt.Errorf("invalid token")
	}

	// Set cwd for this execution
	ctx.Server.fs.SetCWD(agent.CWD)

	// Build payload without token
	var payload strings.Builder
	for k, v := range args {
		if k == "token" {
			continue
		}
		payload.WriteString(k)
		payload.WriteString("=")
		payload.WriteString(v)
		payload.WriteString("\n")
	}

	_, err := ctx.Server.fs.HandleProcNewBg(context.Background(), payload.String())
	return err
}

// procBindings returns bindings for each running process.
func procBindings(ctx ToolsrvCtx) ([]Binding, error) {
	pids := ctx.Server.fs.ListProcs()
	var bindings []Binding

	for _, pid := range pids {
		proc := ctx.Server.fs.GetProc(pid)
		if proc == nil {
			continue
		}
		p := proc // capture for closure
		bindings = append(bindings, Binding{
			Name: strconv.Itoa(pid),
			Applier: func(c ToolsrvCtx) ToolsrvCtx {
				c.Proc = p
				return c
			},
		})
	}

	return bindings, nil
}

func handleProcOutRead(ctx ToolsrvCtx) ([]byte, error) {
	if ctx.Proc == nil {
		return nil, fmt.Errorf("no process context")
	}
	return []byte(ctx.Proc.Output()), nil
}

func handleProcWait(ctx ToolsrvCtx, cctx context.Context, base string) ([]byte, string, error) {
	if ctx.Proc == nil {
		return nil, "", fmt.Errorf("no process context")
	}
	exitCode := ctx.Proc.Wait()
	return []byte(fmt.Sprintf("%d\n", exitCode)), "", nil
}

func handleProcStatRead(ctx ToolsrvCtx) ([]byte, error) {
	if ctx.Proc == nil {
		return nil, fmt.Errorf("no process context")
	}
	return []byte(ctx.Proc.Stat()), nil
}

func handleProcCtlWrite(ctx ToolsrvCtx, data []byte) error {
	if ctx.Proc == nil {
		return fmt.Errorf("no process context")
	}
	return ctx.Server.fs.HandleProcCtl(ctx.Proc.PID, strings.TrimSpace(string(data)))
}

// Mode constants for convenience
const (
	ModeReadOnly  os.FileMode = 0444
	ModeWriteOnly os.FileMode = 0222
	ModeReadWrite os.FileMode = 0666
)

// parseKV parses key=value lines into a map.
func parseKV(input string) map[string]string {
	m := make(map[string]string)
	for _, line := range strings.Split(input, "\n") {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		idx := strings.Index(line, "=")
		if idx < 0 {
			continue
		}
		k := line[:idx]
		v := line[idx+1:]
		// Unescape \n and \\
		v = strings.ReplaceAll(v, "\\n", "\n")
		v = strings.ReplaceAll(v, "\\\\", "\\")
		m[k] = v
	}
	return m
}