ollie/toolsrv/client9p.go

// client9p.go - 9P client for toolsrv.
package toolsrv

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"syscall"

	"9fans.net/go/plan9"
	p9client "9fans.net/go/plan9/client"
)

// Conn is a 9P connection to a toolsrv.
type Conn struct {
	fsys           *p9client.Fsys
	conn           *p9client.Conn
	token          string // session token from auth
	secret         string // secret used for auth (for reconnect)
	onToolsChanged func() // callback for tool changes (client must poll)
}

// Dial connects to a toolsrv at the given Unix socket path and authenticates.
// If secret is empty, a random one is generated (first connection).
// Returns the connection and the secret used (caller should save for reconnect).
func Dial(socketPath string, secret string) (*Conn, error) {
	conn, err := p9client.Dial("unix", socketPath)
	if err != nil {
		return nil, err
	}

	// Generate secret if not provided
	if secret == "" {
		secret = randomSecret()
	}

	// Authenticate via Tauth
	afid, err := conn.Auth("agent", "")
	if err != nil {
		conn.Close()
		return nil, fmt.Errorf("auth: %w", err)
	}

	// Write secret to auth fid
	if _, err := afid.Write([]byte(secret)); err != nil {
		afid.Close()
		conn.Close()
		return nil, fmt.Errorf("write secret: %w", err)
	}

	// Read token from auth fid (use ReadAt to read from offset 0,
	// since Write advances the file offset)
	tokenBuf := make([]byte, 1024)
	n, err := afid.ReadAt(tokenBuf, 0)
	if err != nil && err != io.EOF {
		afid.Close()
		conn.Close()
		return nil, fmt.Errorf("read token: %w", err)
	}
	token := string(tokenBuf[:n])
	// Trim newline
	if len(token) > 0 && token[len(token)-1] == '\n' {
		token = token[:len(token)-1]
	}

	// Attach using the authenticated afid
	fsys, err := conn.Attach(afid, "agent", "")
	if err != nil {
		afid.Close()
		conn.Close()
		return nil, fmt.Errorf("attach: %w", err)
	}

	afid.Close()

	return &Conn{
		fsys:   fsys,
		conn:   conn,
		token:  token,
		secret: secret,
	}, nil
}

// Secret returns the secret used for authentication.
// Save this for reconnecting to the same toolsrv instance.
func (c *Conn) Secret() string {
	return c.secret
}

// Token returns the session token.
func (c *Conn) Token() string {
	return c.token
}

// ListTools returns the list of loaded tools.
func (c *Conn) ListTools() ([]ToolInfo, error) {
	fid, err := c.fsys.Open("tools", plan9.OREAD)
	if err != nil {
		return nil, err
	}
	defer fid.Close()

	data, err := io.ReadAll(fid)
	if err != nil {
		return nil, err
	}

	var tools []ToolInfo
	if err := json.Unmarshal(data, &tools); err != nil {
		return nil, fmt.Errorf("parse tools: %w", err)
	}
	return tools, nil
}

// CallTool executes a tool and returns the result.
func (c *Conn) CallTool(ctx context.Context, name string, args json.RawMessage) (json.RawMessage, error) {
	fid, err := c.fsys.Open("proc/new", plan9.ORDWR)
	if err != nil {
		return nil, fmt.Errorf("open proc/new: %w", err)
	}
	defer fid.Close()

	// Convert JSON args to key=value format
	var argMap map[string]interface{}
	if err := json.Unmarshal(args, &argMap); err != nil {
		argMap = make(map[string]interface{})
	}

	payload := fmt.Sprintf("token=%s\ntool=%s\n", c.token, name)
	for k, v := range argMap {
		escaped := escapeValue(fmt.Sprintf("%v", v))
		payload += fmt.Sprintf("%s=%s\n", k, escaped)
	}

	if _, err := fid.Write([]byte(payload)); err != nil {
		return nil, fmt.Errorf("write: %w", err)
	}

	// Read result from offset 0 (Write advances the file offset, but
	// for request-response files we need to read the result from the start)
	var result []byte
	buf := make([]byte, 8192)
	for offset := int64(0); ; {
		n, err := fid.ReadAt(buf, offset)
		if n > 0 {
			result = append(result, buf[:n]...)
			offset += int64(n)
		}
		if err == io.EOF || n == 0 {
			break
		}
		if err != nil {
			return nil, fmt.Errorf("read: %w", err)
		}
	}

	return json.RawMessage(result), nil
}

// LoadTool loads a tool by name.
func (c *Conn) LoadTool(name string) error {
	fid, err := c.fsys.Open("ctl", plan9.OWRITE)
	if err != nil {
		return err
	}
	defer fid.Close()
	_, err = fid.Write([]byte("load " + name + "\n"))
	return err
}

// UnloadTool unloads a tool by name.
func (c *Conn) UnloadTool(name string) error {
	fid, err := c.fsys.Open("ctl", plan9.OWRITE)
	if err != nil {
		return err
	}
	defer fid.Close()
	_, err = fid.Write([]byte("unload " + name + "\n"))
	return err
}

// Ping checks if the connection is alive.
func (c *Conn) Ping() error {
	_, err := c.fsys.Stat("info")
	return err
}

// Close closes the connection.
func (c *Conn) Close() {
	if c.conn != nil {
		c.conn.Close()
	}
}

// HostInfo returns host information from the server.
func (c *Conn) HostInfo() (HostInfo, error) {
	fid, err := c.fsys.Open("info", plan9.OREAD)
	if err != nil {
		return HostInfo{}, err
	}
	defer fid.Close()

	data, err := io.ReadAll(fid)
	if err != nil {
		return HostInfo{}, err
	}

	info := HostInfo{}
	for _, line := range splitLines(string(data)) {
		if idx := indexOf(line, '='); idx >= 0 {
			key := line[:idx]
			val := line[idx+1:]
			switch key {
			case "platform":
				info.Platform = val
			case "git":
				info.IsGitRepo = val == "true"
			}
		}
	}
	return info, nil
}

// --- Environment and Configuration ---

// SetEnv sets an environment variable via ctl command.
func (c *Conn) SetEnv(key, value string) error {
	fid, err := c.fsys.Open("ctl", plan9.OWRITE)
	if err != nil {
		return err
	}
	defer fid.Close()
	_, err = fid.Write([]byte(fmt.Sprintf("env %s=%s\n", key, value)))
	return err
}

// SetCWD sets the working directory via ctl command.
func (c *Conn) SetCWD(dir string) error {
	fid, err := c.fsys.Open("ctl", plan9.OWRITE)
	if err != nil {
		return err
	}
	defer fid.Close()
	_, err = fid.Write([]byte(fmt.Sprintf("cwd %s\n", dir)))
	return err
}

// SetOnToolsChanged stores a callback that is invoked when tools change.
// In 9P model, this is handled by polling ToolRegistryRevision().
// The callback is stored but you must call CheckToolsChanged() periodically.
func (c *Conn) SetOnToolsChanged(fn func()) {
	c.onToolsChanged = fn
}

// ToolRegistryRevision returns the current tool registry revision.
// Returns 0 if not available (always refresh).
func (c *Conn) ToolRegistryRevision() uint64 {
	fid, err := c.fsys.Open("info", plan9.OREAD)
	if err != nil {
		return 0
	}
	defer fid.Close()

	data, err := io.ReadAll(fid)
	if err != nil {
		return 0
	}

	for _, line := range splitLines(string(data)) {
		if idx := indexOf(line, '='); idx >= 0 {
			if line[:idx] == "tools_rev" {
				var rev uint64
				fmt.Sscanf(line[idx+1:], "%d", &rev)
				return rev
			}
		}
	}
	return 0
}

// --- Process management ---

// Detach detaches the current foreground process to background.
// Returns true if successful.
func (c *Conn) Detach() bool {
	fid, err := c.fsys.Open("ctl", plan9.OWRITE)
	if err != nil {
		return false
	}
	defer fid.Close()
	_, err = fid.Write([]byte("detach\n"))
	return err == nil
}

// ListDetachedRaw returns raw info about detached processes.
func (c *Conn) ListDetachedRaw() []any {
	fid, err := c.fsys.Open("proc", plan9.OREAD)
	if err != nil {
		return nil
	}
	defer fid.Close()

	var result []any
	// Read directory entries
	for {
		dirs, err := fid.Dirread()
		if err != nil || len(dirs) == 0 {
			break
		}
		for _, d := range dirs {
			if d.Name == "new" || d.Name == "new.bg" {
				continue
			}
			info := c.readProcInfo(d.Name)
			if info != nil {
				result = append(result, info)
			}
		}
	}
	return result
}

// readProcInfo reads info about a single proc.
func (c *Conn) readProcInfo(name string) map[string]any {
	path := fmt.Sprintf("proc/%s/stat", name)
	fid, err := c.fsys.Open(path, plan9.OREAD)
	if err != nil {
		return nil
	}
	defer fid.Close()

	data, err := io.ReadAll(fid)
	if err != nil {
		return nil
	}

	info := make(map[string]any)
	for _, line := range splitLines(string(data)) {
		if idx := indexOf(line, '='); idx >= 0 {
			key := line[:idx]
			val := line[idx+1:]
			switch key {
			case "pid":
				var pid int
				fmt.Sscanf(val, "%d", &pid)
				info["pid"] = pid
			case "command":
				info["command"] = val
			case "started":
				var ts int64
				fmt.Sscanf(val, "%d", &ts)
				info["started"] = ts
			case "exited":
				info["exited"] = val == "true"
			case "exit_code":
				var code int
				fmt.Sscanf(val, "%d", &code)
				info["exit_code"] = code
			}
		}
	}
	return info
}

// SignalDetached sends a signal to a detached process.
func (c *Conn) SignalDetached(pid int, sig syscall.Signal) error {
	path := fmt.Sprintf("proc/%d/ctl", pid)
	fid, err := c.fsys.Open(path, plan9.OWRITE)
	if err != nil {
		return err
	}
	defer fid.Close()
	_, err = fid.Write([]byte(fmt.Sprintf("signal %d\n", sig)))
	return err
}

// DismissDetached dismisses a detached process.
func (c *Conn) DismissDetached(pid int) bool {
	path := fmt.Sprintf("proc/%d/ctl", pid)
	fid, err := c.fsys.Open(path, plan9.OWRITE)
	if err != nil {
		return false
	}
	defer fid.Close()
	_, err = fid.Write([]byte("dismiss\n"))
	return err == nil
}

// GetDetachedOutput gets output from a detached process.
func (c *Conn) GetDetachedOutput(pid int) (string, error) {
	path := fmt.Sprintf("proc/%d/out", pid)
	fid, err := c.fsys.Open(path, plan9.OREAD)
	if err != nil {
		return "", err
	}
	defer fid.Close()
	data, err := io.ReadAll(fid)
	return string(data), err
}

// --- Helper types ---

// HostInfo contains information about the host running toolsrv.
type HostInfo struct {
	Platform  string
	IsGitRepo bool
}

// RateLimitedError indicates a rate limit was hit.
type RateLimitedError struct {
	Err       error
	Remaining int
}

func (e RateLimitedError) Error() string {
	return e.Err.Error()
}

// WithOutputStream returns a context with streaming output function attached.
func WithOutputStream(ctx context.Context, fn func(string)) context.Context {
	return WithStreamFunc(ctx, fn)
}

// --- Helpers ---

func randomSecret() string {
	b := make([]byte, 16)
	rand.Read(b)
	return hex.EncodeToString(b)
}

func splitLines(s string) []string {
	var lines []string
	start := 0
	for i := 0; i < len(s); i++ {
		if s[i] == '\n' {
			lines = append(lines, s[start:i])
			start = i + 1
		}
	}
	if start < len(s) {
		lines = append(lines, s[start:])
	}
	return lines
}

func indexOf(s string, c byte) int {
	for i := 0; i < len(s); i++ {
		if s[i] == c {
			return i
		}
	}
	return -1
}

func escapeValue(s string) string {
	// Escape newlines and backslashes
	result := make([]byte, 0, len(s))
	for i := 0; i < len(s); i++ {
		switch s[i] {
		case '\n':
			result = append(result, '\\', 'n')
		case '\\':
			result = append(result, '\\', '\\')
		default:
			result = append(result, s[i])
		}
	}
	return string(result)
}