llmgw/internal/join/serve.go

package join

import (
	"bufio"
	"bytes"
	"context"
	"crypto/sha256"
	"crypto/subtle"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"os"
	"os/exec"
	"strings"
	"time"
)

// Serve dials the setup a code names and runs what it sends until it says it
// is done. Each script is reported to log by its first line, so whoever ran
// `llmgw join` sees what is being done to their machine.
func Serve(ctx context.Context, c Code, log func(string)) error {
	d := tls.Dialer{
		NetDialer: &net.Dialer{Timeout: 10 * time.Second},
		Config: &tls.Config{
			MinVersion: tls.VersionTLS13,
			// Trusted by the fingerprint in the code, not by an authority:
			// the listener's certificate is made for this setup alone.
			InsecureSkipVerify: true, //nolint:gosec // verified below, by fingerprint
		},
	}
	want := c.Fingerprint[:]
	d.Config.VerifyPeerCertificate = func(raw [][]byte, _ [][]*x509.Certificate) error {
		if len(raw) == 0 {
			return errors.New("the setup sent no certificate")
		}
		sum := sha256.Sum256(raw[0])
		if subtle.ConstantTimeCompare(sum[:16], want) != 1 {
			return errors.New("the machine at that address is not the setup this code is for")
		}
		return nil
	}
	conn, err := d.DialContext(ctx, "tcp", c.Addr)
	if err != nil {
		return fmt.Errorf("cannot reach the setup at %s: %w; it listens there only while it waits for this machine", c.Addr, err)
	}
	defer conn.Close()

	host, _ := os.Hostname()
	local, _, _ := net.SplitHostPort(conn.LocalAddr().String())
	b, _ := json.Marshal(hello{Secret: c.Secret[:], Hostname: host, Addr: local})
	if _, err := conn.Write(append(b, '\n')); err != nil {
		return err
	}
	r := bufio.NewReader(conn)
	line, err := r.ReadBytes('\n')
	if err != nil {
		return fmt.Errorf("the setup closed the connection: %w", err)
	}
	var ack response
	if err := json.Unmarshal(line, &ack); err != nil {
		return errors.New("the setup answered something unreadable")
	}
	if ack.Error != "" {
		return errors.New(ack.Error)
	}
	log(fmt.Sprintf("joined the setup at %s; it is setting this machine up", c.Addr))

	for {
		line, err := r.ReadBytes('\n')
		if err != nil {
			return fmt.Errorf("the setup went away before it said it was done: %w", err)
		}
		var req request
		if err := json.Unmarshal(line, &req); err != nil {
			return errors.New("the setup sent something unreadable")
		}
		if req.Bye {
			log("the setup is done with this machine")
			return nil
		}
		log(summary(req.Script))
		resp := run(ctx, req)
		out, _ := json.Marshal(resp)
		if _, err := conn.Write(append(out, '\n')); err != nil {
			return err
		}
	}
}

// run runs one script as the setup asked.
func run(ctx context.Context, req request) response {
	cmd := exec.CommandContext(ctx, "sh", "-c", req.Script)
	if req.Stdin != nil {
		cmd.Stdin = bytes.NewReader(req.Stdin)
	}
	var stdout, stderr bytes.Buffer
	cmd.Stdout, cmd.Stderr = &stdout, &stderr
	err := cmd.Run()
	resp := response{Stdout: stdout.Bytes(), Stderr: stderr.Bytes()}
	var exit *exec.ExitError
	switch {
	case errors.As(err, &exit):
		resp.Exit = exit.ExitCode()
	case err != nil:
		resp.Error = err.Error()
	}
	return resp
}

// summary is a script's first line that says something.
func summary(script string) string {
	for _, l := range strings.Split(script, "\n") {
		l = strings.TrimSpace(l)
		if l == "" || l == "set -e" || strings.HasPrefix(l, "#") {
			continue
		}
		if len(l) > 100 {
			l = l[:100] + "…"
		}
		return "  " + l
	}
	return "  (an empty script)"
}