toolsrv: add 9P server implementation
New files: - server9p.go: 9P server for tool execution - Agent registration with HMAC-signed tokens - Token → cwd mapping for secure execution context - Full 9P protocol handling (version, attach, walk, open, read, write, stat, clunk) - Namespace: /register, /ctl, /tools, /proc/*, /info - auth9p.go: Authentication helpers - GenerateSecret(): 32-byte random secret - ComputeRegistrationSig(): HMAC-SHA256 signature for registration - server9p_test.go: Unit tests for auth and server Security model: - olliesrv spawns toolsrv with shared secret (env or file) - Agent registration: olliesrv signs (agentID, cwd) → token - Every tool call includes token; toolsrv verifies and looks up cwd - Prevents unauthorized cwd manipulation from untrusted actors (LLM)
This commit is contained in:
Levi Neely
parent
1cd158926a
commit
d37404ebd3
|
|
@ -0,0 +1,26 @@
|
|||
// auth9p.go - Authentication helpers for 9P toolsrv
|
||||
package toolsrv
|
||||
|
||||
import (
|
||||
"crypto/hmac"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
)
|
||||
|
||||
// GenerateSecret generates a random shared secret for toolsrv auth.
|
||||
func GenerateSecret() ([]byte, error) {
|
||||
secret := make([]byte, 32)
|
||||
if _, err := rand.Read(secret); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return secret, nil
|
||||
}
|
||||
|
||||
// ComputeRegistrationSig computes the HMAC signature for agent registration.
|
||||
// This is used by olliesrv to sign agent registration requests.
|
||||
func ComputeRegistrationSig(secret []byte, agentID, cwd string) string {
|
||||
h := hmac.New(sha256.New, secret)
|
||||
h.Write([]byte(agentID + "|" + cwd))
|
||||
return hex.EncodeToString(h.Sum(nil))
|
||||
}
|
||||
|
|
@ -0,0 +1,778 @@
|
|||
// 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)
|
||||
}
|
||||
|
|
@ -0,0 +1,203 @@
|
|||
package toolsrv
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestServer9P_RegisterAgent(t *testing.T) {
|
||||
secret := []byte("test-secret-key-1234567890123456")
|
||||
srv := NewServer9P(secret)
|
||||
|
||||
agentID := "agent-uuid-123"
|
||||
cwd := "/home/user/project"
|
||||
sig := ComputeRegistrationSig(secret, agentID, cwd)
|
||||
|
||||
// Register should succeed with valid signature
|
||||
token, err := srv.RegisterAgent(agentID, cwd, sig)
|
||||
if err != nil {
|
||||
t.Fatalf("RegisterAgent failed: %v", err)
|
||||
}
|
||||
if token == "" {
|
||||
t.Fatal("RegisterAgent returned empty token")
|
||||
}
|
||||
|
||||
// Should be able to look up the agent
|
||||
agent, ok := srv.GetAgent(token)
|
||||
if !ok {
|
||||
t.Fatal("GetAgent returned false")
|
||||
}
|
||||
if agent.ID != agentID {
|
||||
t.Errorf("agent.ID = %q, want %q", agent.ID, agentID)
|
||||
}
|
||||
if agent.CWD != cwd {
|
||||
t.Errorf("agent.CWD = %q, want %q", agent.CWD, cwd)
|
||||
}
|
||||
}
|
||||
|
||||
func TestServer9P_RegisterAgentBadSig(t *testing.T) {
|
||||
secret := []byte("test-secret-key-1234567890123456")
|
||||
srv := NewServer9P(secret)
|
||||
|
||||
agentID := "agent-uuid-123"
|
||||
cwd := "/home/user/project"
|
||||
badSig := "invalid-signature"
|
||||
|
||||
// Register should fail with invalid signature
|
||||
_, err := srv.RegisterAgent(agentID, cwd, badSig)
|
||||
if err == nil {
|
||||
t.Fatal("RegisterAgent should fail with bad signature")
|
||||
}
|
||||
}
|
||||
|
||||
func TestServer9P_RegisterAgentWrongSecret(t *testing.T) {
|
||||
secret := []byte("test-secret-key-1234567890123456")
|
||||
wrongSecret := []byte("wrong-secret-key-abcdefghijklmnop")
|
||||
srv := NewServer9P(secret)
|
||||
|
||||
agentID := "agent-uuid-123"
|
||||
cwd := "/home/user/project"
|
||||
// Sign with wrong secret
|
||||
sig := ComputeRegistrationSig(wrongSecret, agentID, cwd)
|
||||
|
||||
// Register should fail
|
||||
_, err := srv.RegisterAgent(agentID, cwd, sig)
|
||||
if err == nil {
|
||||
t.Fatal("RegisterAgent should fail with wrong secret")
|
||||
}
|
||||
}
|
||||
|
||||
func TestServer9P_GetAgentNotFound(t *testing.T) {
|
||||
secret := []byte("test-secret-key-1234567890123456")
|
||||
srv := NewServer9P(secret)
|
||||
|
||||
_, ok := srv.GetAgent("nonexistent-token")
|
||||
if ok {
|
||||
t.Fatal("GetAgent should return false for nonexistent token")
|
||||
}
|
||||
}
|
||||
|
||||
func TestServer9P_QidFor(t *testing.T) {
|
||||
secret := []byte("test-secret-key-1234567890123456")
|
||||
srv := NewServer9P(secret)
|
||||
|
||||
tests := []struct {
|
||||
path string
|
||||
wantDir bool
|
||||
wantErr bool
|
||||
}{
|
||||
{"/", true, false},
|
||||
{"/register", false, false},
|
||||
{"/ctl", false, false},
|
||||
{"/tools", false, false},
|
||||
{"/info", false, false},
|
||||
{"/proc", true, false},
|
||||
{"/proc/new", false, false},
|
||||
{"/proc/new.bg", false, false},
|
||||
{"/proc/123", true, false},
|
||||
{"/proc/123/out", false, false},
|
||||
{"/proc/123/wait", false, false},
|
||||
{"/proc/123/stat", false, false},
|
||||
{"/proc/123/ctl", false, false},
|
||||
{"/nonexistent", false, true},
|
||||
{"/proc/123/nonexistent", false, true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
qt, _, err := srv.qidFor(tt.path)
|
||||
if tt.wantErr {
|
||||
if err == nil {
|
||||
t.Errorf("qidFor(%q) should error", tt.path)
|
||||
}
|
||||
continue
|
||||
}
|
||||
if err != nil {
|
||||
t.Errorf("qidFor(%q) error: %v", tt.path, err)
|
||||
continue
|
||||
}
|
||||
isDir := qt == 0x80 // QTDIR
|
||||
if isDir != tt.wantDir {
|
||||
t.Errorf("qidFor(%q) isDir = %v, want %v", tt.path, isDir, tt.wantDir)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseKV(t *testing.T) {
|
||||
tests := []struct {
|
||||
input string
|
||||
want map[string]string
|
||||
}{
|
||||
{
|
||||
input: "key=value\n",
|
||||
want: map[string]string{"key": "value"},
|
||||
},
|
||||
{
|
||||
input: "a=1\nb=2\nc=3\n",
|
||||
want: map[string]string{"a": "1", "b": "2", "c": "3"},
|
||||
},
|
||||
{
|
||||
input: "path=/home/user/project\n",
|
||||
want: map[string]string{"path": "/home/user/project"},
|
||||
},
|
||||
{
|
||||
input: "multi=line\\none\\ntwo\n",
|
||||
want: map[string]string{"multi": "line\none\ntwo"},
|
||||
},
|
||||
{
|
||||
input: "empty=\n",
|
||||
want: map[string]string{"empty": ""},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
got := parseKV(tt.input)
|
||||
for k, v := range tt.want {
|
||||
if got[k] != v {
|
||||
t.Errorf("parseKV(%q)[%q] = %q, want %q", tt.input, k, got[k], v)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestComputeRegistrationSig(t *testing.T) {
|
||||
secret := []byte("test-secret")
|
||||
|
||||
sig1 := ComputeRegistrationSig(secret, "agent1", "/path1")
|
||||
sig2 := ComputeRegistrationSig(secret, "agent1", "/path1")
|
||||
sig3 := ComputeRegistrationSig(secret, "agent1", "/path2")
|
||||
sig4 := ComputeRegistrationSig(secret, "agent2", "/path1")
|
||||
|
||||
// Same inputs should produce same signature
|
||||
if sig1 != sig2 {
|
||||
t.Error("same inputs should produce same signature")
|
||||
}
|
||||
|
||||
// Different cwd should produce different signature
|
||||
if sig1 == sig3 {
|
||||
t.Error("different cwd should produce different signature")
|
||||
}
|
||||
|
||||
// Different agent should produce different signature
|
||||
if sig1 == sig4 {
|
||||
t.Error("different agent should produce different signature")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerateSecret(t *testing.T) {
|
||||
s1, err := GenerateSecret()
|
||||
if err != nil {
|
||||
t.Fatalf("GenerateSecret failed: %v", err)
|
||||
}
|
||||
if len(s1) != 32 {
|
||||
t.Errorf("secret length = %d, want 32", len(s1))
|
||||
}
|
||||
|
||||
s2, err := GenerateSecret()
|
||||
if err != nil {
|
||||
t.Fatalf("GenerateSecret failed: %v", err)
|
||||
}
|
||||
|
||||
// Should be different each time
|
||||
if string(s1) == string(s2) {
|
||||
t.Error("GenerateSecret should produce different values")
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue