internal/mcpserver/sshrun.go

Ref: Size: 5.6 KiB History
package mcpserver

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"os"
	"path"
	"time"

	"github.com/pkg/sftp"
	"golang.org/x/crypto/ssh"
)

// outputCap bounds captured exec/file bytes returned to the model.
const outputCap = 1 << 20 // 1 MiB

// VMDialer reaches a VM by name and names it the way the caller's world spells
// it. It is the seam between reaching a VM and doing anything to one: the
// control plane tunnels over the VM host's own sync connection with the
// credential that tenant has delegated to it (internal/server/vmssh).
// Everything above this line — exec, SFTP, output capping — is transport-blind.
type VMDialer interface {
	Dial(ctx context.Context, vmName string) (*ssh.Client, error)
	// ConnectName maps a bare VM name to the <tenant>.<name> form the gate
	// resolves and the VM's host certificate names.
	ConnectName(ctx context.Context, vmName string) (string, error)
}

// ExecResult is a completed remote command.
type ExecResult struct {
	Stdout    string
	Stderr    string
	ExitCode  int
	Truncated bool
}

// Runner executes commands and transfers files on VMs over SSH, reaching each
// VM by NAME through its dialer. There is no TOFU/known_hosts anywhere on
// either path: every host key is a certificate verified against the eitri host
// CA, and the client authenticates with a short-lived CA-signed user cert.
type Runner struct {
	dial VMDialer
}

func NewRunner(d VMDialer) *Runner { return &Runner{dial: d} }

// ConnectName returns the connect name for vmName, so the Tools layer can build
// a correct `ssh -J` hint without dialing.
func (r *Runner) ConnectName(ctx context.Context, vmName string) (string, error) {
	return r.dial.ConnectName(ctx, vmName)
}

// Exec runs cmd on the VM named vmName. A non-zero remote exit is NOT an error —
// it's in ExitCode.
func (r *Runner) Exec(ctx context.Context, vmName, cmd string, timeout time.Duration) (ExecResult, error) {
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()
	client, err := r.dial.Dial(ctx, vmName)
	if err != nil {
		return ExecResult{}, err
	}
	defer client.Close()

	sess, err := client.NewSession()
	if err != nil {
		return ExecResult{}, fmt.Errorf("ssh session: %w", err)
	}
	defer sess.Close()

	var stdout, stderr cappedBuf
	sess.Stdout, sess.Stderr = &stdout, &stderr

	done := make(chan error, 1)
	go func() { done <- sess.Run(cmd) }()
	select {
	case <-ctx.Done():
		_ = sess.Close()
		return ExecResult{}, fmt.Errorf("exec timed out after %s", timeout)
	case err = <-done:
	}
	res := ExecResult{Stdout: stdout.String(), Stderr: stderr.String(), Truncated: stdout.truncated || stderr.truncated}
	if exitErr, ok := errors.AsType[*ssh.ExitError](err); ok {
		res.ExitCode = exitErr.ExitStatus()
		return res, nil
	}
	if err != nil {
		return res, fmt.Errorf("exec: %w", err)
	}
	return res, nil
}

// cappedBuf captures at most outputCap bytes and records truncation.
type cappedBuf struct {
	buf       bytes.Buffer
	truncated bool
}

func (b *cappedBuf) Write(p []byte) (int, error) {
	room := outputCap - b.buf.Len()
	if room <= 0 {
		b.truncated = true
		return len(p), nil
	}
	if len(p) > room {
		b.buf.Write(p[:room])
		b.truncated = true
		return len(p), nil
	}
	return b.buf.Write(p)
}

func (b *cappedBuf) String() string { return b.buf.String() }

// WriteFile writes data to remotePath on the VM named vmName over SFTP. Missing
// parent dirs are created at the server default mode; the file itself is set to
// mode. The chmod is applied to the freshly-created file BEFORE any bytes are
// written, so a restrictive mode (e.g. 0600) is never briefly world-readable on
// the guest during the write (pkg/sftp's OpenFile takes no mode argument).
func (r *Runner) WriteFile(ctx context.Context, vmName, remotePath string, data []byte, mode fs.FileMode) error {
	client, sf, err := r.sftp(ctx, vmName)
	if err != nil {
		return err
	}
	defer client.Close()
	defer sf.Close()
	if dir := path.Dir(remotePath); dir != "." && dir != "/" {
		if err := sf.MkdirAll(dir); err != nil {
			return fmt.Errorf("mkdir %s: %w", dir, err)
		}
	}
	f, err := sf.OpenFile(remotePath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC)
	if err != nil {
		return fmt.Errorf("create %s: %w", remotePath, err)
	}
	if err := f.Chmod(mode); err != nil {
		f.Close()
		return fmt.Errorf("chmod %s: %w", remotePath, err)
	}
	if _, err := f.Write(data); err != nil {
		f.Close()
		return fmt.Errorf("write %s: %w", remotePath, err)
	}
	return f.Close()
}

// ReadFile reads at most outputCap bytes from remotePath on the VM named
// vmName; truncated reports whether the file was larger.
func (r *Runner) ReadFile(ctx context.Context, vmName, remotePath string) (data []byte, truncated bool, err error) {
	client, sf, err := r.sftp(ctx, vmName)
	if err != nil {
		return nil, false, err
	}
	defer client.Close()
	defer sf.Close()
	f, err := sf.Open(remotePath)
	if err != nil {
		return nil, false, fmt.Errorf("open %s: %w", remotePath, err)
	}
	defer f.Close()
	buf := make([]byte, outputCap+1)
	n, rerr := io.ReadFull(f, buf)
	if rerr != nil && rerr != io.ErrUnexpectedEOF && rerr != io.EOF {
		return nil, false, fmt.Errorf("read %s: %w", remotePath, rerr)
	}
	if n > outputCap {
		return buf[:outputCap], true, nil
	}
	return buf[:n], false, nil
}

// sftp dials the VM named vmName and wraps the connection in an sftp.Client.
// The caller must Close both the returned *ssh.Client and *sftp.Client.
func (r *Runner) sftp(ctx context.Context, vmName string) (*ssh.Client, *sftp.Client, error) {
	client, err := r.dial.Dial(ctx, vmName)
	if err != nil {
		return nil, nil, err
	}
	sf, err := sftp.NewClient(client)
	if err != nil {
		client.Close()
		return nil, nil, fmt.Errorf("sftp: %w", err)
	}
	return client, sf, nil
}