feat: join-key enrollment (v1.5.0, protocol v1.2.0 1.1)

Installing the agent with one key is now all it takes to add a host. New `join_key` config field, presented while auth_token is empty. The SSO exchanges it for this agent's own token and the public key it must pin, both delivered in the config frame; the agent persists them and blanks the join key. Nothing has to be copied between two machines by hand. PersistEnrollment rewrites only the credential lines -- line-based rather than a YAML round-trip -- so operator comments, the capability matrix and formatting survive. It re-reads afterwards so the credential is live without a restart, and keeps the file 0600. The connect URL carries ?hostname= so a self-enrolling host is named after itself, and the agent refuses to connect at all (with a long back-off) when it has no credential rather than presenting an empty one. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-08-06 10:34:06 -04:00
parent dc14274edc
commit 0a5011cd42
8 changed files with 296 additions and 18 deletions
@@ -3,6 +3,7 @@ package main
 import (
 	"os"
 	"path/filepath"
+	"strings"
 	"testing"
 )

@@ -96,3 +97,108 @@ func TestCanManageService(t *testing.T) {
 		})
 	}
 }
+
+// PersistEnrollment rewrites a hand-edited file, so it must replace exactly the
+// credential lines and leave everything else -- comments, capabilities,
+// formatting -- untouched.
+func TestPersistEnrollmentPreservesFile(t *testing.T) {
+	dir := t.TempDir()
+	path := dir + "/agent.yml"
+	original := `# theta-agent configuration file
+server_url: "https://sso.example.com"
+
+# Bootstrap credential, exchanged on first connect.
+join_key: "tjk_abc123"
+public_key: ""
+location: "rack-4"
+
+capabilities:
+  telemetry: true
+  # keep this comment
+  reboot: false
+  service_control: ["nginx"]
+`
+	if err := os.WriteFile(path, []byte(original), 0600); err != nil {
+		t.Fatal(err)
+	}
+	cm, err := NewConfigManager(path)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if err := cm.PersistEnrollment("issued-token-xyz", "PUBKEYBASE64"); err != nil {
+		t.Fatalf("PersistEnrollment: %v", err)
+	}
+
+	out, _ := os.ReadFile(path)
+	got := string(out)
+
+	for _, want := range []string{
+		`auth_token: "issued-token-xyz"`,
+		`public_key: "PUBKEYBASE64"`,
+		`join_key: ""`, // blanked: a fleet-wide key must not linger once unneeded
+		"# theta-agent configuration file",
+		"# keep this comment",
+		`location: "rack-4"`,
+		`service_control: ["nginx"]`,
+	} {
+		if !strings.Contains(got, want) {
+			t.Errorf("expected %q in rewritten config, got:\n%s", want, got)
+		}
+	}
+
+	// and the in-memory config is live without a restart
+	if cm.Get().AuthToken != "issued-token-xyz" {
+		t.Errorf("config not reloaded: AuthToken = %q", cm.Get().AuthToken)
+	}
+	if cm.Get().Credential() != "issued-token-xyz" {
+		t.Errorf("Credential() should prefer the issued token, got %q", cm.Get().Credential())
+	}
+
+	// file must stay 0600 -- it now holds a credential
+	fi, _ := os.Stat(path)
+	if fi.Mode().Perm() != 0600 {
+		t.Errorf("expected mode 0600, got %o", fi.Mode().Perm())
+	}
+}
+
+func TestPersistEnrollmentAddsMissingKeys(t *testing.T) {
+	dir := t.TempDir()
+	path := dir + "/agent.yml"
+	// No auth_token or public_key lines at all.
+	if err := os.WriteFile(path, []byte("server_url: \"https://sso.example.com\"\njoin_key: \"tjk_x\"\n"), 0600); err != nil {
+		t.Fatal(err)
+	}
+	cm, err := NewConfigManager(path)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := cm.PersistEnrollment("tok", "pk"); err != nil {
+		t.Fatalf("PersistEnrollment: %v", err)
+	}
+	if cm.Get().AuthToken != "tok" || cm.Get().PublicKey != "pk" {
+		out, _ := os.ReadFile(path)
+		t.Errorf("keys not appended; file:\n%s", out)
+	}
+}
+
+func TestCredentialPrefersAuthToken(t *testing.T) {
+	c := &Config{JoinKey: "tjk_x"}
+	if c.Credential() != "tjk_x" {
+		t.Errorf("unenrolled agent should present the join key, got %q", c.Credential())
+	}
+	c.AuthToken = "own-token"
+	if c.Credential() != "own-token" {
+		t.Errorf("enrolled agent must present its own token, not the join key, got %q", c.Credential())
+	}
+}
+
+func TestPersistEnrollmentRejectsEmptyToken(t *testing.T) {
+	dir := t.TempDir()
+	path := dir + "/agent.yml"
+	os.WriteFile(path, []byte("server_url: \"x\"\n"), 0600)
+	cm, _ := NewConfigManager(path)
+	if err := cm.PersistEnrollment("", "pk"); err == nil {
+		t.Error("expected an error when the server sends no token")
+	}
+}