Merge pull request #5 from theta42/feat/join-key-enrollment

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

CHANGELOG.md

@@ -5,6 +5,20 @@ All notable changes to the `theta-agent` daemon will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [v1.5.0] - 2026-08-06
+
+Join-key enrollment (protocol v1.2.0 §1.1). Installing the agent with one key is now all it takes to add a host.
+
+### Added
+- **`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 writes them into `agent.yml` and blanks the join key. No value has to be copied between two machines by hand any more.
+- `ConfigManager.PersistEnrollment` rewrites only the credential lines, line-based rather than a YAML round-trip, so operator comments, the capability matrix and formatting survive. Re-reads the file afterwards, so the new credential is live without a restart, and keeps the file at `0600`.
+- `Config.Credential()` — the agent's own token when it has one, otherwise the join key.
+- The connect URL carries `?hostname=`, so a self-enrolling host is named after itself instead of a generated placeholder.
+- `install.sh --join-key`.
+
+### Fixed
+- The agent now refuses to connect (with a clear message and a long back-off) when it has neither an `auth_token` nor a `join_key`, rather than repeatedly presenting an empty credential.
+
 ## [v1.4.0] - 2026-08-05
 
 Implements **Protocol v1.2.0**. See `PROTOCOL.md` §1.1, §5.1–5.3.

PROTOCOL.md

@@ -12,6 +12,27 @@ The agent establishes a persistent outbound WebSocket connection.
 
 ### 1.1 Enrollment (changed in v1.2.0)
 
+Two credentials can appear in `agent.yml`. The agent presents `auth_token` when
+it has one, otherwise `join_key`:
+
+| Field | Meaning |
+| :--- | :--- |
+| `auth_token` | This agent's own token, issued by the server. Long-term identity. |
+| `join_key` | Bootstrap credential (`tjk_…`), exchanged for an `auth_token` on first connect. |
+
+**Join-key flow.** The agent connects presenting a join key and
+`?hostname=<its hostname>`. The server enrolls the host and answers with a
+`config` frame carrying `enrolled: true`, `auth_token` and `public_key`. The
+agent writes both into `agent.yml`, blanks `join_key`, and uses its own token
+from then on. This is what makes "install the agent with a key" sufficient to
+add a host — no value has to be copied between two machines by hand.
+
+The public key is accepted on first connect (trust on first use) over the same
+channel that issued the token. Pre-register the host instead if you need the
+trust anchor pinned out of band.
+
+
+
 The token **must be issued by the server**. An administrator enrolls the agent in
 the SSO (Directory → Agents, or `POST /api/agent/enroll`), which mints the token,
 stores only its SHA-256, and displays the raw value once. That value goes into
@@ -28,7 +49,7 @@ authentication failure arrives as a **close frame**, not an HTTP status:
 
 | Code | Meaning | Agent behaviour |
 | :--- | :--- | :--- |
-| `4001` | Token unknown, or not issued by this server | Back off (5 min); the credential will not fix itself |
+| `4001` | Credential unknown — neither an issued token nor a valid join key | Back off (5 min); the credential will not fix itself |
 | `4002` | Superseded — another connection authenticated as this agent | Normal reconnect |
 | `4003` | Enrollment revoked or deleted by an administrator | Back off (5 min) |
 | `4004` | Token rotated — `agent.yml` holds the superseded value | Back off (5 min); re-copy the token |
@@ -97,6 +118,20 @@ Sent in response to any command received from the server.
 
 ## 4. Server $\rightarrow$ Client Messages
 
+### 4.0 `config`
+
+Sent immediately on a successful connection.
+
+- **Type**: `config`
+- **Payload**:
+  - `message`: (string) human-readable greeting.
+  - `protocol_version`: (string) the server's protocol version.
+  - `agent_id`: (string) this agent's id in the SSO.
+  - `enrolled`: (bool, optional) present and `true` only when this connection
+    just enrolled via a join key.
+  - `auth_token`: (string, optional) the issued per-agent token — **persist it**.
+  - `public_key`: (string, optional) the key to pin — **persist it**.
+
 ### 4.1 Standard Commands
 These commands are executed if the corresponding capability is enabled in `agent.yml`.
 

agent.yml.example

@@ -3,18 +3,26 @@
 
 server_url: "https://sso.example.com"
 
-# Issued by the SSO when you enroll this host (Directory -> Install Agent, or
-# POST /api/agent/enroll). The server records only its hash and rejects any
-# token it did not issue, so a value invented locally will not connect.
-auth_token: "REPLACE_WITH_ISSUED_AGENT_TOKEN"
+# This agent's own token. Leave EMPTY when installing with a join key -- the
+# agent fills it in itself once the SSO enrolls it. The server records only a
+# hash and rejects any token it did not issue, so a locally invented value will
+# never connect.
+auth_token: ""
 
-# Base64 of the SSO's RAW 32-byte Ed25519 public key -- the `publicKey` value
-# from enrollment or GET /api/agent/nodes. This is NOT a PEM body.
+# The one credential you need to add a host. Used only while auth_token is
+# empty: the SSO exchanges it for this agent's own token + public key on first
+# connect, and the agent then blanks this line. Get one from the SSO
+# (Directory -> Install Agent, or POST /api/agent/join-keys).
+join_key: ""
+
+# Base64 of the SSO's RAW 32-byte Ed25519 public key (NOT a PEM body). Filled in
+# automatically when enrolling with a join key; set it by hand only if you
+# pre-registered this host.
 #
 # Required for any high-risk command. Without it the agent still reports
 # telemetry, but REFUSES reboot / service_restart / configure_ldap /
 # arbitrary_bash / update_binary, because it has no way to verify them.
-public_key: "REPLACE_WITH_SSO_PUBLIC_KEY"
+public_key: ""
 
 location: "default" # Location identifier (e.g., site, datacenter) for naming
 

config.go

@@ -3,6 +3,8 @@ package main
 import (
 	"fmt"
 	"os"
+	"regexp"
+	"strings"
 	"sync"
 
 	"gopkg.in/yaml.v3"
@@ -17,17 +19,31 @@ type Capabilities struct {
 }
 
 type Config struct {
-	ServerURL    string       `yaml:"server_url"`
-	AuthToken    string       `yaml:"auth_token"`
+	ServerURL string `yaml:"server_url"`
+	AuthToken string `yaml:"auth_token"`
+	// A join key is the one credential an operator hands out. On first connect
+	// the server exchanges it for a per-agent AuthToken (written back to this
+	// file), so it is a bootstrap value, not a long-term credential. Used only
+	// when AuthToken is empty.
+	JoinKey      string       `yaml:"join_key"`
 	Location     string       `yaml:"location"`
 	PublicKey    string       `yaml:"public_key"` // Ed25519 public key for signed commands
 	Capabilities Capabilities `yaml:"capabilities"`
 }
 
+// Credential returns the value to present when connecting: our own token once
+// enrolled, otherwise the join key.
+func (c *Config) Credential() string {
+	if c.AuthToken != "" {
+		return c.AuthToken
+	}
+	return c.JoinKey
+}
+
 // ConfigManager handles thread-safe access and reloading of the agent configuration.
 type ConfigManager struct {
-	mu       sync.RWMutex
-	current  *Config
+	mu         sync.RWMutex
+	current    *Config
 	configPath string
 }
 
@@ -61,6 +77,62 @@ func (cm *ConfigManager) Reload() error {
 	return nil
 }
 
+// PersistEnrollment writes the credentials the server issued during join-key
+// enrollment back into agent.yml, then reloads. Only the auth_token and
+// public_key lines are rewritten (added if absent); every other line, including
+// operator comments and the capability matrix, is preserved -- this file is
+// hand-edited, so a naive marshal-and-write would destroy it.
+//
+// The join key is blanked once we hold our own token: leaving a fleet-wide
+// credential on every host after it has stopped being needed is exactly the
+// blast radius the per-agent token exists to avoid.
+func (cm *ConfigManager) PersistEnrollment(token, publicKey string) error {
+	if token == "" {
+		return fmt.Errorf("server reported enrollment but sent no token")
+	}
+
+	cm.mu.Lock()
+	defer cm.mu.Unlock()
+
+	raw, err := os.ReadFile(cm.configPath)
+	if err != nil {
+		return fmt.Errorf("read %s: %w", cm.configPath, err)
+	}
+
+	out := setYamlScalar(string(raw), "auth_token", token)
+	if publicKey != "" {
+		out = setYamlScalar(out, "public_key", publicKey)
+	}
+	out = setYamlScalar(out, "join_key", "")
+
+	// Same permissions the installer sets: this file now holds a credential.
+	if err := os.WriteFile(cm.configPath, []byte(out), 0600); err != nil {
+		return fmt.Errorf("write %s: %w", cm.configPath, err)
+	}
+
+	cfg, err := LoadConfig(cm.configPath)
+	if err != nil {
+		return fmt.Errorf("reload after enrollment: %w", err)
+	}
+	cm.current = cfg
+	return nil
+}
+
+// setYamlScalar replaces the value of a top-level `key: "..."` line, or appends
+// the key when it is absent. Deliberately line-based rather than a YAML
+// round-trip so comments and formatting survive.
+func setYamlScalar(doc, key, value string) string {
+	re := regexp.MustCompile(`(?m)^[ \t]*` + regexp.QuoteMeta(key) + `[ \t]*:.*$`)
+	line := fmt.Sprintf("%s: %q", key, value)
+	if re.MatchString(doc) {
+		return re.ReplaceAllString(doc, line)
+	}
+	if !strings.HasSuffix(doc, "\n") {
+		doc += "\n"
+	}
+	return doc + line + "\n"
+}
+
 func LoadConfig(path string) (*Config, error) {
 	file, err := os.Open(path)
 	if err != nil {

config_test.go

@@ -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")
+	}
+}

install.sh

@@ -50,6 +50,7 @@ install_sssd_deps() {
 # 2. Argument Parsing
 URL=""
 TOKEN=""
+JOIN_KEY=""
 PUBLIC_KEY=""
 B64_CONFIG=""
 INSTALL_SSSD=0
@@ -72,6 +73,13 @@ while [[ $# -gt 0 ]]; do
       PUBLIC_KEY="$2"
       shift 2
       ;;
+    # The one credential an operator hands out. The server exchanges it for a
+    # per-agent token on first connect, which the agent writes back into
+    # agent.yml -- so this is all you need to add a host.
+    --join-key)
+      JOIN_KEY="$2"
+      shift 2
+      ;;
     --install-sssd|--ldap)
       INSTALL_SSSD=1
       shift
@@ -84,14 +92,16 @@ while [[ $# -gt 0 ]]; do
 done
 
 # Validation
-if [ -z "$B64_CONFIG" ] && [ -z "$URL" ] || [ -z "$B64_CONFIG" ] && [ -z "$TOKEN" ]; then
-  error "Missing required configuration. Either provide a base64 encoded config, or both --url and --token."
+if [ -z "$B64_CONFIG" ] && { [ -z "$URL" ] || { [ -z "$TOKEN" ] && [ -z "$JOIN_KEY" ]; }; }; then
+  error "Missing required configuration. Provide a base64 encoded config, or --url with either --join-key or --token."
   echo "Usage examples:"
   echo "  sh install.sh \"BASE64_CONFIG\""
-  echo "  sh install.sh --url \"https://sso.local\" --token \"ISSUED_TOKEN\" --public-key \"BASE64_KEY\" --install-sssd"
+  echo "  sh install.sh --url \"https://sso.local\" --join-key \"tjk_...\" --install-sssd"
+  echo "  sh install.sh --url \"https://sso.local\" --token \"ISSUED_TOKEN\" --public-key \"BASE64_KEY\""
   echo ""
-  echo "The token must be issued by the SSO (Directory -> Install Agent enrolls"
-  echo "the host and mints it). Tokens the server did not issue are rejected."
+  echo "--join-key is the normal path: the host enrolls itself on first connect"
+  echo "and the SSO issues it its own token + public key, which the agent writes"
+  echo "back into agent.yml. Get a key from Directory -> Install Agent."
   exit 1
 fi
 
@@ -116,6 +126,7 @@ else
   cat <<EOF > "$CONFIG_FILE"
 server_url: "$URL"
 auth_token: "$TOKEN"
+join_key: "$JOIN_KEY"
 public_key: "$PUBLIC_KEY"
 location: "unknown"
 capabilities:

websocket.go

@@ -121,7 +121,21 @@ func connectWebSocket(cm *ConfigManager, exec Executor) {
 			log.Fatalf("Invalid ServerURL: %v", err)
 		}
 		u.Path = "/api/agent/ws"
-		u.RawQuery = "token=" + url.QueryEscape(cfg.AuthToken)
+		// Our own token once enrolled, else the join key. The hostname lets the
+		// server name a self-enrolling host something meaningful instead of a
+		// generated placeholder.
+		q := url.Values{}
+		q.Set("token", cfg.Credential())
+		if hn, err := os.Hostname(); err == nil && hn != "" {
+			q.Set("hostname", hn)
+		}
+		u.RawQuery = q.Encode()
+
+		if cfg.Credential() == "" {
+			log.Printf("No auth_token or join_key in %s -- nothing to authenticate with. Retrying in %s.", cm.configPath, authRetryInterval)
+			time.Sleep(authRetryInterval)
+			continue
+		}
 
 		// Never log u.String(): RawQuery carries the auth token, and agent logs
 		// are routinely shipped around and pasted into issues.
@@ -292,6 +306,24 @@ func handleCommand(cm *ConfigManager, msg WSMessage, c MessageWriter, exec Execu
 		sendResponse("ok", "update applied successfully; restarting agent...")
 		os.Exit(0)
 	case "config":
+		// A config frame carrying credentials means the server accepted our
+		// join key and enrolled this host. Persist what it issued -- our own
+		// per-agent token and the public key to pin -- so the next connection
+		// authenticates as this agent rather than re-enrolling, and so signed
+		// commands can be verified. This is what lets an install ship with only
+		// a join key and still end up fully configured.
+		if enrolled, _ := msg.Payload["enrolled"].(bool); enrolled {
+			token, _ := msg.Payload["auth_token"].(string)
+			pubKey, _ := msg.Payload["public_key"].(string)
+			if err := cm.PersistEnrollment(token, pubKey); err != nil {
+				log.Printf("Enrolled, but could not persist credentials: %v", err)
+				log.Printf("This agent will re-enroll on every reconnect until %s is writable.", cm.configPath)
+			} else {
+				log.Printf("Enrolled with the SSO. Credentials written to %s; the join key is no longer needed.", cm.configPath)
+			}
+			sendResponse("ok", "enrollment stored")
+			return
+		}
 		log.Printf("Received config payload: %v", msg.Payload)
 		sendResponse("ok", "Configuration received")
 	case "reboot":