sherlock/cmd/gssh-mcp/tools_ssh.go

package main

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/coder/websocket"
	"github.com/modelcontextprotocol/go-sdk/mcp"

	"gitea.alexandru.macocian.me/amacocian/sherlock/internal/wsdatagram"
)

// hardClientTimeout caps the entire run_command tool call. The server
// itself bounds the remote command at 60s; we add 30s headroom for
// WebSocket teardown and to surface a Timeout datagram cleanly.
const hardClientTimeout = 90 * time.Second

// stdoutCap and stderrCap bound what we return to the LLM. Anything
// past this point is truncated with a flag the agent can act on.
const (
	stdoutCap = 256 * 1024
	stderrCap = 64 * 1024
)

func registerSshTools(server *mcp.Server, api *gsshAPI) {
	mcp.AddTool(server, &mcp.Tool{
		Name: "run_command",
		Description: "Run a single shell command on the named SSH host through the Gssh gateway (no PTY). " +
			"Returns the captured stdout/stderr and the exit code, or a timeout flag if the server-side " +
			"60-second cap fires first. stdout is truncated at 256 KiB, stderr at 64 KiB.",
	}, runCommandHandler(api))
}

type runCommandInput struct {
	Host    string `json:"host" jsonschema:"the SSH host (must be in list_hosts output)"`
	Command string `json:"command" jsonschema:"the shell command to execute on the remote host"`
}

type runCommandOutput struct {
	Stdout          string `json:"stdout"`
	Stderr          string `json:"stderr"`
	ExitCode        *int   `json:"exit_code,omitempty"`
	TimedOut        bool   `json:"timed_out"`
	StdoutTruncated bool   `json:"stdout_truncated,omitempty"`
	StderrTruncated bool   `json:"stderr_truncated,omitempty"`
}

func runCommandHandler(api *gsshAPI) mcp.ToolHandlerFor[runCommandInput, runCommandOutput] {
	return func(ctx context.Context, _ *mcp.CallToolRequest, in runCommandInput) (*mcp.CallToolResult, runCommandOutput, error) {
		if strings.TrimSpace(in.Host) == "" {
			return nil, runCommandOutput{}, errors.New("run_command: host is required")
		}
		if strings.TrimSpace(in.Command) == "" {
			return nil, runCommandOutput{}, errors.New("run_command: command is required")
		}
		out, err := execOverWebSocket(ctx, api, in.Host, in.Command)
		if err != nil {
			return nil, runCommandOutput{}, err
		}
		return nil, out, nil
	}
}

// execOverWebSocket dials the Gssh exec endpoint, sends a Start
// datagram with the command, demuxes Stdout/Stderr datagrams into
// capped buffers, and returns when an Exit or Timeout datagram
// arrives (or the WS closes for any other reason).
func execOverWebSocket(parent context.Context, api *gsshAPI, host, command string) (runCommandOutput, error) {
	ctx, cancel := context.WithTimeout(parent, hardClientTimeout)
	defer cancel()

	wsURL := buildWSURL(api.baseURL, host)
	dbg("WS dial %s", wsURL)

	tok, err := api.token(ctx)
	if err != nil {
		return runCommandOutput{}, err
	}
	headers := http.Header{}
	headers.Set("Authorization", "Bearer "+tok)

	conn, resp, err := websocket.Dial(ctx, wsURL, &websocket.DialOptions{
		HTTPHeader: headers,
		HTTPClient: &http.Client{
			// Don't reuse api.client — it has a 30s timeout, but the
			// WS handshake should complete in well under that, and we
			// don't want the timeout to fire mid-stream on the
			// upgraded connection.
			Timeout: 30 * time.Second,
		},
	})
	if err != nil {
		if resp != nil && resp.StatusCode == http.StatusUnauthorized {
			return runCommandOutput{}, fmt.Errorf("%w (Gssh rejected the bearer token on WS upgrade for %s)", ErrUnauthorized, host)
		}
		if resp != nil {
			return runCommandOutput{}, fmt.Errorf("gssh: WS dial %s: HTTP %d: %w", host, resp.StatusCode, err)
		}
		return runCommandOutput{}, fmt.Errorf("gssh: WS dial %s: %w", host, err)
	}
	// Per coder/websocket convention, always Close to release the
	// goroutine. StatusNormalClosure is the polite signal back to the
	// server even on errors; the server already finalised by then.
	defer func() { _ = conn.Close(websocket.StatusNormalClosure, "") }()

	// Lift the per-message size limit. The default 32 KiB is enough
	// for one datagram (2 KiB) but the lift is cheap defensive code
	// in case the server is upgraded to send larger frames later.
	conn.SetReadLimit(64 * 1024)

	// Send the Start datagram.
	startFrame, err := wsdatagram.Encode(wsdatagram.Datagram{
		Op:      wsdatagram.OpStart,
		Payload: []byte(command),
	})
	if err != nil {
		return runCommandOutput{}, fmt.Errorf("gssh: encode Start: %w", err)
	}
	if err := conn.Write(ctx, websocket.MessageBinary, startFrame); err != nil {
		return runCommandOutput{}, fmt.Errorf("gssh: send Start: %w", err)
	}
	dbg("WS sent Start (%d bytes command)", len(command))

	var (
		stdoutBuf bytes.Buffer
		stderrBuf bytes.Buffer
		out       runCommandOutput
	)

	for {
		msgType, data, err := conn.Read(ctx)
		if err != nil {
			// Normal closure after we've already received Exit or
			// Timeout: return what we have. Otherwise, surface the
			// error along with whatever buffers we accumulated.
			status := websocket.CloseStatus(err)
			if status == websocket.StatusNormalClosure || status == websocket.StatusGoingAway {
				out.Stdout = stdoutBuf.String()
				out.Stderr = stderrBuf.String()
				if !out.TimedOut && out.ExitCode == nil {
					// Server closed cleanly but never sent a terminal
					// datagram. Treat as a Gssh protocol violation
					// rather than silently making something up.
					return out, fmt.Errorf("gssh: WS closed without Exit or Timeout datagram")
				}
				return out, nil
			}
			return runCommandOutput{
				Stdout:          stdoutBuf.String(),
				Stderr:          stderrBuf.String(),
				StdoutTruncated: out.StdoutTruncated,
				StderrTruncated: out.StderrTruncated,
			}, fmt.Errorf("gssh: WS read: %w", err)
		}
		if msgType != websocket.MessageBinary {
			dbg("WS ignored non-binary frame (type=%v)", msgType)
			continue
		}
		dg, err := wsdatagram.Decode(data)
		if err != nil {
			return runCommandOutput{
				Stdout: stdoutBuf.String(),
				Stderr: stderrBuf.String(),
			}, fmt.Errorf("gssh: decode datagram: %w", err)
		}
		switch dg.Op {
		case wsdatagram.OpStdout:
			appendCapped(&stdoutBuf, dg.Payload, stdoutCap, &out.StdoutTruncated)
		case wsdatagram.OpStderr:
			appendCapped(&stderrBuf, dg.Payload, stderrCap, &out.StderrTruncated)
		case wsdatagram.OpExit:
			code, perr := strconv.Atoi(string(dg.Payload))
			if perr != nil {
				return runCommandOutput{
					Stdout: stdoutBuf.String(),
					Stderr: stderrBuf.String(),
				}, fmt.Errorf("gssh: parse Exit payload %q: %w", dg.Payload, perr)
			}
			out.ExitCode = &code
			out.Stdout = stdoutBuf.String()
			out.Stderr = stderrBuf.String()
			dbg("WS got Exit %d", code)
			return out, nil
		case wsdatagram.OpTimeout:
			out.TimedOut = true
			out.Stdout = stdoutBuf.String()
			out.Stderr = stderrBuf.String()
			dbg("WS got Timeout")
			return out, nil
		default:
			dbg("WS ignored unknown op %v", dg.Op)
		}
	}
}

// appendCapped writes payload into buf, stopping when cap is reached
// and setting *truncated. Once truncated, further appends are no-ops.
func appendCapped(buf *bytes.Buffer, payload []byte, cap int, truncated *bool) {
	if *truncated {
		return
	}
	remaining := cap - buf.Len()
	if remaining <= 0 {
		*truncated = true
		return
	}
	if len(payload) > remaining {
		buf.Write(payload[:remaining])
		*truncated = true
		return
	}
	buf.Write(payload)
}

// buildWSURL turns the https://gssh.example/.../api/v1/exec/{host}
// REST URL into its wss:// equivalent.
func buildWSURL(baseURL, host string) string {
	u, _ := url.Parse(baseURL)
	switch u.Scheme {
	case "https":
		u.Scheme = "wss"
	case "http":
		u.Scheme = "ws"
	}
	u.Path = strings.TrimRight(u.Path, "/") + "/api/v1/exec/" + url.PathEscape(host)
	return u.String()
}