oikos/internal/actuator/client.go

package actuator

import (
	"bytes"
	"context"
	"fmt"
	"net"
	"os"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
)

// defaultDialTimeout bounds an SSH dial when the caller leaves Timeout unset.
// 10s matches the previous hardcoded value at every dial site.
const defaultDialTimeout = 10 * time.Second

// LoadSigner reads and parses the private key at keyPath.
func LoadSigner(keyPath string) (ssh.Signer, error) {
	key, err := os.ReadFile(keyPath)
	if err != nil {
		return nil, fmt.Errorf("read ssh key: %w", err)
	}
	return LoadSignerFromBytes(key)
}

// LoadSignerFromBytes parses an in-memory private key into an ssh.Signer.
func LoadSignerFromBytes(key []byte) (ssh.Signer, error) {
	signer, err := ssh.ParsePrivateKey(key)
	if err != nil {
		return nil, fmt.Errorf("parse ssh key: %w", err)
	}
	return signer, nil
}

// DialOptions configures an SSH dial.
type DialOptions struct {
	Host    string
	Port    int           // 0 means 22
	User    string
	Signer  ssh.Signer
	Timeout time.Duration // dial timeout; <=0 means defaultDialTimeout
}

// Dial opens a crypto/ssh connection through the centralized HostKeyCallback.
// The connection itself is bounded by Timeout; ctx is respected by callers
// via RunCombinedOutput once the session is running.
func Dial(ctx context.Context, opts DialOptions) (*ssh.Client, error) {
	port := opts.Port
	if port <= 0 {
		port = 22
	}
	timeout := opts.Timeout
	if timeout <= 0 {
		timeout = defaultDialTimeout
	}
	cfg := &ssh.ClientConfig{
		User:            opts.User,
		Auth:            []ssh.AuthMethod{ssh.PublicKeys(opts.Signer)},
		HostKeyCallback: HostKeyCallback(),
		Timeout:         timeout,
	}
	addr := net.JoinHostPort(opts.Host, fmt.Sprintf("%d", port))
	client, err := ssh.Dial("tcp", addr, cfg)
	if err != nil {
		return nil, fmt.Errorf("ssh dial %s:%d: %w", opts.Host, port, err)
	}
	return client, nil
}

// RunCombinedOutput runs cmd on an established client and returns its combined
// stdout/stderr. Context cancellation closes the session to abort the remote
// command instead of blocking until it finishes — the same goroutine+select
// pattern the actuator, mcp, and scheduler each reimplemented before.
func RunCombinedOutput(ctx context.Context, client *ssh.Client, cmd string) ([]byte, error) {
	session, err := client.NewSession()
	if err != nil {
		return nil, fmt.Errorf("create session: %w", err)
	}
	defer session.Close()

	type result struct {
		out []byte
		err error
	}
	ch := make(chan result, 1)
	go func() {
		out, err := session.CombinedOutput(cmd)
		ch <- result{out: out, err: err}
	}()

	select {
	case <-ctx.Done():
		session.Close()
		return nil, ctx.Err()
	case res := <-ch:
		if res.err != nil {
			return res.out, fmt.Errorf("command: %w", res.err)
		}
		return res.out, nil
	}
}

// RunOutput runs cmd on an established client and returns stdout only.
// Stderr is folded into the returned error so callers that parse stdout
// as JSON (e.g. the scheduler's check scripts) don't get interleaved
// stderr in the output stream.
func RunOutput(ctx context.Context, client *ssh.Client, cmd string) ([]byte, error) {
	session, err := client.NewSession()
	if err != nil {
		return nil, fmt.Errorf("create session: %w", err)
	}
	defer session.Close()

	var outBuf, errBuf bytes.Buffer
	session.Stdout = &outBuf
	session.Stderr = &errBuf

	type result struct {
		runErr error
	}
	ch := make(chan result, 1)
	go func() {
		ch <- result{runErr: session.Run(cmd)}
	}()

	select {
	case <-ctx.Done():
		session.Close()
		return nil, ctx.Err()
	case res := <-ch:
		if res.runErr != nil {
			if errBuf.Len() > 0 {
				return outBuf.Bytes(), fmt.Errorf("command: %w\nstderr: %s", res.runErr, strings.TrimSpace(errBuf.String()))
			}
			return outBuf.Bytes(), fmt.Errorf("command: %w", res.runErr)
		}
		return outBuf.Bytes(), nil
	}
}