ollie/toolsrv/server9p.go

// server9p.go - 9P server implementation for toolsrv
//
// Namespace:
//   /register      rdwr: register agent (aname + cwd + sig → token)
//   /ctl           write: load/unload tool
//   /tools         read: list loaded tools
//   /proc/new      rdwr: token + tool + args → result (blocking)
//   /proc/new.bg   write: token + tool + args → pid (detached)
//   /proc/{pid}/out   read: output stream
//   /proc/{pid}/wait  read: blocks until exit, returns code
//   /proc/{pid}/stat  read: status
//   /proc/{pid}/ctl   write: signal, dismiss
//   /info          read: platform info
package toolsrv

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"net"
	"path"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"sync"

	"9fans.net/go/plan9"
)

// Server9P implements a 9P server for tool execution.
type Server9P struct {
	mu sync.RWMutex

	secret   []byte            // shared secret for auth
	registry *Registry         // tool registry
	sessID   string            // session ID for registry scoping
	yolo     bool              // skip sandbox

	// Agent registration: token → agent info
	agents map[string]*AgentInfo

	// Process management
	fs *FS9P
}

// AgentInfo holds registered agent data.
type AgentInfo struct {
	ID    string // immutable agent ID (UUID)
	CWD   string // working directory
	Token string // auth token
}

// NewServer9P creates a new 9P tool server.
func NewServer9P(secret []byte) *Server9P {
	return &Server9P{
		secret: secret,
		agents: make(map[string]*AgentInfo),
		fs:     NewFS9P(""), // cwd set per-agent
	}
}

// SetRegistry configures the tool registry and session ID.
func (s *Server9P) SetRegistry(r *Registry, sessID string) {
	s.mu.Lock()
	s.registry = r
	s.sessID = sessID
	s.fs.SetRegistry(r, sessID)
	s.mu.Unlock()
}

// SetYolo enables/disables sandbox bypass.
func (s *Server9P) SetYolo(yolo bool) {
	s.mu.Lock()
	s.yolo = yolo
	s.fs.SetYolo(yolo)
	s.mu.Unlock()
}

// RegisterAgent registers an agent and returns a token.
// Verifies the signature using the shared secret.
func (s *Server9P) RegisterAgent(id, cwd, sig string) (string, error) {
	// Verify signature
	expected := s.computeSig(id, cwd)
	if !hmac.Equal([]byte(sig), []byte(expected)) {
		return "", fmt.Errorf("invalid signature")
	}

	// Generate token
	tokenBytes := make([]byte, 16)
	if _, err := rand.Read(tokenBytes); err != nil {
		return "", err
	}
	token := hex.EncodeToString(tokenBytes)

	// Store agent info
	s.mu.Lock()
	s.agents[token] = &AgentInfo{
		ID:    id,
		CWD:   cwd,
		Token: token,
	}
	s.mu.Unlock()

	return token, nil
}

// GetAgent returns agent info for a token.
func (s *Server9P) GetAgent(token string) (*AgentInfo, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	agent, ok := s.agents[token]
	return agent, ok
}

// computeSig computes HMAC signature for agent registration.
func (s *Server9P) computeSig(id, cwd string) string {
	h := hmac.New(sha256.New, s.secret)
	h.Write([]byte(id + "|" + cwd))
	return hex.EncodeToString(h.Sum(nil))
}

// Serve starts serving 9P on the given listener.
func (s *Server9P) Serve(ln net.Listener) error {
	for {
		conn, err := ln.Accept()
		if err != nil {
			return err
		}
		go s.serveConn(conn)
	}
}

// serveConn handles a single 9P connection.
func (s *Server9P) serveConn(conn net.Conn) {
	defer conn.Close()

	cs := &connState9P{
		srv:  s,
		fids: make(map[uint32]*fid9P),
	}

	for {
		fcall, err := readFcall(conn)
		if err != nil {
			if err != io.EOF {
				// log error
			}
			return
		}

		resp := cs.handle(fcall)
		if err := writeFcall(conn, resp); err != nil {
			return
		}
	}
}

// connState9P holds per-connection state.
type connState9P struct {
	srv  *Server9P
	mu   sync.Mutex
	fids map[uint32]*fid9P
	auth bool // connection authenticated
}

// fid9P represents an open file handle.
type fid9P struct {
	path   string
	qid    plan9.Qid
	opened bool
	mode   uint8

	// For rdwr operations
	wrBuf []byte
	rdBuf []byte
	rdOff int64
}

func (cs *connState9P) handle(fc *plan9.Fcall) *plan9.Fcall {
	resp := &plan9.Fcall{Tag: fc.Tag}

	switch fc.Type {
	case plan9.Tversion:
		resp.Type = plan9.Rversion
		resp.Msize = fc.Msize
		if fc.Msize > 65536 {
			resp.Msize = 65536
		}
		resp.Version = "9P2000"

	case plan9.Tauth:
		// We use shared secret, not per-connection auth
		// Accept any auth request, actual verification happens at /register
		resp.Type = plan9.Rauth
		resp.Aqid = plan9.Qid{Type: plan9.QTAUTH}
		cs.auth = true

	case plan9.Tattach:
		if !cs.auth && fc.Afid != plan9.NOFID {
			// Auth was provided, mark as authenticated
			cs.auth = true
		}
		cs.mu.Lock()
		cs.fids[fc.Fid] = &fid9P{
			path: "/",
			qid:  plan9.Qid{Type: plan9.QTDIR, Path: 0},
		}
		cs.mu.Unlock()
		resp.Type = plan9.Rattach
		resp.Qid = plan9.Qid{Type: plan9.QTDIR, Path: 0}

	case plan9.Twalk:
		resp = cs.handleWalk(fc)

	case plan9.Topen:
		resp = cs.handleOpen(fc)

	case plan9.Tread:
		resp = cs.handleRead(fc)

	case plan9.Twrite:
		resp = cs.handleWrite(fc)

	case plan9.Tclunk:
		cs.mu.Lock()
		delete(cs.fids, fc.Fid)
		cs.mu.Unlock()
		resp.Type = plan9.Rclunk

	case plan9.Tstat:
		resp = cs.handleStat(fc)

	case plan9.Tremove:
		resp.Type = plan9.Rerror
		resp.Ename = "remove not supported"

	case plan9.Twstat:
		resp.Type = plan9.Rerror
		resp.Ename = "wstat not supported"

	default:
		resp.Type = plan9.Rerror
		resp.Ename = fmt.Sprintf("unknown message type %d", fc.Type)
	}

	return resp
}

func (cs *connState9P) handleWalk(fc *plan9.Fcall) *plan9.Fcall {
	resp := &plan9.Fcall{Tag: fc.Tag}

	cs.mu.Lock()
	f, ok := cs.fids[fc.Fid]
	cs.mu.Unlock()
	if !ok {
		resp.Type = plan9.Rerror
		resp.Ename = "unknown fid"
		return resp
	}

	newPath := f.path
	var qids []plan9.Qid

	for _, name := range fc.Wname {
		if name == ".." {
			newPath = path.Dir(newPath)
		} else {
			newPath = path.Join(newPath, name)
		}

		qt, qp, err := cs.srv.qidFor(newPath)
		if err != nil {
			if len(qids) == 0 {
				resp.Type = plan9.Rerror
				resp.Ename = err.Error()
				return resp
			}
			break
		}
		qids = append(qids, plan9.Qid{Type: qt, Path: qp})
	}

	resp.Type = plan9.Rwalk
	resp.Wqid = qids

	if len(qids) == len(fc.Wname) {
		cs.mu.Lock()
		if fc.Newfid != fc.Fid {
			cs.fids[fc.Newfid] = &fid9P{
				path: newPath,
				qid:  qids[len(qids)-1],
			}
		} else if len(qids) > 0 {
			f.path = newPath
			f.qid = qids[len(qids)-1]
		}
		cs.mu.Unlock()
	}

	return resp
}

func (cs *connState9P) handleOpen(fc *plan9.Fcall) *plan9.Fcall {
	resp := &plan9.Fcall{Tag: fc.Tag}

	cs.mu.Lock()
	f, ok := cs.fids[fc.Fid]
	cs.mu.Unlock()
	if !ok {
		resp.Type = plan9.Rerror
		resp.Ename = "unknown fid"
		return resp
	}

	qt, qp, err := cs.srv.qidFor(f.path)
	if err != nil {
		resp.Type = plan9.Rerror
		resp.Ename = err.Error()
		return resp
	}

	f.opened = true
	f.mode = fc.Mode
	f.qid = plan9.Qid{Type: qt, Path: qp}
	f.wrBuf = nil
	f.rdBuf = nil
	f.rdOff = 0

	resp.Type = plan9.Ropen
	resp.Qid = f.qid
	resp.Iounit = 8192

	return resp
}

func (cs *connState9P) handleRead(fc *plan9.Fcall) *plan9.Fcall {
	resp := &plan9.Fcall{Tag: fc.Tag}

	cs.mu.Lock()
	f, ok := cs.fids[fc.Fid]
	cs.mu.Unlock()
	if !ok {
		resp.Type = plan9.Rerror
		resp.Ename = "unknown fid"
		return resp
	}

	if !f.opened {
		resp.Type = plan9.Rerror
		resp.Ename = "not open"
		return resp
	}

	// For rdwr files, check if we have a result buffer
	if f.rdBuf != nil {
		data := cs.readFromBuf(f, fc.Offset, fc.Count)
		resp.Type = plan9.Rread
		resp.Data = data
		return resp
	}

	// Regular read
	data, err := cs.srv.readFile(f.path, int64(fc.Offset), int(fc.Count))
	if err != nil {
		resp.Type = plan9.Rerror
		resp.Ename = err.Error()
		return resp
	}

	resp.Type = plan9.Rread
	resp.Data = data
	return resp
}

func (cs *connState9P) readFromBuf(f *fid9P, offset uint64, count uint32) []byte {
	if int64(offset) >= int64(len(f.rdBuf)) {
		return nil
	}
	end := int64(offset) + int64(count)
	if end > int64(len(f.rdBuf)) {
		end = int64(len(f.rdBuf))
	}
	return f.rdBuf[offset:end]
}

func (cs *connState9P) handleWrite(fc *plan9.Fcall) *plan9.Fcall {
	resp := &plan9.Fcall{Tag: fc.Tag}

	cs.mu.Lock()
	f, ok := cs.fids[fc.Fid]
	cs.mu.Unlock()
	if !ok {
		resp.Type = plan9.Rerror
		resp.Ename = "unknown fid"
		return resp
	}

	if !f.opened {
		resp.Type = plan9.Rerror
		resp.Ename = "not open"
		return resp
	}

	// Handle rdwr files - buffer writes, execute on read
	if isRdwrPath(f.path) {
		f.wrBuf = append(f.wrBuf, fc.Data...)
		
		// Check if we have a complete request (ends with newline or is a clunk)
		if len(f.wrBuf) > 0 && f.wrBuf[len(f.wrBuf)-1] == '\n' {
			result, err := cs.srv.handleRdwr(f.path, string(f.wrBuf))
			if err != nil {
				f.rdBuf = []byte("error: " + err.Error() + "\n")
			} else {
				f.rdBuf = []byte(result)
			}
			f.wrBuf = nil
		}

		resp.Type = plan9.Rwrite
		resp.Count = uint32(len(fc.Data))
		return resp
	}

	// Regular write
	n, err := cs.srv.writeFile(f.path, fc.Data)
	if err != nil {
		resp.Type = plan9.Rerror
		resp.Ename = err.Error()
		return resp
	}

	resp.Type = plan9.Rwrite
	resp.Count = uint32(n)
	return resp
}

func (cs *connState9P) handleStat(fc *plan9.Fcall) *plan9.Fcall {
	resp := &plan9.Fcall{Tag: fc.Tag}

	cs.mu.Lock()
	f, ok := cs.fids[fc.Fid]
	cs.mu.Unlock()
	if !ok {
		resp.Type = plan9.Rerror
		resp.Ename = "unknown fid"
		return resp
	}

	dir, err := cs.srv.statFile(f.path)
	if err != nil {
		resp.Type = plan9.Rerror
		resp.Ename = err.Error()
		return resp
	}

	resp.Type = plan9.Rstat
	resp.Stat = dir
	return resp
}

// --- File system operations ---

func (s *Server9P) qidFor(p string) (uint8, uint64, error) {
	switch p {
	case "/":
		return plan9.QTDIR, 0, nil
	case "/register", "/ctl", "/tools", "/info":
		return plan9.QTFILE, pathHash(p), nil
	case "/proc":
		return plan9.QTDIR, pathHash(p), nil
	case "/proc/new", "/proc/new.bg":
		return plan9.QTFILE, pathHash(p), nil
	}

	if strings.HasPrefix(p, "/proc/") {
		parts := strings.Split(p[6:], "/")
		if len(parts) == 1 {
			// /proc/{pid}
			return plan9.QTDIR, pathHash(p), nil
		}
		if len(parts) == 2 {
			// /proc/{pid}/{file}
			switch parts[1] {
			case "out", "wait", "stat", "ctl":
				return plan9.QTFILE, pathHash(p), nil
			}
		}
	}

	return 0, 0, fmt.Errorf("not found: %s", p)
}

func pathHash(p string) uint64 {
	var h uint64 = 5381
	for _, c := range p {
		h = h*33 + uint64(c)
	}
	return h
}

func isRdwrPath(p string) bool {
	return p == "/register" || p == "/proc/new"
}

func (s *Server9P) readFile(p string, offset int64, count int) ([]byte, error) {
	var data []byte

	switch p {
	case "/":
		data = []byte("ctl\ninfo\nproc\nregister\ntools\n")

	case "/tools":
		data = []byte(s.fs.HandleToolsRead())

	case "/info":
		data = []byte(s.hostInfo())

	case "/proc":
		var entries []string
		entries = append(entries, "new", "new.bg")
		for _, pid := range s.fs.ListProcs() {
			entries = append(entries, strconv.Itoa(pid))
		}
		sort.Strings(entries)
		data = []byte(strings.Join(entries, "\n") + "\n")

	default:
		if strings.HasPrefix(p, "/proc/") {
			return s.readProcFile(p, offset, count)
		}
		return nil, fmt.Errorf("cannot read: %s", p)
	}

	if offset >= int64(len(data)) {
		return nil, nil
	}
	end := offset + int64(count)
	if end > int64(len(data)) {
		end = int64(len(data))
	}
	return data[offset:end], nil
}

func (s *Server9P) readProcFile(p string, offset int64, count int) ([]byte, error) {
	parts := strings.Split(p[6:], "/") // strip "/proc/"
	if len(parts) < 2 {
		return nil, fmt.Errorf("invalid proc path")
	}

	pid, err := strconv.Atoi(parts[0])
	if err != nil {
		return nil, fmt.Errorf("invalid pid")
	}

	proc := s.fs.GetProc(pid)
	if proc == nil {
		return nil, fmt.Errorf("process not found")
	}

	var data []byte
	switch parts[1] {
	case "out":
		data = []byte(proc.Output())
	case "wait":
		exitCode := proc.Wait()
		data = []byte(fmt.Sprintf("%d\n", exitCode))
	case "stat":
		data = []byte(proc.Stat())
	default:
		return nil, fmt.Errorf("not found: %s", parts[1])
	}

	if offset >= int64(len(data)) {
		return nil, nil
	}
	end := offset + int64(count)
	if end > int64(len(data)) {
		end = int64(len(data))
	}
	return data[offset:end], nil
}

func (s *Server9P) writeFile(p string, data []byte) (int, error) {
	input := strings.TrimSpace(string(data))

	switch p {
	case "/ctl":
		if err := s.fs.HandleCtl(input); err != nil {
			return 0, err
		}
		return len(data), nil

	case "/tools":
		if err := s.fs.HandleToolsWrite(input); err != nil {
			return 0, err
		}
		return len(data), nil

	case "/proc/new.bg":
		return s.handleProcNewBg(input)

	default:
		if strings.HasPrefix(p, "/proc/") && strings.HasSuffix(p, "/ctl") {
			return s.handleProcCtl(p, input)
		}
		return 0, fmt.Errorf("cannot write: %s", p)
	}
}

func (s *Server9P) handleRdwr(p string, input string) (string, error) {
	switch p {
	case "/register":
		return s.handleRegister(input)
	case "/proc/new":
		return s.handleProcNew(input)
	default:
		return "", fmt.Errorf("not rdwr: %s", p)
	}
}

func (s *Server9P) handleRegister(input string) (string, error) {
	args := parseKV(input)
	id := args["aname"]
	cwd := args["cwd"]
	sig := args["sig"]

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

	token, err := s.RegisterAgent(id, cwd, sig)
	if err != nil {
		return "", err
	}
	return token + "\n", nil
}

func (s *Server9P) handleProcNew(input string) (string, error) {
	args := parseKV(input)
	token := args["token"]

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

	// Set cwd for this execution
	s.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")
	}

	return s.fs.HandleProcNew(nil, payload.String()) // TODO: pass context
}

func (s *Server9P) handleProcNewBg(input string) (int, error) {
	args := parseKV(input)
	token := args["token"]

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

	// Set cwd for this execution
	s.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 := s.fs.HandleProcNewBg(nil, payload.String()) // TODO: pass context
	if err != nil {
		return 0, err
	}
	return len(result), nil
}

func (s *Server9P) handleProcCtl(p string, input string) (int, error) {
	// Extract pid from path: /proc/{pid}/ctl
	parts := strings.Split(p[6:], "/")
	if len(parts) < 2 {
		return 0, fmt.Errorf("invalid path")
	}
	pid, err := strconv.Atoi(parts[0])
	if err != nil {
		return 0, fmt.Errorf("invalid pid")
	}

	if err := s.fs.HandleProcCtl(pid, input); err != nil {
		return 0, err
	}
	return len(input), nil
}

func (s *Server9P) statFile(p string) ([]byte, error) {
	qt, qp, err := s.qidFor(p)
	if err != nil {
		return nil, err
	}

	name := path.Base(p)
	if p == "/" {
		name = "/"
	}

	mode := plan9.Perm(0644)
	if qt == plan9.QTDIR {
		mode = plan9.Perm(plan9.DMDIR | 0755)
	}

	dir := plan9.Dir{
		Name: name,
		Qid:  plan9.Qid{Type: qt, Path: qp},
		Mode: mode,
		Uid:  "toolsrv",
		Gid:  "toolsrv",
	}
	data, err := dir.Bytes()
	return data, err
}

func (s *Server9P) hostInfo() string {
	return fmt.Sprintf("platform=%s\narch=%s\n", runtime.GOOS, runtime.GOARCH)
}

// parseKV parses key=value lines.
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
		v = strings.ReplaceAll(v, "\\n", "\n")
		v = strings.ReplaceAll(v, "\\\\", "\\")
		m[k] = v
	}
	return m
}

// --- 9P message I/O ---

func readFcall(r io.Reader) (*plan9.Fcall, error) {
	return plan9.ReadFcall(r)
}

func writeFcall(w io.Writer, fc *plan9.Fcall) error {
	return plan9.WriteFcall(w, fc)
}