mesh-catalog/modules/ssh-client/cmd/ssh-client-tools/keykeeper_bus_test.go

package main

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/nats-io/nats-server/v2/server"
	"github.com/nats-io/nats.go"
	"github.com/nats-io/nats.go/jetstream"
)

// Issue 955, area S4. The acceptance: "a machine's account makes its forge key, and the key's public half is
// in ssh-client's keys state under <node>.<account>" — proved over a real bus. A JetStream server is linked into
// the test and started for it alone; the module is this test binary started again as the bundle, speaking MCP on
// its standard streams to a runtime played here, which keeps the state in buckets of the bus as the tool runner
// does. It runs in the forge's merge check, not on the way: `go test -short` leaves it out.

const bundleVar = "MESH_SSH_CLIENT_TEST_BUNDLE"

func TestMain(m *testing.M) {
	if os.Getenv(bundleVar) == "1" {
		main()
		os.Exit(0)
	}
	os.Exit(m.Run())
}

type keysRuntime struct {
	t       *testing.T
	js      jetstream.JetStream
	in      io.WriteCloser
	mu      sync.Mutex
	buckets map[string]jetstream.KeyValue
}

func (r *keysRuntime) bucket(name string) jetstream.KeyValue {
	r.mu.Lock()
	defer r.mu.Unlock()
	if kv, ok := r.buckets[name]; ok {
		return kv
	}
	kv, err := r.js.CreateOrUpdateKeyValue(context.Background(), jetstream.KeyValueConfig{Bucket: name})
	if err != nil {
		r.t.Fatal(err)
	}
	r.buckets[name] = kv
	return kv
}

func (r *keysRuntime) say(m map[string]any) {
	m["jsonrpc"] = "2.0"
	raw, _ := json.Marshal(m)
	r.mu.Lock()
	defer r.mu.Unlock()
	_, _ = r.in.Write(append(raw, '\n'))
}

func (r *keysRuntime) read(out io.Reader) {
	sc := bufio.NewScanner(out)
	sc.Buffer(make([]byte, 1<<20), 1<<24)
	for sc.Scan() {
		var m struct {
			ID     json.RawMessage `json:"id"`
			Method string          `json:"method"`
			Params json.RawMessage `json:"params"`
		}
		if json.Unmarshal(sc.Bytes(), &m) != nil || m.Method == "" {
			continue
		}
		go r.runtime(m.ID, m.Method, m.Params)
	}
}

// runtime is what the tool runner does for the bundle: only the state this module puts.
func (r *keysRuntime) runtime(id json.RawMessage, method string, params json.RawMessage) {
	var p struct {
		State string          `json:"state"`
		Key   string          `json:"key"`
		Value json.RawMessage `json:"value"`
	}
	_ = json.Unmarshal(params, &p)
	if method != "mesh/state.put" {
		r.say(map[string]any{"id": id, "error": map[string]any{"code": -32000, "message": fmt.Sprintf("the test runtime offers no %s", method)}})
		return
	}
	rev, err := r.bucket(p.State).Put(context.Background(), p.Key, p.Value)
	if err != nil {
		r.say(map[string]any{"id": id, "error": map[string]any{"code": -32000, "message": err.Error()}})
		return
	}
	r.say(map[string]any{"id": id, "result": map[string]any{"revision": rev}})
}

type keysLog struct{ t *testing.T }

func (l keysLog) Write(p []byte) (int, error) { l.t.Logf("bundle: %s", p); return len(p), nil }

func TestTheKeyProcessMakesTheKeyAndPutsItsEntryOverTheBus(t *testing.T) {
	if testing.Short() {
		t.Skip("integration test: runs in the forge's merge check")
	}
	s, err := server.NewServer(&server.Options{Host: "127.0.0.1", Port: server.RANDOM_PORT, JetStream: true, StoreDir: t.TempDir(), NoLog: true, NoSigs: true})
	if err != nil {
		t.Fatal(err)
	}
	go s.Start()
	if !s.ReadyForConnections(30 * time.Second) {
		t.Fatal("no bus within 30s")
	}
	t.Cleanup(func() { s.Shutdown(); s.WaitForShutdown() })
	nc, err := nats.Connect(s.ClientURL())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(nc.Close)
	js, _ := jetstream.New(nc)
	r := &keysRuntime{t: t, js: js, buckets: map[string]jetstream.KeyValue{}}

	account, err := user.Current()
	if err != nil {
		t.Fatal(err)
	}
	home := t.TempDir()
	path := os.Getenv("PATH")
	if _, err := exec.LookPath("ssh-keygen"); err != nil {
		// A build host without OpenSSH: a stand-in that writes the pair the command names.
		bin := t.TempDir()
		stub := "#!/bin/sh\nwhile [ $# -gt 0 ]; do case $1 in -f) f=$2;; -C) c=$2;; esac; shift; done\n" +
			"printf 'PRIVATE\\n' > \"$f\"; chmod 600 \"$f\"\n" +
			"printf 'ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA %s\\n' \"$c\" > \"$f.pub\"\n"
		if err := os.WriteFile(filepath.Join(bin, "ssh-keygen"), []byte(stub), 0o755); err != nil {
			t.Fatal(err)
		}
		path = bin + string(os.PathListSeparator) + path
		t.Log("no ssh-keygen on this host: using a stand-in")
	}

	cmd := exec.Command(os.Args[0], "-test.run=^$")
	cmd.Env = append(os.Environ(), bundleVar+"=1", "PATH="+path, "MESH_OPERATOR_HOME="+home,
		"MESH_SERVED_MODULE=ssh-client", "MESH_MODULE=ssh-client", "MESH_NODE=test")
	in, _ := cmd.StdinPipe()
	out, _ := cmd.StdoutPipe()
	cmd.Stderr = keysLog{t}
	if err := cmd.Start(); err != nil {
		t.Fatal(err)
	}
	r.in = in
	go r.read(out)
	t.Cleanup(func() { _ = in.Close(); _ = cmd.Wait() })

	key := "test." + account.Username
	var raw []byte
	for deadline := time.Now().Add(60 * time.Second); time.Now().Before(deadline); time.Sleep(100 * time.Millisecond) {
		if e, err := r.bucket("keys").Get(context.Background(), key); err == nil {
			raw = e.Value()
			break
		}
	}
	if raw == nil {
		t.Fatalf("no entry %s in the keys state within 60 s", key)
	}
	var e KeyEntry
	if err := json.Unmarshal(raw, &e); err != nil {
		t.Fatal(err)
	}
	if !strings.HasPrefix(e.Public, "ssh-ed25519 ") || e.Node != "test" || e.Account != account.Username || e.Kind != "machine" || !strings.HasPrefix(e.Fingerprint, "SHA256:") {
		t.Errorf("entry = %s", raw)
	}
	if strings.Contains(strings.ToUpper(string(raw)), "PRIVATE") {
		t.Errorf("the entry carries private material: %s", raw)
	}
	info, err := os.Stat(filepath.Join(home, ".ssh", KeyFile))
	if err != nil {
		t.Fatal(err)
	}
	if info.Mode().Perm() != 0o600 {
		t.Errorf("the private key is %04o, want 0600", info.Mode().Perm())
	}
}