Merge pull request #4 from theta42/fix/agent-enrollment-auth

sec: fail-closed verification + server-issued enrollment (v1.4.0)

CHANGELOG.md

@@ -0,0 +1,52 @@
+# Changelog
+
+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.4.0] - 2026-08-05
+
+Implements **Protocol v1.2.0**. See `PROTOCOL.md` §1.1, §5.1–5.3.
+
+### Security
+- **Fail-closed signature verification.** `verifySignature` returned `true` when no `public_key` was configured, logging "skipping signature verification". Combined with an installer that never wrote a `public_key`, that meant a default install would execute `reboot`, `service_restart`, `configure_ldap`, `arbitrary_bash` and `update_binary` **unverified** from anything that could reach its socket. An agent that cannot verify a high-risk command now refuses it.
+- **The token must be issued by the server.** The SSO now rejects tokens it did not mint (close code `4001`). Agents carrying a token generated by the old browser-side installer will not connect until re-enrolled.
+
+### Fixed
+- **Canonicalization mismatch broke signatures for most real scripts.** Go's `encoding/json` escapes `<`, `>` and `&` by default; the server's `JSON.stringify` does not. Any payload containing them — an `arbitrary_bash` script using `>` redirection or `&&`, which is most of them — hashed differently on each side and failed verification. `canonicalize()` now uses `json.Encoder` with `SetEscapeHTML(false)` and trims the encoder's trailing newline.
+- **Auth failures no longer hot-loop.** A rejected credential was retried every 5 seconds forever, flooding the SSO and its audit log. Close codes `4001`/`4003`/`4004` now back off for 5 minutes and log what to do about it.
+- **The auth token no longer appears in logs.** The connect line logged the full URL, including `?token=...`. It now logs only host + path, and the token is URL-escaped.
+
+### Added
+- Close-code handling for the SSO's enrollment signals: `4001` unauthorized, `4002` superseded, `4003` revoked, `4004` token rotated.
+- `install.sh --public-key <base64>`, written into the generated `agent.yml`. The installer warns loudly when no public key is configured, since such an agent can report telemetry but will refuse every high-risk command.
+- Tests: fail-closed with no key, wrong key, payload tampered after signing, shell metacharacters (`>`, `&&`, `<`) round-tripping, and canonical-form equality with the server. `interop_check_test.go` verifies a signature produced by the live SSO against the agent's own verifier (skipped unless `INTEROP_FIXTURE` is set).
+
+### Changed
+- Existing tests no longer rely on verification being skipped; high-risk cases now sign with a real test key.
+- `agent.yml.example`, `README.md`, `INSTALL.md`: enrollment is a prerequisite, and `public_key` is the base64 of the **raw 32-byte** Ed25519 key — not a PEM body. The previous documented example (`MCowBQYDK2VwAyEA...`) decodes to 44 bytes and would have been rejected.
+
+## [v1.3.0] - 2026-08-04
+
+### Fixed
+- **Silently ignore `heartbeat_ack`** — the server replies to the agent's own periodic heartbeat with `heartbeat_ack`. The agent had no case for it, so it fell through to the unknown-command handler, logged `Unknown command type: heartbeat_ack` every minute, and answered with a spurious error response. Heartbeat acks are fire-and-forget; the agent now ignores them silently.
+
+## [v1.2.0] - 2026-08-03
+
+### Added
+- **Protocol v1.1.0 Compliance**: Full alignment with `PROTOCOL.md` (v1.1.0) specification.
+- **Ed25519 Cryptographic Verification**: Verification of Ed25519 Base64 signatures for high-risk C2 commands (`reboot`, `service_restart`, `configure_ldap`, `arbitrary_bash`, `update_binary`).
+- **Enhanced Journal Log Fetcher (`fetch_logs`)**: Support for querying service-specific systemd logs (`service` parameter) with configurable line count (`lines` parameter).
+- **Pure Go Self-Update Engine (`update_binary`)**: Replaced shell script execution with pure Go HTTP client fetching, SHA256 verification, atomic file replacement, and clean daemon restart.
+
+### Fixed
+- **Goroutine Leak Prevention**: Added `stopCh` lifecycle management to terminate background telemetry and heartbeat tickers upon WebSocket disconnection.
+- **Dynamic Config Rerenders**: Resolved data races when toggling capabilities or reloading `agent.yml` via `reload_config`.
+- **Config Path Uniformity**: Standardized canonical configuration file path across codebase, installer, and documentation to `/etc/theta42/agent.yml`.
+
+## [v0.1.0] - 2026-08-01
+
+### Added
+- Initial release of `theta-agent` Go daemon replacing legacy bash metric scripts.
+- Persistent outbound WebSocket telemetry and local capability matrix enforcement.

INSTALL.md

@@ -2,6 +2,18 @@
 
 Theta Agent is designed for rapid deployment across the fleet. The recommended method is via the "One-Liner" install, which marries the agent to a specific SSO Manager instance.
 
+## Prerequisite: enroll the host
+
+The agent's token is issued by the SSO, not chosen by you. In the SSO open
+**Directory → Install Agent**, name the host, bind it to a host resource, and
+press **Enroll & issue token**. You get:
+
+- the **agent token** — shown once; only its hash is stored
+- the **SSO public key** — pinned by the agent to verify high-risk commands
+
+The modal builds the install command below with both already filled in. A token
+the SSO did not issue is rejected at connect time with close code `4001`.
+
 ## Quick Start (The One-Liner)
 
 The SSO Manager provides a pre-generated installation command. Copy and paste it into your terminal as root:
@@ -15,7 +27,8 @@ curl -fsSL https://sso.example.com/resources/theta-agent/install.sh | sh -s -- "
 ### Option B: Minimal Setup
 Use this for rapid deployment with basic telemetry:
 ```bash
-curl -fsSL https://sso.example.com/resources/theta-agent/install.sh | sh -s -- --url "https://sso.example.com" --token "your-host-token"
+curl -fsSL https://sso.example.com/resources/theta-agent/install.sh | sh -s -- \
+  --url "https://sso.example.com" --token "<ISSUED_TOKEN>" --public-key "<BASE64_PUBLIC_KEY>"
 ```
 
 ### What this does:

PROTOCOL.md

@@ -1,4 +1,4 @@
-# Theta Agent Protocol Specification (v1.1.0)
+# Theta Agent Protocol Specification (v1.2.0)
 
 This document defines the communication protocol between the `theta-agent` (Client) and the `sso-manager` (Server).
 
@@ -7,8 +7,34 @@ This document defines the communication protocol between the `theta-agent` (Clie
 The agent establishes a persistent outbound WebSocket connection.
 
 - **Endpoint**: `wss://<manager-url>/api/agent/ws`
-- **Authentication**: The agent must provide a unique host token as a query parameter:
-  - `wss://<manager-url>/api/agent/ws?token=<HOST_TOKEN>`
+- **Authentication**: The agent must provide its enrollment token as a query parameter:
+  - `wss://<manager-url>/api/agent/ws?token=<AGENT_TOKEN>`
+
+### 1.1 Enrollment (changed in v1.2.0)
+
+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
+`auth_token` in `agent.yml`.
+
+Up to v1.1.0 the token was generated in the browser and never recorded
+server-side, so the server accepted *any* string: anyone who could reach
+`/api/agent/ws` could register as a node, publish discovery/telemetry, and
+receive commands addressed to a token they guessed. Tokens the server did not
+issue are now rejected.
+
+The server accepts the WebSocket upgrade before authenticating, so an
+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 |
+| `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 |
+
+Revocation and rotation both drop any live socket immediately, so they take
+effect without waiting for the agent to reconnect.
 
 ## 2. Message Format
 
@@ -97,9 +123,50 @@ These commands **require** an Ed25519 signature in the payload. The agent verifi
 
 To send a high-risk command:
 1. Create the payload (e.g., `{"script": "uptime"}`).
-2. Canonicalize the JSON (sort keys alphabetically, remove whitespace).
+2. Canonicalize the JSON (see 5.1).
 3. Sign the canonical bytes using the private Ed25519 key.
 4. Add the base64 signature to the payload: `{"script": "uptime", "signature": "..."}`.
 5. Send as a `WSMessage`.
 
 The agent performs the reverse process to verify authenticity before execution.
+
+### 5.1 Canonical form
+
+Both sides must produce **byte-identical** input to sign/verify:
+
+- keys sorted alphabetically
+- no insignificant whitespace
+- the `signature` key omitted
+- **no HTML escaping** — `<`, `>` and `&` are emitted literally
+- no trailing newline
+
+The escaping rule is load-bearing. Go's `encoding/json` escapes those three
+characters by default while JavaScript's `JSON.stringify` does not, so a payload
+containing any of them hashed differently on each side and verification failed.
+For `arbitrary_bash` that is most real scripts (`>` redirection, `&&`). The Go
+client uses `json.Encoder` with `SetEscapeHTML(false)`.
+
+Example — payload `{"script": "echo a > b && c", "comment": "x&y"}` canonicalizes to:
+
+```
+{"comment":"x&y","script":"echo a > b && c"}
+```
+
+### 5.2 The server signing key (changed in v1.2.0)
+
+The server's Ed25519 key pair is **persistent**, stored in OpenBao at
+`secret/agent/signing-key`. `public_key` in `agent.yml` is the base64-encoded raw
+32-byte public key, available from the enrollment response or
+`GET /api/agent/nodes`.
+
+Previously the pair was generated in memory at process start, so it changed on
+every restart and no agent could meaningfully pin it. If the server cannot load
+or persist a key it now **refuses to send high-risk commands** rather than
+signing with a key no agent has seen.
+
+### 5.3 Agent-side verification is fail-closed (changed in v1.2.0)
+
+An agent with no `public_key` configured **rejects** every high-risk command.
+Until v1.1.0 it logged "skipping signature verification" and executed them,
+which meant an agent installed without a key would run `reboot`,
+`configure_ldap` and `arbitrary_bash` from anything that reached its socket.

README.md

@@ -39,6 +39,17 @@ The agent will **only** execute commands that are explicitly enabled in its loca
 ### Cryptographic Hardening
 All high-risk commands require an Ed25519 signature. The agent verifies the signature against the `public_key` provided in the local config. If the signature is missing or invalid, the command is rejected regardless of the capability matrix.
 
+Verification is **fail-closed**: an agent with no `public_key` configured rejects
+every high-risk command. (Before protocol v1.2.0 it logged "skipping signature
+verification" and executed them, so an agent installed without a key would run
+`reboot`, `configure_ldap` and `arbitrary_bash` unverified.)
+
+### Enrollment
+The agent's token must be **issued by the SSO**. The server stores only its
+SHA-256 and rejects anything else at the WebSocket handshake, so a token cannot
+be minted client-side, and an enrollment can be revoked or rotated centrally —
+either drops the agent's live connection immediately. See `PROTOCOL.md` §1.1.
+
 ### Capability Matrix
 
 | Capability | Risk Level | Description | Impact |
@@ -56,7 +67,7 @@ Configuration is stored in YAML format at `/etc/theta42/agent.yml`.
 ### Example `agent.yml`
 ```yaml
 server_url: "wss://sso.theta42.local"
-auth_token: "your-unique-host-token"
+auth_token: "issued-by-the-sso-at-enrollment"
 public_key: "base64-encoded-ed25519-public-key"
 location: "dc-01-rack-12"
 capabilities:
@@ -71,7 +82,12 @@ capabilities:
 
 1. **Build**: Compile for your target architecture (see CI/CD artifacts).
 2. **Deploy**: Place the binary in `/usr/local/bin/theta-agent`.
-3. **Configure**: Create `/etc/theta42/agent.yml` with the required token and capabilities.
+3. **Enroll**: In the SSO, open **Directory → Install Agent**, name the host, bind
+   it to a host resource, and press **Enroll & issue token**. The SSO mints the
+   token (shown once) and gives you its public key. Tokens the server did not
+   issue are rejected.
+4. **Configure**: Create `/etc/theta42/agent.yml` with the issued `auth_token`,
+   the SSO's `public_key`, and your capabilities.
 4. **Service**: Set up as a systemd unit (example: `/etc/systemd/system/theta-agent.service`).
 
 For the fastest deployment, use the installation script:
@@ -79,6 +95,14 @@ For the fastest deployment, use the installation script:
 curl -fsSL https://sso.example.com/resources/theta-agent/install.sh | sh -s -- "BASE64_ENCODED_CONFIG"
 ```
 
+The **Install Agent** modal generates that command for you after enrollment,
+with the token and public key already embedded. The equivalent flag form is:
+
+```bash
+curl -fsSL https://sso.example.com/resources/theta-agent/install.sh | sh -s -- \
+  --url "https://sso.example.com" --token "<ISSUED_TOKEN>" --public-key "<BASE64_PUBLIC_KEY>"
+```
+
 ## Development & Testing
 
 The agent uses a decoupled execution engine for safety and testability.

agent.yml.example

@@ -2,7 +2,20 @@
 # Default location: /etc/theta42/agent.yml
 
 server_url: "https://sso.example.com"
-auth_token: "REPLACE_WITH_AGENT_TOKEN"
+
+# 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"
+
+# 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.
+#
+# 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"
+
 location: "default" # Location identifier (e.g., site, datacenter) for naming
 
 capabilities:

install.sh

@@ -23,10 +23,36 @@ if [ "$EUID" -ne 0 ]; then
   error "This script must be run as root."
 fi
 
+# Install SSSD and PAM integration packages if missing
+install_sssd_deps() {
+  if ! command -v sssd >/dev/null 2>&1; then
+    log "Installing SSSD and PAM integration dependencies..."
+    if command -v apt-get >/dev/null 2>&1; then
+      DEBIAN_FRONTEND=noninteractive apt-get update -qq || true
+      DEBIAN_FRONTEND=noninteractive apt-get install -y -qq sssd sssd-ldap libnss-sss libpam-sss libsss-sudo pam-auth-update || true
+      if command -v pam-auth-update >/dev/null 2>&1; then
+        pam-auth-update --enable mkhomedir || true
+      fi
+    elif command -v dnf >/dev/null 2>&1; then
+      dnf install -y sssd sssd-ldap sssd-tools || true
+    elif command -v yum >/dev/null 2>&1; then
+      yum install -y sssd sssd-ldap sssd-tools || true
+    elif command -v pacman >/dev/null 2>&1; then
+      pacman -S --noconfirm sssd || true
+    elif command -v zypper >/dev/null 2>&1; then
+      zypper in -y sssd || true
+    fi
+  else
+    log "SSSD is already installed."
+  fi
+}
+
 # 2. Argument Parsing
 URL=""
 TOKEN=""
+PUBLIC_KEY=""
 B64_CONFIG=""
+INSTALL_SSSD=0
 
 while [[ $# -gt 0 ]]; do
   case $1 in
@@ -38,6 +64,18 @@ while [[ $# -gt 0 ]]; do
       TOKEN="$2"
       shift 2
       ;;
+    # Base64 of the SSO's raw Ed25519 public key. The agent verifies high-risk
+    # commands (reboot, configure_ldap, arbitrary_bash, update_binary) against
+    # it and REFUSES them when it is absent, so an install without this key can
+    # stream telemetry but cannot be acted on.
+    --public-key)
+      PUBLIC_KEY="$2"
+      shift 2
+      ;;
+    --install-sssd|--ldap)
+      INSTALL_SSSD=1
+      shift
+      ;;
     *)
       B64_CONFIG="$1"
       shift
@@ -50,7 +88,10 @@ if [ -z "$B64_CONFIG" ] && [ -z "$URL" ] || [ -z "$B64_CONFIG" ] && [ -z "$TOKEN
   error "Missing required configuration. Either provide a base64 encoded config, or both --url and --token."
   echo "Usage examples:"
   echo "  sh install.sh \"BASE64_CONFIG\""
-  echo "  sh install.sh --url \"https://sso.local\" --token \"secret-token\""
+  echo "  sh install.sh --url \"https://sso.local\" --token \"ISSUED_TOKEN\" --public-key \"BASE64_KEY\" --install-sssd"
+  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."
   exit 1
 fi
 
@@ -75,6 +116,7 @@ else
   cat <<EOF > "$CONFIG_FILE"
 server_url: "$URL"
 auth_token: "$TOKEN"
+public_key: "$PUBLIC_KEY"
 location: "unknown"
 capabilities:
   telemetry: true
@@ -86,6 +128,20 @@ EOF
 fi
 chmod 600 "$CONFIG_FILE"
 
+# An agent with no public_key cannot verify signed commands and will refuse
+# every one of them. That is the safe default, but it is silent at run time, so
+# say it plainly here where the operator is watching.
+if ! grep -qE '^public_key:[[:space:]]*"[^"]+"' "$CONFIG_FILE" 2>/dev/null; then
+  log "WARNING: no public_key configured — this agent will report telemetry but"
+  log "         REFUSE reboot / configure_ldap / arbitrary_bash / update_binary."
+  log "         Re-run with --public-key \"<base64 key>\" (shown at enrollment)."
+fi
+
+# 4b. Ensure SSSD dependencies are installed if configure_ldap is enabled
+if [ "$INSTALL_SSSD" -eq 1 ] || grep -q -i "configure_ldap:\s*true" "$CONFIG_FILE" 2>/dev/null; then
+  install_sssd_deps
+fi
+
 # 5. Setup systemd service
 log "Creating systemd service unit..."
 cat <<EOF > "$SERVICE_FILE"

interop_check_test.go

@@ -0,0 +1,36 @@
+package main
+
+import (
+	"encoding/json"
+	"os"
+	"testing"
+)
+
+// Cross-implementation check: a payload signed by the Node server (utils/
+// agent_manager.js) must verify with the agent's own verifySignature. Skips
+// unless the fixture is present, so it never breaks a normal `go test`.
+func TestInteropWithServerSignature(t *testing.T) {
+	raw, err := os.ReadFile(os.Getenv("INTEROP_FIXTURE"))
+	if err != nil {
+		t.Skip("no INTEROP_FIXTURE provided")
+	}
+	var fx struct {
+		Pub     string                 `json:"pub"`
+		Payload map[string]interface{} `json:"payload"`
+		Sig     string                 `json:"sig"`
+	}
+	if err := json.Unmarshal(raw, &fx); err != nil {
+		t.Fatalf("bad fixture: %v", err)
+	}
+	payload := map[string]interface{}{}
+	for k, v := range fx.Payload {
+		payload[k] = v
+	}
+	payload["signature"] = fx.Sig
+
+	cfg := &Config{PublicKey: fx.Pub}
+	if !verifySignature(cfg, WSMessage{Type: "arbitrary_bash", Payload: payload}) {
+		t.Fatal("agent REJECTED a signature produced by the SSO server")
+	}
+	t.Log("agent accepted the server-produced signature")
+}

telemetry.go

@@ -114,38 +114,57 @@ func collectGPUUsage(exec Executor) float64 {
 }
 
 // StartTelemetryLoop manages the initial discovery push and the periodic telemetry stream.
-func StartTelemetryLoop(c *websocket.Conn, cfg *Config, exec Executor) {
+func StartTelemetryLoop(c MessageWriter, cm *ConfigManager, exec Executor, stopCh <-chan struct{}) {
+	cfg := cm.Get()
+
 	// 1. Immediate Discovery Push
 	pushDiscovery(c, cfg)
 
+	// If telemetry capability is disabled in agent.yml, return early after discovery
+	if !cfg.Capabilities.Telemetry {
+		log.Println("Telemetry capability is disabled in agent.yml; skipping telemetry stream.")
+		return
+	}
+
 	// 2. Periodic Telemetry Stream
 	ticker := time.NewTicker(30 * time.Second)
 	go func() {
+		defer ticker.Stop()
 		var lastIPs []string
-		for range ticker.C {
-			// Network Change Detection
-			currentIPs := collectIPs()
-			if !equalSlices(lastIPs, currentIPs) {
-				log.Println("Network change detected. Pushing discovery update...")
-				pushDiscovery(c, cfg)
-				lastIPs = currentIPs
-			}
-
-			telemetry := CollectTelemetryData(exec)
-			payload, _ := json.Marshal(WSMessage{
-				Type: "telemetry",
-				Payload: map[string]interface{}{
-					"cpu_usage_percent":  telemetry.CPUUsagePercent,
-					"ram_usage_percent":  telemetry.RAMUsagePercent,
-					"disk_usage_percent": telemetry.DiskUsagePercent,
-					"zfs_health":         telemetry.ZFSHealth,
-					"gpu_usage_percent":   telemetry.GPUUsage,
-					"timestamp":          telemetry.Timestamp,
-				},
-			})
-			if err := c.WriteMessage(websocket.TextMessage, payload); err != nil {
-				log.Printf("Failed to stream telemetry: %v", err)
+		for {
+			select {
+			case <-stopCh:
 				return
+			case <-ticker.C:
+				currentCFG := cm.Get()
+				if !currentCFG.Capabilities.Telemetry {
+					continue
+				}
+
+				// Network Change Detection
+				currentIPs := collectIPs()
+				if !equalSlices(lastIPs, currentIPs) {
+					log.Println("Network change detected. Pushing discovery update...")
+					pushDiscovery(c, currentCFG)
+					lastIPs = currentIPs
+				}
+
+				telemetry := CollectTelemetryData(exec)
+				payload, _ := json.Marshal(WSMessage{
+					Type: "telemetry",
+					Payload: map[string]interface{}{
+						"cpu_usage_percent":  telemetry.CPUUsagePercent,
+						"ram_usage_percent":  telemetry.RAMUsagePercent,
+						"disk_usage_percent": telemetry.DiskUsagePercent,
+						"zfs_health":         telemetry.ZFSHealth,
+						"gpu_usage_percent":   telemetry.GPUUsage,
+						"timestamp":          telemetry.Timestamp,
+					},
+				})
+				if err := c.WriteMessage(websocket.TextMessage, payload); err != nil {
+					log.Printf("Failed to stream telemetry: %v", err)
+					return
+				}
 			}
 		}
 	}()
@@ -176,7 +195,7 @@ func equalSlices(a, b []string) bool {
 	return true
 }
 
-func pushDiscovery(c *websocket.Conn, cfg *Config) {
+func pushDiscovery(c MessageWriter, cfg *Config) {
 	discovery := CollectDiscoveryData(cfg)
 	discoveryPayload, _ := json.Marshal(discovery)
 

websocket.go

@@ -1,12 +1,18 @@
 package main
 
 import (
+	"bytes"
 	"crypto/ed25519"
+	"crypto/sha256"
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
+	"io"
 	"log"
+	"net/http"
 	"net/url"
+	"os"
+	"path/filepath"
 	"strings"
 	"time"
 
@@ -18,14 +24,55 @@ type WSMessage struct {
 	Payload map[string]interface{} `json:"payload"`
 }
 
+// Application close codes the SSO uses to say "your enrollment is the problem"
+// (PROTOCOL.md §1.1). All three mean retrying quickly is pointless.
+const (
+	closeUnauthorized = 4001 // token was never issued, or is unknown
+	closeSuperseded   = 4002 // another connection took over this enrollment
+	closeRevoked      = 4003 // enrollment revoked or deleted by an admin
+	closeTokenRotated = 4004 // token rotated; agent.yml holds the old one
+)
+
+// How long to wait before retrying after the server rejects our credential.
+// Short enough that a re-enrollment is picked up without a restart, long enough
+// that a decommissioned agent is not a permanent load on the SSO.
+const authRetryInterval = 5 * time.Minute
+
 type MessageWriter interface {
 	WriteMessage(messageType int, data []byte) error
 }
 
+// canonicalize produces the exact bytes the server signed (PROTOCOL.md §5):
+// keys sorted alphabetically, no whitespace, `signature` omitted.
+//
+// encoding/json sorts map keys for us, but by default it also escapes <, > and
+// & as <, > and & -- which Node's JSON.stringify on the server
+// does not. Any payload containing those characters therefore hashed
+// differently on each side and the signature failed. For arbitrary_bash that is
+// most real scripts: `>` redirection and `&&` are everywhere. SetEscapeHTML
+// (false) is what makes the two encoders agree.
+//
+// Encoder.Encode also appends a trailing newline, which must be trimmed or it
+// is signed-over data the server never produced.
+func canonicalize(payload map[string]interface{}) ([]byte, error) {
+	var buf bytes.Buffer
+	enc := json.NewEncoder(&buf)
+	enc.SetEscapeHTML(false)
+	if err := enc.Encode(payload); err != nil {
+		return nil, err
+	}
+	return bytes.TrimRight(buf.Bytes(), "\n"), nil
+}
+
 func verifySignature(cfg *Config, msg WSMessage) bool {
+	// Fail CLOSED. This used to return true when no public key was configured,
+	// which meant an agent installed without a `public_key` would execute
+	// reboot / configure_ldap / arbitrary_bash from anything that could reach
+	// its socket, with no verification at all -- the exact commands the
+	// signature exists to protect. An agent that cannot verify must not act.
 	if cfg.PublicKey == "" {
-		log.Println("No public key configured; skipping signature verification")
-		return true
+		log.Println("Refusing high-risk command: no public_key configured in agent.yml")
+		return false
 	}
 
 	sigB64, ok := msg.Payload["signature"].(string)
@@ -47,7 +94,11 @@ func verifySignature(cfg *Config, msg WSMessage) bool {
 			payloadCopy[k] = v
 		}
 	}
-	canonicalPayload, _ := json.Marshal(payloadCopy)
+	canonicalPayload, err := canonicalize(payloadCopy)
+	if err != nil {
+		log.Printf("Could not canonicalize payload for verification: %v", err)
+		return false
+	}
 
 	pubKeyBytes, err := base64.StdEncoding.DecodeString(cfg.PublicKey)
 	if err != nil || len(pubKeyBytes) != ed25519.PublicKeySize {
@@ -70,12 +121,22 @@ func connectWebSocket(cm *ConfigManager, exec Executor) {
 			log.Fatalf("Invalid ServerURL: %v", err)
 		}
 		u.Path = "/api/agent/ws"
-		u.RawQuery = "token=" + cfg.AuthToken
+		u.RawQuery = "token=" + url.QueryEscape(cfg.AuthToken)
 
-		log.Printf("Connecting to %s", u.String())
+		// Never log u.String(): RawQuery carries the auth token, and agent logs
+		// are routinely shipped around and pasted into issues.
+		log.Printf("Connecting to %s%s", u.Host, u.Path)
 
-		c, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
+		c, resp, err := websocket.DefaultDialer.Dial(u.String(), nil)
 		if err != nil {
+			// The server now rejects tokens it did not issue. Retrying a bad
+			// credential every 5s just floods the SSO and its audit log
+			// forever, so back off hard and say plainly what is wrong.
+			if resp != nil && (resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden) {
+				log.Printf("Server rejected our token (HTTP %d). Enroll this agent in the SSO Directory and put the issued token in agent.yml. Retrying in %s.", resp.StatusCode, authRetryInterval)
+				time.Sleep(authRetryInterval)
+				continue
+			}
 			log.Printf("Dial error: %v. Retrying in 5 seconds...", err)
 			time.Sleep(5 * time.Second)
 			continue
@@ -83,26 +144,47 @@ func connectWebSocket(cm *ConfigManager, exec Executor) {
 
 		log.Println("Successfully connected to SSO Manager.")
 
-		// Start telemetry and discovery
-		StartTelemetryLoop(c, cfg, exec)
+		stopCh := make(chan struct{})
+
+		// Start telemetry and discovery with stopCh lifecycle control
+		StartTelemetryLoop(c, cm, exec, stopCh)
 
 		// Heartbeat loop
 		go func() {
 			ticker := time.NewTicker(60 * time.Second)
-			for range ticker.C {
-				hb := WSMessage{Type: "heartbeat", Payload: map[string]interface{}{"timestamp": time.Now().Format(time.RFC3339)}}
-				payload, _ := json.Marshal(hb)
-				if err := c.WriteMessage(websocket.TextMessage, payload); err != nil {
+			defer ticker.Stop()
+			for {
+				select {
+				case <-stopCh:
 					return
+				case <-ticker.C:
+					hb := WSMessage{Type: "heartbeat", Payload: map[string]interface{}{"timestamp": time.Now().Format(time.RFC3339)}}
+					payload, _ := json.Marshal(hb)
+					if err := c.WriteMessage(websocket.TextMessage, payload); err != nil {
+						return
+					}
 				}
 			}
 		}()
 
+		// Set when the server closes us for an enrollment problem rather than a
+		// transient fault, so the reconnect below can back off instead of
+		// spinning on a credential that will not start working by itself.
+		authRejected := false
+
 		// Read loop
 		for {
 			_, message, err := c.ReadMessage()
 			if err != nil {
-				log.Println("WebSocket read error:", err)
+				// The SSO accepts the upgrade and only then closes with an
+				// application code, so an auth failure surfaces here rather
+				// than at Dial.
+				if websocket.IsCloseError(err, closeUnauthorized, closeRevoked, closeTokenRotated) {
+					authRejected = true
+					log.Printf("Server closed the connection: %v. This agent's token is not valid for that SSO — re-enroll it and update agent.yml.", err)
+				} else {
+					log.Println("WebSocket read error:", err)
+				}
 				break // break read loop, reconnect
 			}
 
@@ -116,8 +198,15 @@ func connectWebSocket(cm *ConfigManager, exec Executor) {
 		}
 
 		// Cleanup on disconnect
+		close(stopCh)
 		c.Close()
 
+		if authRejected {
+			log.Printf("Reconnecting in %s.", authRetryInterval)
+			time.Sleep(authRetryInterval)
+			continue
+		}
+
 		log.Println("WebSocket disconnected. Reconnecting in 5 seconds...")
 		time.Sleep(5 * time.Second)
 	}
@@ -125,7 +214,11 @@ func connectWebSocket(cm *ConfigManager, exec Executor) {
 
 func handleCommand(cm *ConfigManager, msg WSMessage, c MessageWriter, exec Executor) {
 	cfg := cm.Get()
-	log.Printf("Received command: %s", msg.Type)
+	// Don't log the server's fire-and-forget heartbeat ack — it arrives every
+	// 60s and is not a command to act on; logging it is pure per-minute noise.
+	if msg.Type != "heartbeat_ack" {
+		log.Printf("Received command: %s", msg.Type)
+	}
 
 	sendResponse := func(status string, message string) {
 		resp, _ := json.Marshal(map[string]string{"status": status, "message": message})
@@ -142,15 +235,33 @@ func handleCommand(cm *ConfigManager, msg WSMessage, c MessageWriter, exec Execu
 			sendResponse("ok", "configuration reloaded")
 		}
 	case "fetch_logs":
-		out, err := exec.Execute("journalctl", "-u", "theta-agent", "-n", "100")
+		serviceName, _ := msg.Payload["service"].(string)
+		if serviceName == "" {
+			serviceName = "theta-agent"
+		}
+
+		if serviceName != "theta-agent" && !cfg.Capabilities.CanManageService(serviceName) {
+			log.Printf("Fetch logs rejected for '%s': not in allowed service list", serviceName)
+			sendResponse("error", "service log fetch rejected")
+			return
+		}
+
+		linesCount := 100
+		if l, ok := msg.Payload["lines"].(float64); ok && l > 0 {
+			linesCount = int(l)
+		}
+
+		log.Printf("Fetching logs for service %s (%d lines)...", serviceName, linesCount)
+		out, err := exec.Execute("journalctl", "-u", serviceName, "-n", fmt.Sprintf("%d", linesCount), "--no-pager")
 		if err != nil {
 			log.Printf("Log fetch failed: %v", err)
 			sendResponse("error", "failed to fetch logs")
 			return
 		}
-		resp := map[string]string{
-			"status": "ok",
-			"logs":   string(out),
+		resp := map[string]interface{}{
+			"status":  "ok",
+			"service": serviceName,
+			"logs":    string(out),
 		}
 		respPayload, _ := json.Marshal(resp)
 		c.WriteMessage(websocket.TextMessage, respPayload)
@@ -160,28 +271,25 @@ func handleCommand(cm *ConfigManager, msg WSMessage, c MessageWriter, exec Execu
 			sendResponse("error", "signature verification failed")
 			return
 		}
-		if !cfg.Capabilities.ArbitraryBash { // Use Bash as a proxy for "dangerous update" capability
+		if !cfg.Capabilities.ArbitraryBash {
 			sendResponse("error", "update capability disabled")
 			return
 		}
 
-		url, _ := msg.Payload["url"].(string)
+		urlStr, _ := msg.Payload["url"].(string)
 		checksum, _ := msg.Payload["sha256"].(string)
-		if url == "" || checksum == "" {
-			sendResponse("error", "missing url or checksum")
+		if urlStr == "" || checksum == "" {
+			sendResponse("error", "missing url or sha256 checksum")
 			return
 		}
 
-		log.Printf("Updating binary from %s...", url)
-		// implementation of download and replace
-		// ... (simplified for now, using a shell command via executor for brevity in this turn)
-		script := fmt.Sprintf("curl -fsSL %s -o /tmp/theta-agent.new && sha256sum -c <(echo '%s  /tmp/theta-agent.new') && mv /tmp/theta-agent.new $(readlink -f /proc/self/exe)", url, checksum)
-		if _, err := exec.Execute("bash", "-c", script); err != nil {
+		log.Printf("Updating binary from %s...", urlStr)
+		if err := downloadAndUpdateBinary(urlStr, checksum); err != nil {
 			log.Printf("Update failed: %v", err)
-			sendResponse("error", "update failed")
+			sendResponse("error", fmt.Sprintf("update failed: %v", err))
 			return
 		}
-		sendResponse("ok", "update applied. restarting agent...")
+		sendResponse("ok", "update applied successfully; restarting agent...")
 		os.Exit(0)
 	case "config":
 		log.Printf("Received config payload: %v", msg.Payload)
@@ -282,8 +390,68 @@ func handleCommand(cm *ConfigManager, msg WSMessage, c MessageWriter, exec Execu
 		respPayload, _ := json.Marshal(resp)
 		c.WriteMessage(websocket.TextMessage, respPayload)
 		return
+	// heartbeat_ack is the server's acknowledgement of the agent's own periodic
+	// heartbeat (the agent sends `heartbeat`, the server answers `heartbeat_ack`).
+	// There is nothing to do with it -- it is not a command to run, and answering
+	// an ack with an error response would inject spurious errors into the
+	// command-response channel every minute. Silently ignore.
+	case "heartbeat_ack":
+		return
 	default:
 		log.Printf("Unknown command type: %s", msg.Type)
 		sendResponse("error", "unknown command type")
 	}
 }
+
+func downloadAndUpdateBinary(downloadURL string, expectedSHA256 string) error {
+	resp, err := http.Get(downloadURL)
+	if err != nil {
+		return fmt.Errorf("http fetch failed: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return fmt.Errorf("unexpected http status: %s", resp.Status)
+	}
+
+	tmpFile, err := os.CreateTemp("", "theta-agent-update-*")
+	if err != nil {
+		return fmt.Errorf("failed to create temp file: %w", err)
+	}
+	tmpPath := tmpFile.Name()
+	defer os.Remove(tmpPath)
+
+	hasher := sha256.New()
+	writer := io.MultiWriter(tmpFile, hasher)
+
+	if _, err := io.Copy(writer, resp.Body); err != nil {
+		tmpFile.Close()
+		return fmt.Errorf("failed to save binary: %w", err)
+	}
+	tmpFile.Close()
+
+	actualSHA256 := fmt.Sprintf("%x", hasher.Sum(nil))
+	if !strings.EqualFold(actualSHA256, strings.TrimSpace(expectedSHA256)) {
+		return fmt.Errorf("sha256 mismatch: expected %s, got %s", expectedSHA256, actualSHA256)
+	}
+
+	if err := os.Chmod(tmpPath, 0755); err != nil {
+		return fmt.Errorf("failed to set executable permissions: %w", err)
+	}
+
+	selfPath, err := os.Executable()
+	if err != nil {
+		return fmt.Errorf("failed to resolve current binary path: %w", err)
+	}
+
+	resolvedPath, err := filepath.EvalSymlinks(selfPath)
+	if err == nil {
+		selfPath = resolvedPath
+	}
+
+	if err := os.Rename(tmpPath, selfPath); err != nil {
+		return fmt.Errorf("failed to replace binary: %w", err)
+	}
+
+	return nil
+}

websocket_test.go

@@ -1,11 +1,42 @@
 package main
 
 import (
+	"crypto/ed25519"
+	"encoding/base64"
 	"encoding/json"
 	"os"
 	"testing"
 )
 
+// A fixed key pair for the tests, standing in for the SSO's persisted signing
+// key. High-risk commands must now be genuinely signed: the agent fails closed
+// when no public_key is configured, so these tests sign the way the server
+// does instead of relying on verification being skipped.
+var testPubKey, testPrivKey, _ = ed25519.GenerateKey(nil)
+
+func testPubKeyB64() string {
+	return base64.StdEncoding.EncodeToString(testPubKey)
+}
+
+// sign mirrors the server's canonicalization (sorted keys, no whitespace, no
+// HTML escaping, `signature` omitted) and adds the signature to the payload.
+func sign(t *testing.T, payload map[string]interface{}) map[string]interface{} {
+	t.Helper()
+	if payload == nil {
+		payload = map[string]interface{}{}
+	}
+	canonical, err := canonicalize(payload)
+	if err != nil {
+		t.Fatalf("canonicalize: %v", err)
+	}
+	signed := make(map[string]interface{}, len(payload)+1)
+	for k, v := range payload {
+		signed[k] = v
+	}
+	signed["signature"] = base64.StdEncoding.EncodeToString(ed25519.Sign(testPrivKey, canonical))
+	return signed
+}
+
 type MockConn struct {
 	Messages [][]byte
 }
@@ -33,15 +64,30 @@ func (m *MockExecutor) WriteFile(path string, data []byte, perm os.FileMode) err
 	return nil
 }
 
+func (m *MockExecutor) ReadFile(path string) ([]byte, error) {
+	if m.WrittenFiles != nil {
+		if data, ok := m.WrittenFiles[path]; ok {
+			return data, nil
+		}
+	}
+	return []byte("mock file content"), nil
+}
+
 func TestHandleCommand(t *testing.T) {
 	tests := []struct {
-		name            string
-		cfg             *Config
-		msg             WSMessage
-		expectedStatus  string
-		expectedCmd     []string
-		expectedFile    string
+		name             string
+		cfg              *Config
+		msg              WSMessage
+		expectedStatus   string
+		expectedCmd      []string
+		expectedFile     string
 		expectedFileCont string
+		// sign the payload with the test key before dispatch, the way the SSO
+		// signs high-risk commands
+		signed bool
+		// heartbeat_ack (and any fire-and-forget ack) must be silently ignored —
+		// no response message, no command, no log noise.
+		expectedNoResponse bool
 	}{
 		{
 			name: "config command success",
@@ -57,11 +103,13 @@ func TestHandleCommand(t *testing.T) {
 		{
 			name: "reboot command allowed",
 			cfg: &Config{
+				PublicKey:    testPubKeyB64(),
 				Capabilities: Capabilities{Reboot: true},
 			},
 			msg: WSMessage{
 				Type: "reboot",
 			},
+			signed:         true,
 			expectedStatus: "ok",
 			expectedCmd:    []string{"reboot"},
 		},
@@ -111,6 +159,7 @@ func TestHandleCommand(t *testing.T) {
 		{
 			name: "configure_ldap allowed",
 			cfg: &Config{
+				PublicKey:    testPubKeyB64(),
 				Capabilities: Capabilities{ConfigureLDAP: true},
 			},
 			msg: WSMessage{
@@ -119,10 +168,11 @@ func TestHandleCommand(t *testing.T) {
 					"config": "domain = theta42.local\nserver = sso.local",
 				},
 			},
-			expectedStatus:  "ok",
-			expectedFile:    "/etc/sssd/sssd.conf",
+			signed:           true,
+			expectedStatus:   "ok",
+			expectedFile:     "/etc/sssd/sssd.conf",
 			expectedFileCont: "domain = theta42.local\nserver = sso.local",
-			expectedCmd:     []string{"systemctl", "restart", "sssd"},
+			expectedCmd:      []string{"systemctl", "restart", "sssd"},
 		},
 		{
 			name: "configure_ldap denied",
@@ -141,6 +191,7 @@ func TestHandleCommand(t *testing.T) {
 		{
 			name: "arbitrary_bash allowed",
 			cfg: &Config{
+				PublicKey:    testPubKeyB64(),
 				Capabilities: Capabilities{ArbitraryBash: true},
 			},
 			msg: WSMessage{
@@ -149,6 +200,7 @@ func TestHandleCommand(t *testing.T) {
 					"script": "uptime",
 				},
 			},
+			signed:         true,
 			expectedStatus: "ok",
 			expectedCmd:    []string{"bash", "-c", "uptime"},
 		},
@@ -166,6 +218,16 @@ func TestHandleCommand(t *testing.T) {
 			expectedStatus: "error",
 			expectedCmd:    nil,
 		},
+		{
+			name: "heartbeat_ack is silently ignored",
+			cfg: &Config{
+				Capabilities: Capabilities{},
+			},
+			msg: WSMessage{
+				Type: "heartbeat_ack",
+			},
+			expectedNoResponse: true,
+		},
 		{
 			name: "unknown command",
 			cfg: &Config{
@@ -182,7 +244,22 @@ func TestHandleCommand(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			mockConn := &MockConn{}
 			mockExec := &MockExecutor{}
-			handleCommand(tc.cfg, tc.msg, mockConn, mockExec)
+			cm := &ConfigManager{current: tc.cfg}
+			msg := tc.msg
+			if tc.signed {
+				msg.Payload = sign(t, msg.Payload)
+			}
+			handleCommand(cm, msg, mockConn, mockExec)
+
+			if tc.expectedNoResponse {
+				if len(mockConn.Messages) != 0 {
+					t.Fatalf("expected no response message, got %d: %v", len(mockConn.Messages), mockConn.Messages)
+				}
+				if len(mockExec.ExecutedCommands) > 0 {
+					t.Errorf("expected no commands to be executed, but got %v", mockExec.ExecutedCommands)
+				}
+				return
+			}
 
 			if len(mockConn.Messages) != 1 {
 				t.Fatalf("expected 1 response message, got %d", len(mockConn.Messages))
@@ -226,3 +303,67 @@ func TestHandleCommand(t *testing.T) {
 		})
 	}
 }
+
+// The agent must not execute a high-risk command it cannot verify. This used to
+// return true when no public_key was configured, so an agent installed without
+// one executed reboot / configure_ldap / arbitrary_bash unverified.
+func TestVerifySignatureFailsClosedWithoutPublicKey(t *testing.T) {
+	cfg := &Config{} // no PublicKey
+	msg := WSMessage{Type: "arbitrary_bash", Payload: sign(t, map[string]interface{}{"script": "uptime"})}
+	if verifySignature(cfg, msg) {
+		t.Fatal("verifySignature accepted a command with no public_key configured")
+	}
+}
+
+func TestVerifySignatureRejectsWrongKey(t *testing.T) {
+	otherPub, _, _ := ed25519.GenerateKey(nil)
+	cfg := &Config{PublicKey: base64.StdEncoding.EncodeToString(otherPub)}
+	msg := WSMessage{Type: "arbitrary_bash", Payload: sign(t, map[string]interface{}{"script": "uptime"})}
+	if verifySignature(cfg, msg) {
+		t.Fatal("verifySignature accepted a signature from a different key")
+	}
+}
+
+func TestVerifySignatureRejectsTamperedPayload(t *testing.T) {
+	cfg := &Config{PublicKey: testPubKeyB64()}
+	payload := sign(t, map[string]interface{}{"script": "uptime"})
+	payload["script"] = "rm -rf /" // swap the script, keep the signature
+	if verifySignature(cfg, WSMessage{Type: "arbitrary_bash", Payload: payload}) {
+		t.Fatal("verifySignature accepted a payload modified after signing")
+	}
+}
+
+// Regression: encoding/json escapes <, > and & by default, but the server's
+// JSON.stringify does not. Any script using redirection or && therefore
+// canonicalized differently on each side and failed verification -- which is
+// most real scripts.
+func TestVerifySignatureAcceptsShellMetacharacters(t *testing.T) {
+	cfg := &Config{PublicKey: testPubKeyB64()}
+	for _, script := range []string{
+		"echo hi > /tmp//out.log",
+		"systemctl is-active nginx && systemctl reload nginx",
+		"grep -c . < /etc/passwd",
+		"a=1 && b=2 && echo \"$a<$b\" > /dev/null",
+	} {
+		msg := WSMessage{Type: "arbitrary_bash", Payload: sign(t, map[string]interface{}{"script": script})}
+		if !verifySignature(cfg, msg) {
+			t.Errorf("verifySignature rejected a correctly signed script: %q", script)
+		}
+	}
+}
+
+// The canonical form must be byte-identical to the server's: sorted keys, no
+// whitespace, no HTML escaping, no trailing newline, signature omitted.
+func TestCanonicalizeMatchesServerForm(t *testing.T) {
+	got, err := canonicalize(map[string]interface{}{
+		"script":  "echo a > b && c",
+		"comment": "x&y",
+	})
+	if err != nil {
+		t.Fatalf("canonicalize: %v", err)
+	}
+	want := `{"comment":"x&y","script":"echo a > b && c"}`
+	if string(got) != want {
+		t.Errorf("canonical form mismatch:\n got: %s\nwant: %s", got, want)
+	}
+}