ollie/toolsrv/transport.go

package toolsrv

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync/atomic"
	"syscall"
	"time"

	"ollie/paths"
)

// socketSeq is a monotonic counter to ensure unique socket paths.
var socketSeq atomic.Int64

// transport encapsulates the plumbing for a JSON-RPC connection to a
// toolsrv subprocess. Both local and SSH transports produce the same
// result: a reader, writer, socket path, and cleanup function.
type transport struct {
	r        io.Reader
	w        io.WriteCloser
	sockPath string
	cleanup  func()
}

// localDial spawns ollie-remote as a local subprocess listening on a Unix
// socket, waits for it to be ready, and returns the socket path.
func localDial(ctx context.Context, cwd string, opts ...Option) (*transport, error) {
	// Bootstrap: verify ollie-remote binary is present.
	binPath, err := findLocalBinary()
	if err != nil {
		return nil, fmt.Errorf("bootstrap: ollie-remote not found")
	}

	// Verify tools directory is accessible via ToolsPath().
	toolsDir := ToolsPath()
	if _, err := os.Stat(toolsDir); err != nil {
		return nil, fmt.Errorf("bootstrap: tools directory not found at %s", toolsDir)
	}

	// Generate a unique socket path.
	sockDir := filepath.Join(runtimeDir(), "ollie")
	os.MkdirAll(sockDir, 0700)
	sockPath := filepath.Join(sockDir, fmt.Sprintf("toolsrv-%d-%d.sock", os.Getpid(), socketSeq.Add(1)))

	args := []string{"serve", "--cwd", cwd, "--listen", sockPath}

	probe := &Server{}
	for _, o := range opts {
		o(probe)
	}
	if probe.Yolo {
		args = append(args, "--yolo")
	}

	cmd := exec.CommandContext(ctx, binPath, args...)
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Setpgid:   true,
		Pdeathsig: syscall.SIGTERM,
	}

	// Inherit environment.
	cmd.Env = os.Environ()

	// Watch stderr for readiness signal.
	stderrPipe, err := cmd.StderrPipe()
	if err != nil {
		return nil, fmt.Errorf("local stderr pipe: %w", err)
	}

	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("local start: %w", err)
	}

	// Wait for ListenReady on stderr.
	readyCh := make(chan struct{})
	go func() {
		scanner := bufio.NewScanner(stderrPipe)
		for scanner.Scan() {
			if strings.TrimSpace(scanner.Text()) == "ListenReady" {
				close(readyCh)
				return
			}
		}
	}()

	select {
	case <-readyCh:
	case <-time.After(10 * time.Second):
		syscall.Kill(-cmd.Process.Pid, syscall.SIGTERM)
		cmd.Wait()
		return nil, fmt.Errorf("local toolsrv startup timeout")
	}

	// Connect to the socket.
	conn, err := net.Dial("unix", sockPath)
	if err != nil {
		syscall.Kill(-cmd.Process.Pid, syscall.SIGTERM)
		cmd.Wait()
		return nil, fmt.Errorf("local socket connect: %w", err)
	}

	cleanup := func() {
		conn.Close()
		syscall.Kill(-cmd.Process.Pid, syscall.SIGTERM)
		done := make(chan struct{})
		go func() { cmd.Wait(); close(done) }()
		select {
		case <-done:
		case <-time.After(3 * time.Second):
			syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
			<-done
		}
		os.Remove(sockPath)
	}

	return &transport{r: conn, w: conn, sockPath: sockPath, cleanup: cleanup}, nil
}

// sshDial connects to a remote host via SSH, bootstraps ollie-remote
// with --listen mode, sets up SSH Unix socket forwarding, and returns
// a transport connected to the forwarded local socket.
//
// Assumes pubkey auth (BatchMode=yes). No password handling.
func sshDial(ctx context.Context, cfg RemoteConfig) (*transport, error) {
	// Bootstrap: verify local environment has everything needed.
	cacheDir := paths.CfgDir()
	if _, err := os.Stat(cacheDir); err != nil {
		return nil, fmt.Errorf("bootstrap: local environment not found at %s", cacheDir)
	}

	// Generate socket paths.
	sockDir := filepath.Join(runtimeDir(), "ollie")
	os.MkdirAll(sockDir, 0700)
	localSock := filepath.Join(sockDir, fmt.Sprintf("remote-%s-%d-%d.sock", cfg.SSHTarget, os.Getpid(), socketSeq.Add(1)))
	remoteSock := fmt.Sprintf("/tmp/ollie-toolsrv-%d.sock", os.Getpid())

	// Remote bootstrap: extract tarball, set env, exec serve.
	bootstrap := fmt.Sprintf(`#!/bin/sh
set -e
CACHE_DIR="${XDG_CACHE_HOME:-$HOME/.cache}/ollie"
mkdir -p "$CACHE_DIR"
tar xzf - -C "$CACHE_DIR"
# CACHE_DIR is the extracted ollie config dir; set XDG_CONFIG_HOME so
# $XDG_CONFIG_HOME/ollie resolves to it.
export XDG_CONFIG_HOME="$(dirname "$CACHE_DIR")"
export PATH="$CACHE_DIR/bin:$PATH"
exec "$CACHE_DIR/bin/ollie-remote" serve --cwd %s --listen %s
`, shellEscape(cfg.CWD), shellEscape(remoteSock))

	sshArgs := []string{
		"-o", "BatchMode=yes",
		"-o", "StrictHostKeyChecking=accept-new",
		"-o", "StreamLocalBindUnlink=yes",
		"-L", localSock + ":" + remoteSock,
	}
	sshTarget := cfg.SSHTarget
	if host, port, ok := strings.Cut(sshTarget, ":"); ok && port != "" {
		sshTarget = host
		sshArgs = append(sshArgs, "-p", port)
	}
	sshArgs = append(sshArgs, sshTarget, "bash -s")

	cmd := exec.CommandContext(ctx, "ssh", sshArgs...)
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Setpgid:   true,
		Pdeathsig: syscall.SIGTERM,
	}

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, fmt.Errorf("ssh stdin pipe: %w", err)
	}
	stderrPipe, err := cmd.StderrPipe()
	if err != nil {
		return nil, fmt.Errorf("ssh stderr pipe: %w", err)
	}

	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("ssh start: %w", err)
	}

	// Send tarball over stdin, then the bootstrap script.
	if err := writeTarball(stdin, cacheDir); err != nil {
		cmd.Process.Kill()
		return nil, fmt.Errorf("send tarball: %w", err)
	}
	if _, err := stdin.Write([]byte(bootstrap + "\n")); err != nil {
		cmd.Process.Kill()
		return nil, fmt.Errorf("send bootstrap: %w", err)
	}

	// Monitor stderr for ListenReady.
	readyCh := make(chan struct{}, 1)
	go func() {
		scanner := bufio.NewScanner(stderrPipe)
		for scanner.Scan() {
			if strings.TrimSpace(scanner.Text()) == "ListenReady" {
				close(readyCh)
				return
			}
		}
	}()

	select {
	case <-readyCh:
	case <-time.After(30 * time.Second):
		cmd.Process.Kill()
		return nil, fmt.Errorf("bootstrap timeout waiting for listen ready")
	}

	// Poll for local socket forwarding.
	var conn net.Conn
	for i := 0; i < 20; i++ {
		conn, err = net.Dial("unix", localSock)
		if err == nil {
			break
		}
		time.Sleep(50 * time.Millisecond)
	}
	if err != nil {
		cmd.Process.Kill()
		return nil, fmt.Errorf("ssh socket connect: %w", err)
	}

	cleanup := func() {
		conn.Close()
		stdin.Close()
		syscall.Kill(-cmd.Process.Pid, syscall.SIGTERM)
		done := make(chan struct{})
		go func() { cmd.Wait(); close(done) }()
		select {
		case <-done:
		case <-time.After(3 * time.Second):
			syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
			<-done
		}
		os.Remove(localSock)
	}

	return &transport{r: conn, w: conn, sockPath: localSock, cleanup: cleanup}, nil
}

// writeTarball creates a gzipped tarball of the directory at dir and writes
// it to w. Used to transfer the ollie environment to remote hosts.
func writeTarball(w io.Writer, dir string) error {
	cmd := exec.Command("tar", "czf", "-", "-C", dir, ".")
	cmd.Stdout = w
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

// findLocalBinary locates the ollie-remote binary on the local system.
func findLocalBinary() (string, error) {
	home, _ := os.UserHomeDir()
	candidates := []string{
		filepath.Join(home, "bin", "ollie-remote"),
		filepath.Join(home, ".local", "bin", "ollie-remote"),
	}
	for _, p := range candidates {
		if _, err := os.Stat(p); err == nil {
			return p, nil
		}
	}
	if p, err := exec.LookPath("ollie-remote"); err == nil {
		return p, nil
	}
	return "", fmt.Errorf("ollie-remote not found (checked %s)", strings.Join(candidates, ", "))
}

// shellEscape escapes a string for safe use in a shell command.
func shellEscape(s string) string {
	if s == "" {
		return "''"
	}
	if !strings.ContainsAny(s, " \t\n\r'\"\\$`!#&|;(){}") {
		return s
	}
	return "'" + strings.ReplaceAll(s, "'", "'\\''") + "'"
}

// runtimeDir returns XDG_RUNTIME_DIR or a fallback.
func runtimeDir() string {
	if d := os.Getenv("XDG_RUNTIME_DIR"); d != "" {
		return d
	}
	return filepath.Join(os.TempDir(), fmt.Sprintf("runtime-%d", os.Getuid()))
}