diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 559e4680..d252240a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -149,12 +149,13 @@ jobs: # branch that had touched none of the reported code. Bump this # deliberately, in its own commit, with the fallout fixed alongside it. # - # 2.13.1 is also a FLOOR, not just a pin: .golangci.yml excludes + # There is also a FLOOR below this pin: .golangci.yml excludes # errors.AsType from errcheck by function name, and errcheck before - # 2.13.1 cannot resolve a generic function's name, so on an older + # 2.13.0 cannot resolve a generic function's name, so on an older # golangci-lint that exclusion silently fails to match and every - # errors.AsType call site is reported. Keep docs/development.md's - # required-version note in step with this value. + # errors.AsType call site is reported. The floor is 2.13.0 (what Homebrew + # ships, verified clean on this repo); this pin is one patch above it. + # Keep docs/development.md's required-version note in step with both. # # Explicit path patterns mirror the Makefile LINT_PKGS variable: they # exclude web/node_modules/ (third-party JS packages that happen to @@ -959,6 +960,105 @@ jobs: fi exit $missing + # ── Integration suite (whole package) ──────────────────────────────────────── + # Runs every file under test/integration/ via `make test-integration`, rather + # than naming individual tests in a -run regex the way the LDAP, OIDC, and + # isolation jobs below do. Those three name their tests because each needs + # its own runner matrix or container prerequisites; everything else in the + # package — broker auth, EXPR parameter binding and end-to-end execution, + # failure-reason propagation (both in-process and against a real worker + # binary), retry/auto-park behavior, and product retry-override submission — + # shares no such prerequisite, so a single whole-package run covers a new + # test the moment it is added instead of only if someone also remembers to + # list it in a job. `go build`s a real sqi-worker binary along the way (see + # worker_binary_test.go); no other tooling is required beyond Go itself. + # + # LDAP, OIDC, and isolation tests live in this same package and this job + # exercises them too, but ubuntu-latest ships Docker so they are expected to + # run rather than skip; the LDAP and Keycloak images are pulled explicitly + # first so a registry hiccup fails here with a clear message instead of + # inside the test as a cold-pull timeout that would read as a skip. The + # isolation tests, in contrast, require running as root (see + # isolation-integration below) and are expected to skip on this ordinary + # runner — that is fine, and this job asserts nothing about them by name. + # + # IMPORTANT: a green `go test` here proves nothing on its own if the whole + # suite silently collected zero tests (a build-tag regression, a stray + # -run filter). This job does not trust the exit code alone: it requires a + # test named from every file in the package that has no dedicated CI job of + # its own — one exception is authfixture_test.go, which holds only helpers + # shared by ldap_test.go and oidc_test.go and declares no test function of + # its own to name. A confirmed-skip line for LDAP/OIDC/isolation does not + # fail the job; only a missing PASS among the required names, or a nonzero + # exit, does. + integration-tests: + name: Integration suite + runs-on: ubuntu-latest + needs: test + + steps: + - name: Checkout + uses: actions/checkout@v7 + + - name: Set up Go + uses: actions/setup-go@v7 + with: + go-version-file: go.mod + + # Pulled explicitly for the same reason ldap-integration and + # oidc-integration below pull them: a registry problem fails here with a + # clear message rather than surfacing inside the test as a cold-pull + # timeout, which the test's own Docker-unavailable handling would read + # as an ordinary skip. + - name: Pull OpenLDAP image + run: docker pull osixia/openldap:1.5.0 + + - name: Pull Keycloak image + run: docker pull quay.io/keycloak/keycloak:26.0.7 + + - name: Run the integration suite + run: | + set +e + # go test's own default (10m per package) leaves little headroom + # once the LDAP and Keycloak containers are in the run: up to ~90s + # for LDAP readiness plus up to ~180s for Keycloak, on top of the + # suite's own runtime, which is dominated by that container + # readiness rather than by the tests themselves. 15m is set + # explicitly here, through INTEGRATION_TEST_FLAGS, so a slow runner + # gets a clear failure instead of a default timeout landing close + # enough to the estimate to be a coin flip; local runs are + # unaffected since the Makefile leaves this empty by default. + make test-integration INTEGRATION_TEST_FLAGS="-timeout 15m" 2>&1 | tee integration-output.log + status=${PIPESTATUS[0]} + set -e + + # One name from each file in test/integration/ that has no dedicated + # CI job of its own, so a build-tag or filter regression that let the + # suite silently collect zero tests cannot pass this job. + expected=( + TestDefaultConfig_NoBrokerAuth + TestRevocation_DisconnectsAndReclaims + TestEnrollment_ConnectsToRunningBrokerWithoutRestart + TestEXPRParamTypes_BindAndCarry + TestEXPRJobEndToEnd + TestTaskFailureReason_VisibleEndToEnd + TestWorkerBinaryStagingFailureReason + TestProductSubmit_RetryOverrides + TestAutoRetry_RetryThenSucceed + ) + missing=0 + for name in "${expected[@]}"; do + if ! grep -q -- "--- PASS: $name" integration-output.log; then + echo "::error::$name did not pass" + missing=1 + fi + done + if [ "$status" -ne 0 ] || [ "$missing" -ne 0 ]; then + echo "::error::the integration suite did not all pass (exit=$status)." + exit 1 + fi + echo "Confirmed all ${#expected[@]} named integration tests passed, and the whole suite exited 0." + # Runs the official OpenJD conformance suite (a pinned submodule under # third_party/) against internal/openjd. Tagged `conformance`, so it does not # run in the default `make test`. diff --git a/.golangci.yml b/.golangci.yml index e4d65499..cb82cfbf 100644 --- a/.golangci.yml +++ b/.golangci.yml @@ -121,7 +121,7 @@ linters: # and reports it. Exclude the function rather than turning check-blank # off, so a genuine `_ = f()` is still caught. # - # Requires golangci-lint >= 2.13.1: earlier errcheck cannot resolve a + # Requires golangci-lint >= 2.13.0: earlier errcheck cannot resolve a # generic function's name (it reports "Error return value is not # checked" with no name), so this entry never matches and every # errors.AsType call site is flagged. diff --git a/Makefile b/Makefile index 76fe390c..54fa0ba4 100644 --- a/Makefile +++ b/Makefile @@ -200,8 +200,12 @@ test-cover-html: test-cover ## Open HTML coverage report in the browser go tool cover -html=$(COVERAGE_OUT) .PHONY: test-integration +# INTEGRATION_TEST_FLAGS is empty by default so a local run gets go test's +# ordinary defaults; CI sets it to add a timeout sized for the suite's own +# LDAP/Keycloak containers without changing what a local run does. +INTEGRATION_TEST_FLAGS ?= test-integration: ## Run integration tests (tagged 'integration') - go test $(TEST_FLAGS) -tags integration ./test/... + go test $(TEST_FLAGS) -tags integration -v $(INTEGRATION_TEST_FLAGS) ./test/... .PHONY: test-conformance test-conformance: ## Run the official OpenJD conformance suite (needs the pinned submodule) diff --git a/ROADMAP.md b/ROADMAP.md index d53c8dcf..8775a82c 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -79,7 +79,7 @@ Configuration cascades: farm defaults → queue overrides, with retry policy (ma **Scheduling considers:** job priority, task dependencies, queue and farm policy (concurrency limits, scheduling mode), compute location affinity, worker capability tags (OS, GPU, installed software), and usage pool availability. - *Design:* ready tasks remain `ready` until a worker sends a core-NATS - request to `work.lease.`. The server computes free cores + request to `work.lease..`. The server computes free cores (`CPUCount − Σ committed`), selects a priority-ordered batch that fits, atomically transitions the batch `ready → assigned` (stamping `assigned_at` only now), and replies. The `SQI_WORK` JetStream stream, `work.assign.` @@ -212,7 +212,7 @@ NATS JetStream handles: - Heartbeats and worker registration Work leases use **core NATS** request/reply (not JetStream): the worker requests -work on `work.lease.` and the server replies with a batch it is +work on `work.lease..` and the server replies with a batch it is authorized to run (pull-based). Real-time UI updates reach web clients over WebSocket, fanned out by the server after it ingests the JetStream messages. diff --git a/cliff.toml b/cliff.toml index bab5a404..fa12f693 100644 --- a/cliff.toml +++ b/cliff.toml @@ -38,7 +38,9 @@ body = """ {% for commit in commits %} - {% if commit.scope %}**{{ commit.scope }}:** {% endif %}\ {{ commit.message | upper_first }}\ -{% if commit.breaking %} ⚠️ BREAKING{% endif %} \ +{% if commit.breaking %} ⚠️ BREAKING\ +{% if commit.breaking_description and commit.breaking_description != commit.message %} — {{ commit.breaking_description }}{% endif %}\ +{% endif %} \ ([{{ commit.id | truncate(length=7, end="") }}](https://github.com/uberware/sqi/commit/{{ commit.id }})) {%- endfor %} {% endfor %}\n diff --git a/cmd/sqi-server/backup.go b/cmd/sqi-server/backup.go index 9c430062..c2b4ba7e 100644 --- a/cmd/sqi-server/backup.go +++ b/cmd/sqi-server/backup.go @@ -31,6 +31,12 @@ The backup is produced using SQLite's VACUUM INTO statement, which snapshots the live database without taking an exclusive lock. The server may be running or stopped — either works. The destination file must not already exist. +The source database path defaults to store.sqlite_path from the resolved +configuration (the root -c/--config file and SQI_STORE_SQLITE_PATH), falling +back to the legacy SQI_SQLITE_PATH environment variable and then to "sqi.db". +Pass --db to override it explicitly. The source database must already exist — +this command never creates one. + Example: sqi-server backup --db sqi.db --out sqi-backup-$(date +%Y%m%d).db`, RunE: runBackup, @@ -39,8 +45,8 @@ Example: func init() { backupCmd.Flags().StringVar( &backupFlags.DBPath, - "db", envOr("SQI_SQLITE_PATH", "sqi.db"), - "path to source SQLite database file", + "db", "sqi.db", + "path to source SQLite database file (defaults to store.sqlite_path from config, then the deprecated SQI_SQLITE_PATH, then \"sqi.db\")", ) backupCmd.Flags().StringVar( &backupFlags.OutPath, @@ -52,20 +58,31 @@ func init() { } } -func runBackup(_ *cobra.Command, _ []string) error { - if backupFlags.DBPath == "" { - return errors.New("source database path is empty; use --db or set SQI_SQLITE_PATH") - } +func runBackup(cmd *cobra.Command, _ []string) error { if backupFlags.OutPath == "" { return errors.New("destination path is empty; use --out") } + dbPath, err := resolveDBPath(backupFlags.DBPath, cmd != nil && cmd.Flags().Changed("db")) + if err != nil { + return err + } + if dbPath == "" { + return errors.New("source database path is empty; use --db, set store.sqlite_path, or set SQI_STORE_SQLITE_PATH") + } + if err := requireExistingDB(dbPath); err != nil { + return err + } + if err := requireMigratedDB(dbPath); err != nil { + return err + } + logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelInfo})) ctx := context.Background() - logger.InfoContext(ctx, "backup: opening source database", slog.String("path", backupFlags.DBPath)) - st, err := sqlite.Open(ctx, backupFlags.DBPath, sqlite.Options{AutoMigrate: false}) + logger.InfoContext(ctx, "backup: opening source database", slog.String("path", dbPath)) + st, err := sqlite.Open(ctx, dbPath, sqlite.Options{AutoMigrate: false}) if err != nil { return fmt.Errorf("open source database: %w", err) } @@ -74,7 +91,7 @@ func runBackup(_ *cobra.Command, _ []string) error { start := time.Now() logger.InfoContext( ctx, "backup: starting", - slog.String("src", backupFlags.DBPath), + slog.String("src", dbPath), slog.String("dst", backupFlags.OutPath), ) diff --git a/cmd/sqi-server/dbpath.go b/cmd/sqi-server/dbpath.go new file mode 100644 index 00000000..835c1fab --- /dev/null +++ b/cmd/sqi-server/dbpath.go @@ -0,0 +1,122 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package main + +import ( + "context" + "database/sql" + "errors" + "fmt" + "os" + + _ "modernc.org/sqlite" // register "sqlite" driver + + "github.com/uberware/sqi/internal/config" +) + +// resolveDBPath applies the database-path precedence shared by backup, +// migrate, and worker, highest priority first: +// +// 1. explicit is used as-is when explicitChanged is true (the operator +// passed --db on the command line). +// 2. The config layer: the root -c/--config file and SQI_STORE_SQLITE_PATH, +// i.e. whatever [config.LoadWithSources] resolves store.sqlite_path to, +// when [config.Sources.StoreSQLitePath] reports the file or env var +// actually set it — NOT when the resolved value merely differs from the +// built-in default, which a config file that restates the default +// value (as config/sqi-server.example.yaml does) would defeat. +// 3. The legacy SQI_SQLITE_PATH environment variable, kept working as an +// alias. A deprecation notice naming SQI_STORE_SQLITE_PATH — the +// variable sqi-server itself reads — is printed to stderr when this is +// the layer that decided the path. +// 4. The built-in default ("sqi.db"). +// +// The config layer is always loaded, even when an explicit --db makes its +// result moot, so a malformed --config file is reported as an error rather +// than silently ignored when the operator asked for a specific one (root +// -c/--config was passed explicitly). Without an explicit -c, a config-load +// failure — an auto-discovered but broken /etc/sqi/sqi-server.yaml, or an +// unrelated malformed SQI_* env var with nothing to do with the database +// path — is NOT a hard failure: backup, migrate, and worker are the tools +// reached for when something is already broken, so a warning goes to +// stderr and resolution falls through to the legacy env var and default as +// if the config layer had decided nothing. +func resolveDBPath(explicit string, explicitChanged bool) (string, error) { + cfg, src, err := config.LoadWithSources(persistentFlags.ConfigFile, config.FlagOverrides{}) + if err != nil { + if persistentFlags.ConfigFile != "" { + return "", fmt.Errorf("load config: %w", err) + } + fmt.Fprintf( + os.Stderr, + "warning: could not load configuration (%v); falling back to SQI_SQLITE_PATH or the built-in default\n", + err, + ) + cfg = config.DefaultConfig() + src = config.Sources{} + } + + if explicitChanged { + return explicit, nil + } + + if src.StoreSQLitePath { + return cfg.Store.SQLitePath, nil + } + + if legacy := envOr("SQI_SQLITE_PATH", ""); legacy != "" { + fmt.Fprintln(os.Stderr, + "warning: SQI_SQLITE_PATH is deprecated; set SQI_STORE_SQLITE_PATH instead, "+ + "which is the variable sqi-server itself reads") + return legacy, nil + } + + return cfg.Store.SQLitePath, nil +} + +// requireExistingDB stats path and returns an actionable error naming it and +// how to point elsewhere when no file exists there. Used by backup and +// worker, which must never create a database — unlike migrate, whose job is +// to create one. +func requireExistingDB(path string) error { + if _, err := os.Stat(path); err != nil { + if errors.Is(err, os.ErrNotExist) { + return fmt.Errorf( + "no database at %s; point --db, store.sqlite_path, or SQI_STORE_SQLITE_PATH at the right file, "+ + "or run \"sqi-server migrate up\" to create one there", + path, + ) + } + return fmt.Errorf("stat %s: %w", path, err) + } + return nil +} + +// requireMigratedDB catches the case requireExistingDB's stat check cannot +// see: a file that exists at path but was never migrated (empty, or created +// by something other than "migrate up" or the server's own AutoMigrate). +// sqlite.Open with AutoMigrate: false succeeds against such a file — SQLite +// opens an empty database happily — so the first real query would otherwise +// fail with a raw driver error ("no such table: ..."). Checked directly +// against sqlite_master rather than through the store, so it runs before +// any store method that would surface that error unremediated. +func requireMigratedDB(path string) error { + db, err := sql.Open("sqlite", path) + if err != nil { + return fmt.Errorf("open %s: %w", path, err) + } + defer db.Close() + + var count int + err = db.QueryRowContext( + context.Background(), + `SELECT count(*) FROM sqlite_master WHERE type = 'table' AND name NOT LIKE 'sqlite_%'`, + ).Scan(&count) + if err != nil { + return fmt.Errorf("check schema at %s: %w", path, err) + } + if count == 0 { + return fmt.Errorf("the database at %s has no schema; run \"sqi-server migrate up\" first", path) + } + return nil +} diff --git a/cmd/sqi-server/dbpath_test.go b/cmd/sqi-server/dbpath_test.go new file mode 100644 index 00000000..42e6b6b6 --- /dev/null +++ b/cmd/sqi-server/dbpath_test.go @@ -0,0 +1,286 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package main + +import ( + "os" + "path/filepath" + "strings" + "testing" + + "github.com/pressly/goose/v3" +) + +// withConfigFile points persistentFlags.ConfigFile at path for the duration +// of the test and restores the previous value on cleanup. resolveDBPath +// reads persistentFlags.ConfigFile directly (the same global the "-c" root +// flag populates), so tests exercise it without going through cobra parsing. +func withConfigFile(t *testing.T, path string) { + t.Helper() + orig := persistentFlags.ConfigFile + persistentFlags.ConfigFile = path + t.Cleanup(func() { persistentFlags.ConfigFile = orig }) +} + +// writeStoreConfigFile writes a minimal config file setting store.sqlite_path +// and returns its path. +func writeStoreConfigFile(t *testing.T, sqlitePath string) string { + t.Helper() + path := filepath.Join(t.TempDir(), "sqi-server.yaml") + content := "store:\n sqlite_path: " + sqlitePath + "\n" + if err := os.WriteFile(path, []byte(content), 0o600); err != nil { + t.Fatalf("write config file: %v", err) + } + return path +} + +// writeEmptyServerConfigFile writes an empty config file and returns its +// path, for tests that mean "no config file decided anything". Passing this +// instead of "" for persistentFlags.ConfigFile keeps config.Load's default +// search from falling through to $HOME/.sqi/sqi-server.yaml or +// /etc/sqi/sqi-server.yaml — real paths that could exist on the machine +// running the test, which would silently make the test depend on that +// machine's state. +func writeEmptyServerConfigFile(t *testing.T) string { + t.Helper() + path := filepath.Join(t.TempDir(), "empty-sqi-server.yaml") + if err := os.WriteFile(path, []byte("{}\n"), 0o600); err != nil { + t.Fatalf("write empty config file: %v", err) + } + return path +} + +// unsetStoreSQLitePathEnv neutralizes both env vars resolveDBPath consults, +// so a subtest whose meaning is "env decided nothing" (or "only the legacy +// var decided something") is not at the mercy of whatever the machine +// running the test happens to have exported. t.Setenv to "" is treated as +// unset by both config.applyEnv (SQI_STORE_SQLITE_PATH) and resolveDBPath's +// own envOr check (SQI_SQLITE_PATH). +func unsetStoreSQLitePathEnv(t *testing.T) { + t.Helper() + t.Setenv("SQI_STORE_SQLITE_PATH", "") + t.Setenv("SQI_SQLITE_PATH", "") +} + +// TestResolveDBPath_Precedence walks the four-layer precedence order: +// explicit flag > config layer (file/SQI_STORE_SQLITE_PATH) > legacy +// SQI_SQLITE_PATH > built-in default. Every subtest explicitly neutralizes +// both env vars first and only sets the ones its own scenario needs, so +// none of them silently depends on whatever the test machine happens to +// have exported. +func TestResolveDBPath_Precedence(t *testing.T) { + t.Run("explicit flag beats everything else", func(t *testing.T) { + unsetStoreSQLitePathEnv(t) + withConfigFile(t, writeStoreConfigFile(t, "from-config.db")) + t.Setenv("SQI_STORE_SQLITE_PATH", "from-store-env.db") + t.Setenv("SQI_SQLITE_PATH", "from-legacy-env.db") + + got, err := resolveDBPath("from-flag.db", true) + if err != nil { + t.Fatalf("resolveDBPath: %v", err) + } + if got != "from-flag.db" { + t.Errorf("resolveDBPath = %q; want %q", got, "from-flag.db") + } + }) + + t.Run("config file beats legacy env when flag not passed", func(t *testing.T) { + unsetStoreSQLitePathEnv(t) + withConfigFile(t, writeStoreConfigFile(t, "from-config.db")) + t.Setenv("SQI_SQLITE_PATH", "from-legacy-env.db") + + got, err := resolveDBPath("sqi.db", false) + if err != nil { + t.Fatalf("resolveDBPath: %v", err) + } + if got != "from-config.db" { + t.Errorf("resolveDBPath = %q; want %q", got, "from-config.db") + } + }) + + t.Run("SQI_STORE_SQLITE_PATH env beats legacy env when flag not passed", func(t *testing.T) { + unsetStoreSQLitePathEnv(t) + withConfigFile(t, writeEmptyServerConfigFile(t)) + t.Setenv("SQI_STORE_SQLITE_PATH", "from-store-env.db") + t.Setenv("SQI_SQLITE_PATH", "from-legacy-env.db") + + got, err := resolveDBPath("sqi.db", false) + if err != nil { + t.Fatalf("resolveDBPath: %v", err) + } + if got != "from-store-env.db" { + t.Errorf("resolveDBPath = %q; want %q", got, "from-store-env.db") + } + }) + + t.Run("legacy env used and warned about when nothing else set", func(t *testing.T) { + unsetStoreSQLitePathEnv(t) + withConfigFile(t, writeEmptyServerConfigFile(t)) + t.Setenv("SQI_SQLITE_PATH", "from-legacy-env.db") + + var got string + var err error + stderr := captureStderr(t, func() { + got, err = resolveDBPath("sqi.db", false) + }) + if err != nil { + t.Fatalf("resolveDBPath: %v", err) + } + if got != "from-legacy-env.db" { + t.Errorf("resolveDBPath = %q; want %q", got, "from-legacy-env.db") + } + if !strings.Contains(stderr, "SQI_STORE_SQLITE_PATH") { + t.Errorf("expected a deprecation notice naming SQI_STORE_SQLITE_PATH on stderr; got:\n%s", stderr) + } + if !strings.Contains(stderr, "deprecated") { + t.Errorf("expected the word 'deprecated' on stderr; got:\n%s", stderr) + } + }) + + t.Run("built-in default when nothing set", func(t *testing.T) { + unsetStoreSQLitePathEnv(t) + withConfigFile(t, writeEmptyServerConfigFile(t)) + + var got string + var err error + stderr := captureStderr(t, func() { + got, err = resolveDBPath("sqi.db", false) + }) + if err != nil { + t.Fatalf("resolveDBPath: %v", err) + } + if got != "sqi.db" { + t.Errorf("resolveDBPath = %q; want %q", got, "sqi.db") + } + if strings.Contains(stderr, "deprecated") { + t.Errorf("no deprecation notice expected when nothing legacy is set; got:\n%s", stderr) + } + }) +} + +// TestResolveDBPath_MalformedConfigIsAnError verifies that a malformed +// --config file is surfaced as an error rather than silently falling back — +// even when an explicit --db is also passed, since the operator asked for a +// specific config file (persistentFlags.ConfigFile is non-empty here, i.e. +// -c was explicitly passed). +func TestResolveDBPath_MalformedConfigIsAnError(t *testing.T) { + unsetStoreSQLitePathEnv(t) + badPath := filepath.Join(t.TempDir(), "bad.yaml") + if err := os.WriteFile(badPath, []byte("not: [valid: yaml"), 0o600); err != nil { + t.Fatalf("write bad config: %v", err) + } + withConfigFile(t, badPath) + + if _, err := resolveDBPath("explicit.db", true); err == nil { + t.Fatal("expected an error for a malformed config file even with an explicit --db, got nil") + } + if _, err := resolveDBPath("sqi.db", false); err == nil { + t.Fatal("expected an error for a malformed config file, got nil") + } +} + +// TestResolveDBPath_ConfigLoadErrorWithoutExplicitConfigIsLenient verifies +// that a config-load failure having nothing to do with the database path — +// an unrelated malformed SQI_* variable, standing in for what an +// auto-discovered but broken /etc/sqi/sqi-server.yaml would also trigger — +// does not block backup/migrate/worker when -c was not explicitly passed. +// These are the tools reached for when something is already broken; hard +// failure is reserved for when the operator explicitly named a config file +// (see TestResolveDBPath_MalformedConfigIsAnError). +func TestResolveDBPath_ConfigLoadErrorWithoutExplicitConfigIsLenient(t *testing.T) { + unsetStoreSQLitePathEnv(t) + withConfigFile(t, "") // no -c passed + t.Setenv("SQI_SCHEDULER_TICK_INTERVAL", "not-a-duration") + + var got string + var err error + stderr := captureStderr(t, func() { + got, err = resolveDBPath("sqi.db", false) + }) + if err != nil { + t.Fatalf("resolveDBPath: unexpected error: %v", err) + } + if got != "sqi.db" { + t.Errorf("resolveDBPath = %q; want the built-in default %q", got, "sqi.db") + } + if !strings.Contains(stderr, "could not load configuration") { + t.Errorf("expected a warning about the failed config load on stderr; got:\n%s", stderr) + } +} + +// TestResolveDBPath_ConfigLoadErrorWithExplicitConfigIsHardFailure verifies +// the other half of the same rule: a config-load failure IS a hard failure +// once -c was explicitly passed, malformed-env-var or not. +func TestResolveDBPath_ConfigLoadErrorWithExplicitConfigIsHardFailure(t *testing.T) { + unsetStoreSQLitePathEnv(t) + withConfigFile(t, filepath.Join(t.TempDir(), "does-not-exist.yaml")) + t.Setenv("SQI_SCHEDULER_TICK_INTERVAL", "not-a-duration") + + if _, err := resolveDBPath("sqi.db", false); err == nil { + t.Fatal("expected an error: -c named a config file that does not exist") + } +} + +// TestRequireExistingDB verifies the existence guard shared by backup and +// worker: a missing file is an actionable error, an existing one is fine. +func TestRequireExistingDB(t *testing.T) { + t.Run("missing file is an actionable error", func(t *testing.T) { + path := filepath.Join(t.TempDir(), "does-not-exist.db") + err := requireExistingDB(path) + if err == nil { + t.Fatal("expected an error for a missing database, got nil") + } + if !strings.Contains(err.Error(), path) { + t.Errorf("error should name the resolved path %q; got: %v", path, err) + } + if !strings.Contains(err.Error(), "migrate up") { + t.Errorf("error should point at \"migrate up\" as the remediation; got: %v", err) + } + }) + + t.Run("existing file passes", func(t *testing.T) { + path := filepath.Join(t.TempDir(), "exists.db") + if err := os.WriteFile(path, []byte("x"), 0o600); err != nil { + t.Fatalf("write file: %v", err) + } + if err := requireExistingDB(path); err != nil { + t.Errorf("requireExistingDB(%q) = %v; want nil", path, err) + } + }) +} + +// TestRequireMigratedDB verifies the schema guard that catches what +// requireExistingDB's stat check cannot: a file that exists but has no +// tables, because it was created (e.g. by a plain sqlite.Open with +// AutoMigrate: false, or "touch") rather than migrated. +func TestRequireMigratedDB(t *testing.T) { + t.Run("empty file has no schema", func(t *testing.T) { + path := filepath.Join(t.TempDir(), "empty.db") + if err := os.WriteFile(path, nil, 0o600); err != nil { + t.Fatalf("write empty file: %v", err) + } + err := requireMigratedDB(path) + if err == nil { + t.Fatal("expected an error for an unmigrated database, got nil") + } + if !strings.Contains(err.Error(), "migrate up") { + t.Errorf("error should point at \"migrate up\" as the remediation; got: %v", err) + } + }) + + t.Run("migrated database passes", func(t *testing.T) { + path := filepath.Join(t.TempDir(), "migrated.db") + db, err := openMigrateDB(path) + if err != nil { + t.Fatalf("openMigrateDB: %v", err) + } + t.Cleanup(func() { db.Close() }) + if err := goose.Up(db, "."); err != nil { + t.Fatalf("goose.Up: %v", err) + } + + if err := requireMigratedDB(path); err != nil { + t.Errorf("requireMigratedDB(%q) = %v; want nil", path, err) + } + }) +} diff --git a/cmd/sqi-server/main_test.go b/cmd/sqi-server/main_test.go index 11d083d8..d47dc57e 100644 --- a/cmd/sqi-server/main_test.go +++ b/cmd/sqi-server/main_test.go @@ -9,6 +9,9 @@ import ( "path/filepath" "strings" "testing" + + "github.com/spf13/cobra" + "github.com/spf13/pflag" ) // captureStdout redirects os.Stdout to a pipe for the duration of fn, then @@ -39,10 +42,79 @@ func captureStdout(t *testing.T, fn func()) string { return buf.String() } +// captureStderr redirects os.Stderr to a pipe for the duration of fn, then +// returns everything written to it. Needed alongside captureStdout for +// commands (backup's resolveDBPath deprecation notice, keygen's warnings) +// that deliberately separate their stderr diagnostics from stdout output. +// +// Must NOT be called from parallel sub-tests — the redirect is process-wide, +// same caveat as captureStdout. +func captureStderr(t *testing.T, fn func()) string { + t.Helper() + old := os.Stderr + r, w, err := os.Pipe() + if err != nil { + t.Fatalf("os.Pipe: %v", err) + } + os.Stderr = w + + fn() + + w.Close() + os.Stderr = old + + var buf bytes.Buffer + if _, err := io.Copy(&buf, r); err != nil { + t.Fatalf("io.Copy from stderr pipe: %v", err) + } + r.Close() + return buf.String() +} + +// withFlagUnchanged resets a flag's Changed state to false for the duration +// of the test, then restores it. +// +// pflag.Flag.Changed is sticky: it flips to true the first time a flag is +// parsed from the command line and the library never resets it. Flag +// objects in this binary are package-level singletons registered once in +// init(), so within a single test process a later test asserting "the flag +// was omitted" would otherwise observe a stale true left by an earlier test +// that passed it explicitly, purely as an artifact of test ordering. +func withFlagUnchanged(t *testing.T, fs *pflag.FlagSet, name string) { + t.Helper() + f := fs.Lookup(name) + if f == nil { + t.Fatalf("no such flag: %q", name) + } + f.Changed = false + // Reset to false, not the saved original: restoring a stale true would + // leave exactly the leftover this helper exists to prevent, for + // whichever test runs next. + t.Cleanup(func() { f.Changed = false }) +} + +// resetFlagsChanged clears pflag.Flag.Changed for every flag in cmd's own +// FlagSet and PersistentFlags, and recurses into every subcommand. Called +// from prepareRoot before every Execute() so a stale Changed=true left by an +// earlier test's cobra parse — Changed is sticky, pflag never resets it, and +// every command in this binary is a package-level singleton reused across +// the whole test process — cannot leak into a later test that means "this +// flag was not passed". This is the general form of withFlagUnchanged: every +// flag in the tree, before every Execute(), so no individual test needs to +// know which flags an earlier one touched. +func resetFlagsChanged(cmd *cobra.Command) { + cmd.Flags().VisitAll(func(f *pflag.Flag) { f.Changed = false }) + cmd.PersistentFlags().VisitAll(func(f *pflag.Flag) { f.Changed = false }) + for _, c := range cmd.Commands() { + resetFlagsChanged(c) + } +} + // prepareRoot sets the args that rootCmd will parse on the next Execute() call // and redirects cobra's own output writers (help, usage, error messages) to a // discard buffer so test output stays clean. func prepareRoot(args []string) { + resetFlagsChanged(rootCmd) rootCmd.SetArgs(args) var sink bytes.Buffer rootCmd.SetOut(&sink) @@ -308,50 +380,118 @@ func TestMigrateCmd_WithTempDB(t *testing.T) { }) } -// TestRunBackup_ErrorPaths tests the early-return validation guards in -// runBackup directly (without invoking cobra) to avoid the required --out -// flag check that cobra enforces when routing through the command tree. -func TestRunBackup_ErrorPaths(t *testing.T) { - // Save and restore the package-global backupFlags so other tests are - // not affected. - origDB := backupFlags.DBPath - origOut := backupFlags.OutPath - t.Cleanup(func() { - backupFlags.DBPath = origDB - backupFlags.OutPath = origOut +// TestMigrateCmd_DBPath_ExplicitFlagBeatsConfig verifies that --db wins even +// when a config file names a different database, proving migrate consults +// cmd.Flags().Changed("db") rather than always preferring the config layer. +func TestMigrateCmd_DBPath_ExplicitFlagBeatsConfig(t *testing.T) { + explicitPath := filepath.Join(t.TempDir(), "explicit-migrate.db") + configuredPath := filepath.Join(t.TempDir(), "from-config-migrate.db") + withConfigFile(t, writeStoreConfigFile(t, configuredPath)) + + prepareRoot([]string{"migrate", "up", "--db", explicitPath}) + _ = captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("migrate up: unexpected error: %v", err) + } }) - tests := []struct { - name string - dbPath string - outPath string - errContains string - }{ - { - name: "empty db path returns descriptive error", - dbPath: "", - outPath: "somewhere.db", - errContains: "empty", - }, - { - name: "empty out path returns descriptive error", - dbPath: "sqi.db", - outPath: "", - errContains: "empty", - }, + if _, err := os.Stat(explicitPath); err != nil { + t.Errorf("expected migrate to create the explicit --db path %q: %v", explicitPath, err) } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - backupFlags.DBPath = tt.dbPath - backupFlags.OutPath = tt.outPath + if _, err := os.Stat(configuredPath); err == nil { + t.Errorf("migrate created a database at the configured path %q despite an explicit --db", configuredPath) + } +} - err := runBackup(nil, nil) - if err == nil { - t.Fatal("expected error, got nil") - } - if !strings.Contains(err.Error(), tt.errContains) { - t.Errorf("error should contain %q; got: %v", tt.errContains, err) +// TestMigrateCmd_DBPath_ConfigFileHonoredWhenFlagOmitted verifies that +// omitting --db resolves the database through the config layer, and that +// migrate — unlike backup and worker — creates it there. +func TestMigrateCmd_DBPath_ConfigFileHonoredWhenFlagOmitted(t *testing.T) { + withFlagUnchanged(t, migrateCmd.PersistentFlags(), "db") + unsetStoreSQLitePathEnv(t) + configuredPath := filepath.Join(t.TempDir(), "from-config-migrate.db") + withConfigFile(t, writeStoreConfigFile(t, configuredPath)) + + prepareRoot([]string{"migrate", "up"}) + _ = captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("migrate up: unexpected error: %v", err) + } + }) + + if _, err := os.Stat(configuredPath); err != nil { + t.Errorf("expected migrate to create the configured store.sqlite_path %q: %v", configuredPath, err) + } +} + +// TestBackupCmd_ErrorPaths exercises backup's validation guards through the +// real command tree, so cobra's own --out requirement and resolveDBPath's +// flag-changed detection (which needs a live *cobra.Command, not a nil one) +// are both in play. +func TestBackupCmd_ErrorPaths(t *testing.T) { + t.Run("missing --out is rejected by cobra", func(t *testing.T) { + prepareRoot([]string{"backup", "--db", "sqi.db"}) + _ = captureStdout(t, func() { + if err := Execute(); err == nil { + t.Fatal("expected an error for a missing required --out flag, got nil") } }) + }) + + t.Run("explicit empty --db returns a descriptive error", func(t *testing.T) { + outPath := filepath.Join(t.TempDir(), "out.db") + prepareRoot([]string{"backup", "--db", "", "--out", outPath}) + var runErr error + _ = captureStdout(t, func() { + runErr = Execute() + }) + if runErr == nil { + t.Fatal("expected an error for an explicitly empty --db, got nil") + } + if !strings.Contains(runErr.Error(), "empty") { + t.Errorf("error should mention 'empty'; got: %v", runErr) + } + }) + + t.Run("missing source database is an error, not a fresh empty backup", func(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "does-not-exist.db") + outPath := filepath.Join(t.TempDir(), "out.db") + prepareRoot([]string{"backup", "--db", dbPath, "--out", outPath}) + var runErr error + _ = captureStdout(t, func() { + runErr = Execute() + }) + if runErr == nil { + t.Fatal("expected an error for a missing source database, got nil") + } + if !strings.Contains(runErr.Error(), dbPath) { + t.Errorf("error should name the resolved path %q; got: %v", dbPath, runErr) + } + if _, statErr := os.Stat(outPath); statErr == nil { + t.Error("backup should not have produced an output file when the source database is missing") + } + }) +} + +// TestBackupCmd_DBPath_ConfigFileHonoredWhenFlagOmitted verifies that +// omitting --db resolves the source database through the config layer +// (store.sqlite_path) rather than the built-in "sqi.db" default. +func TestBackupCmd_DBPath_ConfigFileHonoredWhenFlagOmitted(t *testing.T) { + withFlagUnchanged(t, backupCmd.Flags(), "db") + unsetStoreSQLitePathEnv(t) + configuredPath := filepath.Join(t.TempDir(), "from-config-backup.db") + createTestDB(t, configuredPath) + withConfigFile(t, writeStoreConfigFile(t, configuredPath)) + + outPath := filepath.Join(t.TempDir(), "out.db") + prepareRoot([]string{"backup", "--out", outPath}) + _ = captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("backup: unexpected error: %v", err) + } + }) + + if _, err := os.Stat(outPath); err != nil { + t.Errorf("expected a backup file at %q: %v", outPath, err) } } diff --git a/cmd/sqi-server/migrate.go b/cmd/sqi-server/migrate.go index c9657dfb..471cc447 100644 --- a/cmd/sqi-server/migrate.go +++ b/cmd/sqi-server/migrate.go @@ -17,7 +17,7 @@ import ( ) // migrateDBPath is the SQLite file path used by all migrate subcommands. -// It is set via --db or falls back to the same default as server.DefaultConfig. +// It is set via --db, or resolved from configuration — see [resolveDBPath]. var migrateDBPath string // migrateCmd groups SQLite schema migration subcommands. @@ -31,10 +31,13 @@ Subcommands: down Roll back the most recently applied migration. status Show applied and pending migrations. -The --db flag (or SQI_SQLITE_PATH env var) controls which SQLite file is -operated on. It defaults to the same path the server uses ("sqi.db" in the -working directory), so running "migrate up" before "serve" is the standard -deployment initialization step.`, +The database path defaults to store.sqlite_path from the resolved +configuration (the root -c/--config file and SQI_STORE_SQLITE_PATH), falling +back to the legacy SQI_SQLITE_PATH environment variable and then to "sqi.db" +in the working directory — the same path the server uses, so running +"migrate up" before "serve" is the standard deployment initialization step. +Pass --db to override it explicitly. Unlike backup and worker, migrate +creates the database file when it does not already exist — that is its job.`, // No RunE — bare "migrate" prints usage. } @@ -42,8 +45,12 @@ var migrateUpCmd = &cobra.Command{ Use: "up", Short: "Apply all pending migrations", Long: `Apply every migration that has not yet been run against the target database.`, - RunE: func(_ *cobra.Command, _ []string) error { - db, err := openMigrateDB(migrateDBPath) + RunE: func(cmd *cobra.Command, _ []string) error { + path, err := resolveDBPath(migrateDBPath, cmd.Flags().Changed("db")) + if err != nil { + return err + } + db, err := openMigrateDB(path) if err != nil { return err } @@ -60,8 +67,12 @@ var migrateDownCmd = &cobra.Command{ Use: "down", Short: "Roll back the last applied migration", Long: `Roll back exactly one migration — the most recently applied one.`, - RunE: func(_ *cobra.Command, _ []string) error { - db, err := openMigrateDB(migrateDBPath) + RunE: func(cmd *cobra.Command, _ []string) error { + path, err := resolveDBPath(migrateDBPath, cmd.Flags().Changed("db")) + if err != nil { + return err + } + db, err := openMigrateDB(path) if err != nil { return err } @@ -78,8 +89,12 @@ var migrateStatusCmd = &cobra.Command{ Use: "status", Short: "Show applied and pending migrations", Long: `List every migration file with its current state: applied (✓) or pending (○).`, - RunE: func(_ *cobra.Command, _ []string) error { - db, err := openMigrateDB(migrateDBPath) + RunE: func(cmd *cobra.Command, _ []string) error { + path, err := resolveDBPath(migrateDBPath, cmd.Flags().Changed("db")) + if err != nil { + return err + } + db, err := openMigrateDB(path) if err != nil { return err } @@ -96,8 +111,8 @@ func init() { // --db flag on the parent so all three subcommands inherit it. migrateCmd.PersistentFlags().StringVar( &migrateDBPath, - "db", envOr("SQI_SQLITE_PATH", "sqi.db"), - "path to SQLite database file", + "db", "sqi.db", + "path to SQLite database file (defaults to store.sqlite_path from config, then the deprecated SQI_SQLITE_PATH, then \"sqi.db\")", ) migrateCmd.AddCommand(migrateUpCmd, migrateDownCmd, migrateStatusCmd) @@ -107,7 +122,7 @@ func init() { // foreign-key enforcement, and wires goose to use the embedded migration FS. func openMigrateDB(path string) (*sql.DB, error) { if path == "" { - return nil, errors.New("database path is empty; use --db or set SQI_SQLITE_PATH") + return nil, errors.New("database path is empty; use --db, set store.sqlite_path, or set SQI_STORE_SQLITE_PATH") } db, err := sql.Open("sqlite", path) diff --git a/cmd/sqi-server/root.go b/cmd/sqi-server/root.go index 8f540059..ce50b338 100644 --- a/cmd/sqi-server/root.go +++ b/cmd/sqi-server/root.go @@ -59,6 +59,7 @@ func init() { migrateCmd, configCmd, backupCmd, + workerCmd, ) } diff --git a/cmd/sqi-server/serve.go b/cmd/sqi-server/serve.go index 4be6db7b..b60fe49f 100644 --- a/cmd/sqi-server/serve.go +++ b/cmd/sqi-server/serve.go @@ -166,30 +166,34 @@ func runServe(cmd *cobra.Command, _ []string) error { // see TestServerConfig_CarriesTheExprCostBounds. func serverConfig(cfg config.Config, schedCfg scheduler.Config) server.Config { return server.Config{ - HTTPAddr: cfg.HTTP.Addr, - CORSOrigins: cfg.HTTP.CORSOrigins, - NATSAddr: cfg.NATS.Addr, - NATSDataDir: cfg.NATS.DataDir, - NATSMaxStoreMB: cfg.NATS.MaxStoreMB, - SQLitePath: cfg.Store.SQLitePath, - EnablePprof: cfg.HTTP.EnablePprof, - CheckpointInterval: cfg.Store.CheckpointInterval, - DiscoveryEnabled: cfg.Discovery.Enabled, - DiscoveryInstanceName: cfg.Discovery.InstanceName, - EnforceOpenJDLimits: cfg.OpenJD.EnforceLimits, - OpenJDExprLimits: server.ExprLimitsFromConfig(cfg.OpenJD), - OpenJDExprSubmissionDeadline: cfg.OpenJD.ExprSubmissionDeadline, - PresetLibraryURL: cfg.PresetLibrary.URL, - AuthEnabled: cfg.Auth.Enabled, - AuthValidateJobOwner: cfg.Auth.ValidateJobOwner, - AuthSessionTTL: cfg.Auth.Session.TTL, - AuthCookieName: cfg.Auth.Session.CookieName, - AuthCookieSecure: cfg.Auth.Session.CookieSecure, - AuthBootstrapUsername: cfg.Auth.Bootstrap.Username, - AuthBootstrapPassword: cfg.Auth.Bootstrap.Password, - AuthLDAP: cfg.Auth.LDAP, - AuthOIDC: cfg.Auth.OIDC, - Scheduler: schedCfg, + HTTPAddr: cfg.HTTP.Addr, + CORSOrigins: cfg.HTTP.CORSOrigins, + NATSAddr: cfg.NATS.Addr, + NATSAuthEnabled: cfg.NATS.Auth.Enabled, + NATSAuthEnrollmentEndpointEnabled: cfg.NATS.Auth.EnrollmentEndpointEnabled, + NATSAuthJoinTokenTTL: cfg.NATS.Auth.JoinTokenTTL, + NATSAuthJoinTokenSingleUse: cfg.NATS.Auth.JoinTokenSingleUse, + NATSDataDir: cfg.NATS.DataDir, + NATSMaxStoreMB: cfg.NATS.MaxStoreMB, + SQLitePath: cfg.Store.SQLitePath, + EnablePprof: cfg.HTTP.EnablePprof, + CheckpointInterval: cfg.Store.CheckpointInterval, + DiscoveryEnabled: cfg.Discovery.Enabled, + DiscoveryInstanceName: cfg.Discovery.InstanceName, + EnforceOpenJDLimits: cfg.OpenJD.EnforceLimits, + OpenJDExprLimits: server.ExprLimitsFromConfig(cfg.OpenJD), + OpenJDExprSubmissionDeadline: cfg.OpenJD.ExprSubmissionDeadline, + PresetLibraryURL: cfg.PresetLibrary.URL, + AuthEnabled: cfg.Auth.Enabled, + AuthValidateJobOwner: cfg.Auth.ValidateJobOwner, + AuthSessionTTL: cfg.Auth.Session.TTL, + AuthCookieName: cfg.Auth.Session.CookieName, + AuthCookieSecure: cfg.Auth.Session.CookieSecure, + AuthBootstrapUsername: cfg.Auth.Bootstrap.Username, + AuthBootstrapPassword: cfg.Auth.Bootstrap.Password, + AuthLDAP: cfg.Auth.LDAP, + AuthOIDC: cfg.Auth.OIDC, + Scheduler: schedCfg, // Phase 1: always seed. Replace with cfg.Store.SeedDefaults when // internal/config grows a setting for it. SeedDefaults: true, diff --git a/cmd/sqi-server/serve_test.go b/cmd/sqi-server/serve_test.go index 8cfc3558..566a02da 100644 --- a/cmd/sqi-server/serve_test.go +++ b/cmd/sqi-server/serve_test.go @@ -112,3 +112,62 @@ func TestServerConfig_CarriesTheRestOfTheConfig(t *testing.T) { t.Error("SeedDefaults = false; this binary always seeds") } } + +// TestServerConfig_CarriesTheBrokerAuthSettings covers the FIRST of the three +// hops (config.Config -> server.Config -> api.Deps -> a mounted route) that +// carry the nats.auth.* settings to the REST worker-enrollment surface. The +// other two hops are pinned in internal/server (natsAuthDeps and the +// router-mount tests) — this one guards the CLI's own mapping, the same shape +// as TestServerConfig_CarriesTheExprCostBounds above. +// +// It matters for the same reason: nothing else fails when this hop breaks. +// Config validation still runs, "config print" still echoes the operator's +// values, and the server still boots — with POST /api/v1/workers/enroll never +// mounted and, if it were, join tokens minted with a zero TTL (born expired). +func TestServerConfig_CarriesTheBrokerAuthSettings(t *testing.T) { + cfg := config.DefaultConfig() + cfg.NATS.Auth.Enabled = true + cfg.NATS.Auth.EnrollmentEndpointEnabled = true + cfg.NATS.Auth.JoinTokenTTL = 42 * time.Minute + cfg.NATS.Auth.JoinTokenSingleUse = false + + got := serverConfig(cfg, scheduler.Config{}) + + if !got.NATSAuthEnabled { + t.Error("NATSAuthEnabled = false, want the configured true") + } + if !got.NATSAuthEnrollmentEndpointEnabled { + t.Error("NATSAuthEnrollmentEndpointEnabled = false, want the configured true") + } + if want := 42 * time.Minute; got.NATSAuthJoinTokenTTL != want { + t.Errorf("NATSAuthJoinTokenTTL = %s, want %s", got.NATSAuthJoinTokenTTL, want) + } + if got.NATSAuthJoinTokenSingleUse { + t.Error("NATSAuthJoinTokenSingleUse = true, want the configured false") + } +} + +// TestServerConfig_BrokerAuthDefaultsAreTheConfigDefaults mirrors +// TestServerConfig_DefaultsAreTheConfigDefaults: a server started with no +// configuration at all must run with internal/config's own defaults (broker +// auth off, but the join-token TTL and single-use settings still sane if it +// is ever turned on), not a zero value that would mean "born expired". +func TestServerConfig_BrokerAuthDefaultsAreTheConfigDefaults(t *testing.T) { + def := config.DefaultConfig() + got := serverConfig(def, scheduler.Config{}) + + if got.NATSAuthEnabled { + t.Error("at defaults NATSAuthEnabled = true, want false") + } + if got.NATSAuthEnrollmentEndpointEnabled != def.NATS.Auth.EnrollmentEndpointEnabled { + t.Errorf("at defaults NATSAuthEnrollmentEndpointEnabled = %v, want %v", + got.NATSAuthEnrollmentEndpointEnabled, def.NATS.Auth.EnrollmentEndpointEnabled) + } + if got.NATSAuthJoinTokenTTL != def.NATS.Auth.JoinTokenTTL { + t.Errorf("at defaults NATSAuthJoinTokenTTL = %s, want %s", got.NATSAuthJoinTokenTTL, def.NATS.Auth.JoinTokenTTL) + } + if got.NATSAuthJoinTokenSingleUse != def.NATS.Auth.JoinTokenSingleUse { + t.Errorf("at defaults NATSAuthJoinTokenSingleUse = %v, want %v", + got.NATSAuthJoinTokenSingleUse, def.NATS.Auth.JoinTokenSingleUse) + } +} diff --git a/cmd/sqi-server/worker.go b/cmd/sqi-server/worker.go new file mode 100644 index 00000000..09ef661b --- /dev/null +++ b/cmd/sqi-server/worker.go @@ -0,0 +1,340 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package main + +import ( + "context" + "errors" + "fmt" + "os" + "text/tabwriter" + "time" + + "github.com/google/uuid" + "github.com/spf13/cobra" + + "github.com/uberware/sqi/internal/auth/jointoken" + "github.com/uberware/sqi/internal/brokerauth" + "github.com/uberware/sqi/internal/config" + "github.com/uberware/sqi/internal/store" + "github.com/uberware/sqi/internal/store/sqlite" +) + +// workerFlags holds the values bound to the "worker" parent command's +// persistent flags, inherited by every subcommand. +var workerFlags struct { + DBPath string +} + +// workerCmd groups worker broker-credential and enrollment subcommands. +// +// These commands open the SQLite database directly, exactly like backup and +// migrate, and never start an HTTP server or NATS broker. That is +// deliberate: broker authentication (nats.auth.enabled) is independent of +// the user-facing auth.enabled gate, and the REST enrollment endpoints only +// exist when auth.enabled is on (otherwise there would be no RBAC in front +// of them). This CLI is what makes credential minting available regardless +// — an operator with a shell on the server host is already trusted. +var workerCmd = &cobra.Command{ + Use: "worker", + Short: "Manage worker broker credentials and enrollment", + Long: `Manage worker broker credentials and enrollment for NATS broker +authentication (the nats.auth.* settings). + +These subcommands operate directly on the SQLite database file and do not +require sqi-server to be running, and work regardless of whether +auth.enabled (the separate, user-facing auth gate) is on. + +The database path defaults to store.sqlite_path from the resolved +configuration (the root -c/--config file and SQI_STORE_SQLITE_PATH), falling +back to the legacy SQI_SQLITE_PATH environment variable and then to "sqi.db". +Pass --db to override it explicitly. The database must already exist — these +subcommands never create one; run "sqi-server migrate up" first. + +Subcommands: + token issue Issue a one-time join token for self-service enrollment. + enroll Directly register a worker's credential by ID and public key + (offline — see "enroll --help"). + revoke Revoke a worker's credential (offline — see "revoke --help"). + list List active worker credentials.`, +} + +// workerTokenCmd groups worker join-token subcommands. +var workerTokenCmd = &cobra.Command{ + Use: "token", + Short: "Manage worker join tokens", +} + +var workerTokenIssueFlags struct { + TTL time.Duration + Name string +} + +var workerTokenIssueCmd = &cobra.Command{ + Use: "issue", + Short: "Issue a one-time worker join token", + Long: `Issue a new join token that lets a worker enroll itself and obtain a +broker credential (POST /api/v1/workers/enroll). + +The raw token is printed to stdout exactly once, immediately after creation. +Only its SHA-256 hash is stored in the database — the raw value cannot be +recovered or displayed again. If it is lost before the worker uses it, issue +a new one; the old one still works until it expires or is used.`, + RunE: runWorkerTokenIssue, +} + +var workerEnrollFlags struct { + WorkerID string + PublicKey string + Name string +} + +var workerEnrollCmd = &cobra.Command{ + Use: "enroll", + Short: "Directly register a worker's broker credential", + Long: `Register a worker's broker credential by worker ID and public key +directly, without a join token. + +This is the manual path for a worker that cannot reach sqi-server's REST API +to self-enroll (an air-gapped host, or an operator who provisions credentials +by hand): run "sqi-worker keygen" on the worker host to generate a keypair, +then run this command on the server with the worker ID and public key it +prints. + +A RUNNING sqi-server does not see the new credential — it reads the enrolled +set once, at startup, and this command writes the database from a separate +process with no broker handle. The worker's connection is refused until +sqi-server is restarted. To enroll against a running server instead, use the +REST API with a join token: + POST /api/v1/workers/enroll`, + RunE: runWorkerEnroll, +} + +var workerRevokeCmd = &cobra.Command{ + Use: "revoke WORKER_ID", + Short: "Revoke a worker's broker credential", + Long: `Revokes a worker credential in the database. A RUNNING sqi-server does not +see this immediately — it applies at next start. To revoke a worker on a +running server and disconnect it at once, use the REST API: + DELETE /api/v1/workers/{id}/credential`, + Args: cobra.ExactArgs(1), + RunE: runWorkerRevoke, +} + +var workerListCmd = &cobra.Command{ + Use: "list", + Short: "List active worker broker credentials", + Long: `List every worker credential that has not been revoked.`, + RunE: runWorkerList, +} + +func init() { + workerCmd.PersistentFlags().StringVar( + &workerFlags.DBPath, + "db", "sqi.db", + "path to SQLite database file (defaults to store.sqlite_path from config, then the deprecated SQI_SQLITE_PATH, then \"sqi.db\")", + ) + + workerTokenIssueCmd.Flags().DurationVar( + &workerTokenIssueFlags.TTL, + "ttl", config.DefaultConfig().NATS.Auth.JoinTokenTTL, + fmt.Sprintf("how long the token remains valid (between %s and %s)", + config.MinNATSAuthJoinTokenTTL, config.MaxNATSAuthJoinTokenTTL), + ) + workerTokenIssueCmd.Flags().StringVar( + &workerTokenIssueFlags.Name, + "name", "", + "optional human-readable label for this token", + ) + workerTokenCmd.AddCommand(workerTokenIssueCmd) + + workerEnrollCmd.Flags().StringVar(&workerEnrollFlags.WorkerID, "worker-id", "", "the worker's stable ID (required)") + workerEnrollCmd.Flags().StringVar(&workerEnrollFlags.PublicKey, "public-key", "", "the worker's nkey public key, starting with 'U' (required)") + workerEnrollCmd.Flags().StringVar(&workerEnrollFlags.Name, "name", "", "optional human-readable label for this credential") + for _, f := range []string{"worker-id", "public-key"} { + if err := workerEnrollCmd.MarkFlagRequired(f); err != nil { + panic(err) + } + } + + workerCmd.AddCommand(workerTokenCmd, workerEnrollCmd, workerRevokeCmd, workerListCmd) +} + +// openWorkerStore resolves the database path (see [resolveDBPath]) and opens +// it without applying migrations. These commands write rows into an existing +// schema; they never create or migrate the database — a worker subcommand +// pointed at the wrong file must fail with an actionable error, not conjure +// an empty one. Run "sqi-server migrate up" first against a fresh database. +func openWorkerStore(ctx context.Context, cmd *cobra.Command) (*sqlite.Store, error) { + dbPath, err := resolveDBPath(workerFlags.DBPath, cmd != nil && cmd.Flags().Changed("db")) + if err != nil { + return nil, err + } + if dbPath == "" { + return nil, errors.New("database path is empty; use --db, set store.sqlite_path, or set SQI_STORE_SQLITE_PATH") + } + if err := requireExistingDB(dbPath); err != nil { + return nil, err + } + if err := requireMigratedDB(dbPath); err != nil { + return nil, err + } + + st, err := sqlite.Open(ctx, dbPath, sqlite.Options{AutoMigrate: false}) + if err != nil { + return nil, fmt.Errorf("open database: %w", err) + } + return st, nil +} + +func runWorkerTokenIssue(cmd *cobra.Command, _ []string) error { + ttl := workerTokenIssueFlags.TTL + if ttl < config.MinNATSAuthJoinTokenTTL || ttl > config.MaxNATSAuthJoinTokenTTL { + return fmt.Errorf("--ttl must be between %s and %s, got %s", + config.MinNATSAuthJoinTokenTTL, config.MaxNATSAuthJoinTokenTTL, ttl) + } + + ctx := context.Background() + st, err := openWorkerStore(ctx, cmd) + if err != nil { + return err + } + defer st.Close() + + token, hash, prefix, err := jointoken.Generate() + if err != nil { + return fmt.Errorf("generate join token: %w", err) + } + + now := time.Now().UTC() + _, err = st.CreateWorkerJoinToken(ctx, store.WorkerJoinToken{ + ID: uuid.NewString(), + TokenHash: hash, + Prefix: prefix, + Name: workerTokenIssueFlags.Name, + ExpiresAt: now.Add(ttl), + CreatedBy: "cli", + CreatedAt: now, + }) + if err != nil { + return fmt.Errorf("store join token: %w", err) + } + + // The warning goes to stderr and the token alone to stdout, so + // `TOKEN=$(sqi-server worker token issue)` captures exactly the token. + fmt.Fprintln(os.Stderr, "This token will not be shown again — store it securely now.") + fmt.Fprintln(os.Stdout, token) + return nil +} + +func runWorkerEnroll(cmd *cobra.Command, _ []string) error { + if err := brokerauth.ValidatePublicKey(workerEnrollFlags.PublicKey); err != nil { + return err + } + // The recorded worker ID is what this credential's broker grants are + // built from (brokerauth.WorkerPermissions), and those grants are NATS + // subject PATTERNS — so "*" would record a credential allowed to publish + // concrete subjects belonging to any worker on the farm, and ">" would + // put a malformed subject into the broker's key set. MarkFlagRequired + // does not cover this: --worker-id "" counts as supplied. + if !brokerauth.ValidWorkerIDToken(workerEnrollFlags.WorkerID) { + return fmt.Errorf( + "--worker-id %q is not a valid NATS subject token: it must be non-empty and must not contain '.', whitespace, '*' or '>'", + workerEnrollFlags.WorkerID, + ) + } + + ctx := context.Background() + st, err := openWorkerStore(ctx, cmd) + if err != nil { + return err + } + defer st.Close() + + _, err = st.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: uuid.NewString(), + WorkerID: workerEnrollFlags.WorkerID, + PublicKey: workerEnrollFlags.PublicKey, + Name: workerEnrollFlags.Name, + EnrolledAt: time.Now().UTC(), + }) + if err != nil { + if errors.Is(err, store.ErrConflict) { + return fmt.Errorf( + "cannot enroll worker %q: it already has a credential, or this public key is already enrolled to another worker", + workerEnrollFlags.WorkerID, + ) + } + return fmt.Errorf("store credential: %w", err) + } + + // Symmetric with runWorkerRevoke's warning, and for the same reason: this + // command writes the database from a process with no broker handle, and + // the broker's authorized-key set is built once at Broker.Start. Without + // this line an operator who enrolls against a running server sees + // "Enrolled worker" and then a worker that exits fatally on + // nats.ErrAuthorization, with nothing connecting the two. + fmt.Fprintf(os.Stdout, "Enrolled worker %q. A RUNNING sqi-server will not accept this credential until it restarts;"+ + " to enroll against a running server, use POST /api/v1/workers/enroll with a join token instead.\n", + workerEnrollFlags.WorkerID) + return nil +} + +func runWorkerRevoke(cmd *cobra.Command, args []string) error { + workerID := args[0] + + ctx := context.Background() + st, err := openWorkerStore(ctx, cmd) + if err != nil { + return err + } + defer st.Close() + + if err := st.RevokeWorkerCredential(ctx, workerID, time.Now().UTC()); err != nil { + if errors.Is(err, store.ErrNotFound) { + // RevokeWorkerCredential's SQL matches only an active credential + // (revoked_at IS NULL), so this same error covers "never + // enrolled" and "already revoked" — say so rather than claiming + // the worker does not exist, which may be false. + return fmt.Errorf( + "no active credential for worker %q — it may never have been enrolled, or its credential may already be revoked", + workerID, + ) + } + return fmt.Errorf("revoke credential: %w", err) + } + + fmt.Fprintf(os.Stdout, "Revoked credential for worker %q. This takes effect the next time sqi-server starts;"+ + " to disconnect it immediately, use DELETE /api/v1/workers/%s/credential instead.\n", workerID, workerID) + return nil +} + +func runWorkerList(cmd *cobra.Command, _ []string) error { + ctx := context.Background() + st, err := openWorkerStore(ctx, cmd) + if err != nil { + return err + } + defer st.Close() + + creds, err := st.ListActiveWorkerCredentials(ctx) + if err != nil { + return fmt.Errorf("list credentials: %w", err) + } + + w := tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0) + fmt.Fprintln(w, "WORKER ID\tNAME\tPUBLIC KEY\tENROLLED\tLAST SEEN") + for _, c := range creds { + name := c.Name + if name == "" { + name = "-" + } + lastSeen := "never" + if c.LastSeenAt != nil { + lastSeen = c.LastSeenAt.Format(time.RFC3339) + } + fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\n", + c.WorkerID, name, c.PublicKey, c.EnrolledAt.Format(time.RFC3339), lastSeen) + } + return w.Flush() +} diff --git a/cmd/sqi-server/worker_test.go b/cmd/sqi-server/worker_test.go new file mode 100644 index 00000000..dd6e8162 --- /dev/null +++ b/cmd/sqi-server/worker_test.go @@ -0,0 +1,528 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package main + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/uberware/sqi/internal/auth/jointoken" + "github.com/uberware/sqi/internal/brokerauth" + "github.com/uberware/sqi/internal/store/sqlite" +) + +// createTestDB creates and migrates a SQLite database at path. The worker +// subcommands never create a database themselves (see [openWorkerStore]), so +// any test exercising a code path that reaches the store needs one to +// already exist first — exactly as an operator would run "sqi-server migrate +// up" before "sqi-server worker ...". +func createTestDB(t *testing.T, path string) { + t.Helper() + st, err := sqlite.Open(context.Background(), path, sqlite.DefaultOptions()) + if err != nil { + t.Fatalf("create test database at %s: %v", path, err) + } + if err := st.Close(); err != nil { + t.Fatalf("close test database at %s: %v", path, err) + } +} + +// TestWorkerCmd_TokenIssue verifies that "worker token issue" prints the raw +// token to stdout exactly once and that only its hash is ever stored. +func TestWorkerCmd_TokenIssue(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "test.db") + createTestDB(t, dbPath) + + prepareRoot([]string{"worker", "token", "issue", "--db", dbPath, "--name", "ci-runner"}) + out := captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("Execute(worker token issue) error = %v", err) + } + }) + + token := strings.TrimSpace(out) + if token == "" { + t.Fatal("expected a token on stdout, got empty output") + } + if strings.Count(out, token) != 1 { + t.Errorf("expected the token to appear exactly once in stdout output; got:\n%s", out) + } + if !strings.HasPrefix(token, "sqiw_") { + t.Errorf("token %q does not have the expected join-token prefix", token) + } + + // The store must hold a record whose hash matches the printed token, and + // nothing else exposes the raw value: WorkerJoinToken only ever carries + // TokenHash, never the token itself. + st, err := sqlite.Open(context.Background(), dbPath, sqlite.Options{AutoMigrate: false}) + if err != nil { + t.Fatalf("reopen store: %v", err) + } + defer st.Close() + + hash := jointoken.Hash(token) + got, err := st.GetWorkerJoinTokenByHash(context.Background(), hash) + if err != nil { + t.Fatalf("GetWorkerJoinTokenByHash: %v", err) + } + if got.Name != "ci-runner" { + t.Errorf("stored token Name = %q; want %q", got.Name, "ci-runner") + } + if got.TokenHash == token { + t.Error("stored TokenHash equals the raw token; only the hash must be persisted") + } +} + +// TestWorkerCmd_TokenIssue_TTLOutOfBounds verifies that an out-of-range +// --ttl is rejected before anything is written to the store. +func TestWorkerCmd_TokenIssue_TTLOutOfBounds(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "test.db") + + tests := []struct { + name string + ttl string + }{ + {"below floor", "30s"}, + {"above ceiling", "48h"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + prepareRoot([]string{"worker", "token", "issue", "--db", dbPath, "--ttl", tt.ttl}) + _ = captureStdout(t, func() { + if err := Execute(); err == nil { + t.Fatal("expected an error for an out-of-bounds --ttl, got nil") + } + }) + }) + } +} + +// TestWorkerCmd_Enroll_InvalidPublicKey verifies that enrolling with a +// malformed public key fails and names the expected "U" nkey prefix. +func TestWorkerCmd_Enroll_InvalidPublicKey(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "test.db") + + prepareRoot([]string{"worker", "enroll", "--db", dbPath, "--worker-id", "w1", "--public-key", "not-a-valid-key"}) + var runErr error + _ = captureStdout(t, func() { + runErr = Execute() + }) + if runErr == nil { + t.Fatal("expected an error for an invalid public key, got nil") + } + if !strings.Contains(runErr.Error(), "'U'") { + t.Errorf("error should name the expected 'U' nkey prefix; got: %v", runErr) + } +} + +// TestWorkerCmd_Enroll_WarnsThatARunningServerNeedsARestart pins the one +// thing an operator cannot discover from a successful enroll: the command +// opens the SQLite file from a separate process, so it cannot reload a +// running broker's authorized-key set (built once at Broker.Start). Without +// this warning the sequence reads as a success followed by an unrelated +// worker that exits on an authorization error. The revoke command carries +// the mirror-image warning; both must keep saying so. +func TestWorkerCmd_Enroll_WarnsThatARunningServerNeedsARestart(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "test.db") + createTestDB(t, dbPath) + _, pub, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed: %v", err) + } + + prepareRoot([]string{"worker", "enroll", "--db", dbPath, "--worker-id", "w1", "--public-key", pub}) + out := captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("enroll: unexpected error: %v", err) + } + }) + + for _, want := range []string{ + "RUNNING sqi-server", + "restarts", + "/api/v1/workers/enroll", + } { + if !strings.Contains(out, want) { + t.Errorf("enroll output does not mention %q; got: %s", want, out) + } + } +} + +// TestWorkerCmd_Enroll_LongHelpWarnsAboutARunningServer keeps the same +// warning on the command's own help text, so an operator reading +// "worker enroll --help" before running anything learns it too. +func TestWorkerCmd_Enroll_LongHelpWarnsAboutARunningServer(t *testing.T) { + if !strings.Contains(workerEnrollCmd.Long, "RUNNING sqi-server") { + t.Errorf("workerEnrollCmd.Long does not warn that a running server will not see the credential; got:\n%s", workerEnrollCmd.Long) + } + if !strings.Contains(workerEnrollCmd.Long, "POST /api/v1/workers/enroll") { + t.Errorf("workerEnrollCmd.Long does not name the REST alternative; got:\n%s", workerEnrollCmd.Long) + } +} + +// TestWorkerCmd_Enroll_DuplicateWorkerIDFails verifies that enrolling the +// same worker ID twice with two different keys fails the second time. +func TestWorkerCmd_Enroll_DuplicateWorkerIDFails(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "test.db") + createTestDB(t, dbPath) + + _, pub1, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed: %v", err) + } + _, pub2, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed: %v", err) + } + + prepareRoot([]string{"worker", "enroll", "--db", dbPath, "--worker-id", "w1", "--public-key", pub1}) + _ = captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("first enroll: unexpected error: %v", err) + } + }) + + prepareRoot([]string{"worker", "enroll", "--db", dbPath, "--worker-id", "w1", "--public-key", pub2}) + var secondErr error + _ = captureStdout(t, func() { + secondErr = Execute() + }) + if secondErr == nil { + t.Fatal("expected the second enroll for the same worker ID to fail, got nil") + } +} + +// TestWorkerCmd_RotationAfterRevoke walks the whole key-rotation flow: +// enroll -> (keygen --force, stood in for by generating a second local +// keypair, exactly as sqi-worker keygen would produce) -> re-enrolling that +// worker ID must fail while the old credential is still active -> revoke -> +// re-enroll with the new key succeeds -> the new key, not the old one, is +// what's active. It holds only because +// internal/store/migrations/00030_broker_auth.sql scopes worker_id +// uniqueness to active rows, not the whole table; scoped to the whole table, +// revocation is a one-way door and the worker ID can never be used again. +func TestWorkerCmd_RotationAfterRevoke(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "test.db") + createTestDB(t, dbPath) + + _, pub1, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed (first key): %v", err) + } + _, pub2, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed (rotated key): %v", err) + } + + prepareRoot([]string{"worker", "enroll", "--db", dbPath, "--worker-id", "w1", "--public-key", pub1}) + _ = captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("initial enroll: unexpected error: %v", err) + } + }) + + // Rotating the key locally (what "sqi-worker keygen --force" does) does + // not by itself free up the worker ID on the server: the old credential + // is still active, so re-enrolling must still fail here. + prepareRoot([]string{"worker", "enroll", "--db", dbPath, "--worker-id", "w1", "--public-key", pub2}) + var beforeRevokeErr error + _ = captureStdout(t, func() { + beforeRevokeErr = Execute() + }) + if beforeRevokeErr == nil { + t.Fatal("expected re-enrolling an active worker ID with a new key to fail before revoking the old credential") + } + + prepareRoot([]string{"worker", "revoke", "w1", "--db", dbPath}) + _ = captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("revoke: unexpected error: %v", err) + } + }) + + // With the old credential revoked, the same worker ID must be free to + // enroll again with the new key. + prepareRoot([]string{"worker", "enroll", "--db", dbPath, "--worker-id", "w1", "--public-key", pub2}) + _ = captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("re-enroll after revoke: unexpected error: %v", err) + } + }) + + prepareRoot([]string{"worker", "list", "--db", dbPath}) + out := captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("list: unexpected error: %v", err) + } + }) + if !strings.Contains(out, pub2) { + t.Errorf("list output missing the rotated (new) public key; got:\n%s", out) + } + if strings.Contains(out, pub1) { + t.Errorf("list output still shows the revoked (old) public key; got:\n%s", out) + } +} + +// TestWorkerCmd_Revoke_UnknownWorkerFails verifies that revoking a worker +// with no credential exits non-zero with an accurate message — one that +// does not claim the worker itself doesn't exist, since the same error also +// covers "already revoked". +func TestWorkerCmd_Revoke_UnknownWorkerFails(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "test.db") + createTestDB(t, dbPath) + + prepareRoot([]string{"worker", "revoke", "does-not-exist", "--db", dbPath}) + var runErr error + _ = captureStdout(t, func() { + runErr = Execute() + }) + if runErr == nil { + t.Fatal("expected an error revoking an unknown worker, got nil") + } + if strings.Contains(runErr.Error(), "does not exist") { + t.Errorf("error must not claim the worker does not exist (it may also be already-revoked); got: %v", runErr) + } + if !strings.Contains(runErr.Error(), "already be revoked") { + t.Errorf("error should mention the already-revoked possibility; got: %v", runErr) + } +} + +// TestWorkerCmd_Revoke_TwiceFailsTheSecondTime pins the documented behavior +// that RevokeWorkerCredential collapses "unknown worker" and "already +// revoked" into the same not-found outcome. +func TestWorkerCmd_Revoke_TwiceFailsTheSecondTime(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "test.db") + createTestDB(t, dbPath) + + _, pub, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed: %v", err) + } + prepareRoot([]string{"worker", "enroll", "--db", dbPath, "--worker-id", "w1", "--public-key", pub}) + _ = captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("enroll: unexpected error: %v", err) + } + }) + + prepareRoot([]string{"worker", "revoke", "w1", "--db", dbPath}) + _ = captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("first revoke: unexpected error: %v", err) + } + }) + + prepareRoot([]string{"worker", "revoke", "w1", "--db", dbPath}) + var secondErr error + _ = captureStdout(t, func() { + secondErr = Execute() + }) + if secondErr == nil { + t.Fatal("expected the second revoke to fail, got nil") + } +} + +// TestWorkerCmd_List verifies that "worker list" reflects enrollment and +// stops listing a worker once its credential is revoked (list only shows +// active credentials). +func TestWorkerCmd_List(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "test.db") + createTestDB(t, dbPath) + + _, pub, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed: %v", err) + } + prepareRoot([]string{"worker", "enroll", "--db", dbPath, "--worker-id", "w1", "--public-key", pub, "--name", "render-01"}) + _ = captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("enroll: unexpected error: %v", err) + } + }) + + prepareRoot([]string{"worker", "list", "--db", dbPath}) + out := captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("list: unexpected error: %v", err) + } + }) + if !strings.Contains(out, "w1") || !strings.Contains(out, "render-01") || !strings.Contains(out, pub) { + t.Errorf("list output missing enrolled worker fields; got:\n%s", out) + } + + prepareRoot([]string{"worker", "revoke", "w1", "--db", dbPath}) + _ = captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("revoke: unexpected error: %v", err) + } + }) + + prepareRoot([]string{"worker", "list", "--db", dbPath}) + out = captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("list after revoke: unexpected error: %v", err) + } + }) + if strings.Contains(out, "w1") { + t.Errorf("list output should not include a revoked worker; got:\n%s", out) + } +} + +// TestWorkerCmd_OpenWorkerStore_NilCommandDoesNotPanic verifies the +// (cmd *cobra.Command) parameter is safe to omit for direct, non-cobra +// callers: it is treated the same as "the --db flag was not passed", +// falling through to config-layer resolution rather than panicking on a nil +// flag set. +func TestWorkerCmd_OpenWorkerStore_NilCommandDoesNotPanic(t *testing.T) { + origDB := workerFlags.DBPath + t.Cleanup(func() { workerFlags.DBPath = origDB }) + + workerFlags.DBPath = filepath.Join(t.TempDir(), "does-not-exist.db") + _, err := openWorkerStore(context.Background(), nil) + if err == nil { + t.Fatal("expected an error for a database that does not exist, got nil") + } +} + +// TestWorkerCmd_ExplicitEmptyDBPath verifies that an explicitly-passed empty +// --db (as opposed to the flag simply being omitted) is reported as a clear +// validation error rather than silently falling through to config +// resolution. +func TestWorkerCmd_ExplicitEmptyDBPath(t *testing.T) { + prepareRoot([]string{"worker", "list", "--db", ""}) + var runErr error + _ = captureStdout(t, func() { + runErr = Execute() + }) + if runErr == nil { + t.Fatal("expected an error for an explicitly empty --db, got nil") + } + if !strings.Contains(runErr.Error(), "empty") { + t.Errorf("error should mention 'empty'; got: %v", runErr) + } +} + +// TestWorkerCmd_MissingDatabaseIsErrorNotCreation verifies that pointing a +// worker subcommand at a database file that does not exist fails with an +// actionable error and does not create one — unlike migrate, which is +// expected to create a fresh database. +func TestWorkerCmd_MissingDatabaseIsErrorNotCreation(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "does-not-exist.db") + + prepareRoot([]string{"worker", "list", "--db", dbPath}) + var runErr error + _ = captureStdout(t, func() { + runErr = Execute() + }) + if runErr == nil { + t.Fatal("expected an error for a missing database, got nil") + } + if !strings.Contains(runErr.Error(), dbPath) { + t.Errorf("error should name the resolved path %q; got: %v", dbPath, runErr) + } + if !strings.Contains(runErr.Error(), "migrate up") { + t.Errorf("error should point at \"migrate up\" as the remediation; got: %v", runErr) + } + if _, statErr := os.Stat(dbPath); statErr == nil { + t.Error("worker subcommand must not create a database file") + } +} + +// TestWorkerCmd_DBPath_ExplicitFlagBeatsConfig verifies that --db wins even +// when a config file names a different (nonexistent) database — proving the +// flag is actually consulted via cmd.Flags().Changed("db") rather than +// config always winning once loaded. +func TestWorkerCmd_DBPath_ExplicitFlagBeatsConfig(t *testing.T) { + explicitPath := filepath.Join(t.TempDir(), "explicit.db") + createTestDB(t, explicitPath) + + configuredPath := filepath.Join(t.TempDir(), "from-config.db") + // Deliberately never created — if the flag were ignored in favor of the + // config layer, this run would fail with "no database at" the + // configured path instead of succeeding against the explicit one. + withConfigFile(t, writeStoreConfigFile(t, configuredPath)) + + prepareRoot([]string{"worker", "list", "--db", explicitPath}) + _ = captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("worker list: unexpected error: %v", err) + } + }) +} + +// TestWorkerCmd_DBPath_ConfigFileHonoredWhenFlagOmitted verifies that +// omitting --db entirely resolves the database through the config layer +// (store.sqlite_path), not the built-in "sqi.db" default. +func TestWorkerCmd_DBPath_ConfigFileHonoredWhenFlagOmitted(t *testing.T) { + withFlagUnchanged(t, workerCmd.PersistentFlags(), "db") + unsetStoreSQLitePathEnv(t) + configuredPath := filepath.Join(t.TempDir(), "from-config.db") + createTestDB(t, configuredPath) + withConfigFile(t, writeStoreConfigFile(t, configuredPath)) + + prepareRoot([]string{"worker", "list"}) + _ = captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("worker list: unexpected error: %v", err) + } + }) +} + +// TestWorkerCmd_Enroll_InvalidWorkerID rejects a worker ID that is not a +// single NATS subject token, at the offline enrollment boundary. +// +// The recorded worker_id is what brokerauth.WorkerPermissions builds this +// credential's grants from, and those grants are subject PATTERNS: a +// worker ID of "*" yields "task.status.*.*", "worker.deregister.*", +// "work.lease.*.*" and the rest, so one credential could publish concrete +// subjects belonging to any worker on the farm. ">" additionally produces +// the malformed "task.status.>.*" inside Options.Nkeys, which can wedge +// the broker or make every later credential reload fail. +func TestWorkerCmd_Enroll_InvalidWorkerID(t *testing.T) { + cases := []struct { + name string + workerID string + }{ + {"single-token wildcard", "*"}, + {"multi-token wildcard", ">"}, + {"contains a dot", "render.01"}, + {"contains whitespace", "render 01"}, + {"empty", ""}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + dbPath := filepath.Join(t.TempDir(), "test.db") + createTestDB(t, dbPath) + _, pub, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed: %v", err) + } + + prepareRoot([]string{"worker", "enroll", "--db", dbPath, "--worker-id", tc.workerID, "--public-key", pub}) + var runErr error + _ = captureStdout(t, func() { + runErr = Execute() + }) + if runErr == nil { + t.Fatalf("worker id %q: expected an error, got nil", tc.workerID) + } + if !strings.Contains(runErr.Error(), "worker-id") { + t.Errorf("worker id %q: error should name the offending flag; got: %v", tc.workerID, runErr) + } + + // Nothing may have been recorded under it. + prepareRoot([]string{"worker", "list", "--db", dbPath}) + out := captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("list: unexpected error: %v", err) + } + }) + if strings.Contains(out, pub) { + t.Errorf("worker id %q: a credential was recorded anyway:\n%s", tc.workerID, out) + } + }) + } +} diff --git a/cmd/sqi-worker/README.md b/cmd/sqi-worker/README.md index e58d76da..ed90ff8a 100644 --- a/cmd/sqi-worker/README.md +++ b/cmd/sqi-worker/README.md @@ -2,7 +2,7 @@ `sqi-worker` is the distributed task-execution agent for the sqi render farm. It connects to a running `sqi-server`, registers its hardware capabilities, and -requests task assignments over core-NATS work leases (`work.lease.`), +requests task assignments over core-NATS work leases (`work.lease..`), executing them as bare-metal OS processes. --- @@ -15,7 +15,7 @@ executing them as bare-metal OS processes. detectable), plus any manual tags from configuration such as `maya-2025` or `arnold-7`. - **Leases** task assignments over core NATS (a long-polling request/reply on - `work.lease.`) and executes them concurrently; the server gates + `work.lease..`) and executes them concurrently; the server gates concurrency via CPU-core accounting. Task status, logs, heartbeats and registration travel the other way over JetStream. - **Streams** task stdout and stderr back to `sqi-server` in real time so the diff --git a/cmd/sqi-worker/connectbroker_test.go b/cmd/sqi-worker/connectbroker_test.go new file mode 100644 index 00000000..9ebcdb08 --- /dev/null +++ b/cmd/sqi-worker/connectbroker_test.go @@ -0,0 +1,110 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package main + +import ( + "context" + "log/slog" + "net" + "strconv" + "strings" + "testing" + "time" + + "github.com/uberware/sqi/internal/bus" + workerconfig "github.com/uberware/sqi/internal/worker/config" +) + +// freeTestPort asks the OS for an unused loopback TCP port, for booting a +// throwaway embedded broker. +func freeTestPort(t *testing.T) int { + t.Helper() + var lc net.ListenConfig + l, err := lc.Listen(context.Background(), "tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("freeTestPort: listen: %v", err) + } + defer func() { _ = l.Close() }() + addr, ok := l.Addr().(*net.TCPAddr) + if !ok { + t.Fatalf("freeTestPort: listener address is %T, want *net.TCPAddr", l.Addr()) + } + return addr.Port +} + +// startEmbeddedBroker boots a real embedded NATS broker on a throwaway +// loopback port and JetStream dir, with or without auth, and registers +// cleanup. A real broker is used rather than a fake so these tests exercise +// the actual nkey handshake, matching how [connectToBroker] behaves in +// production. +func startEmbeddedBroker(t *testing.T, auth bus.BrokerAuthConfig) *bus.Broker { + t.Helper() + b := bus.New(bus.BrokerConfig{ + Addr: net.JoinHostPort("127.0.0.1", strconv.Itoa(freeTestPort(t))), + DataDir: t.TempDir() + "/nats", + Auth: auth, + }, slog.New(slog.DiscardHandler)) + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + if err := b.Start(ctx); err != nil { + t.Fatalf("startEmbeddedBroker: Start: %v", err) + } + t.Cleanup(b.Shutdown) + return b +} + +// TestConnectToBroker_AuthOffFarmBootsWithNoCredential exercises the exact +// boot-time path runStart takes on the default, auth-off configuration: a +// worker with no credential file and no join token configured must still +// connect, exactly as it did before broker authentication existed. This is +// the regression this feature is not allowed to introduce. +func TestConnectToBroker_AuthOffFarmBootsWithNoCredential(t *testing.T) { + b := startEmbeddedBroker(t, bus.BrokerAuthConfig{Enabled: false}) + + var cfg workerconfig.WorkerConfig + cfg.NATS.URL = b.ClientURL() + cfg.NATS.CredentialFile = t.TempDir() + "/worker.nk" // deliberately does not exist + cfg.NATS.MaxReconnectAttempts = 0 + cfg.NATS.ReconnectWait = 10 * time.Millisecond + + logger := slog.New(slog.DiscardHandler) + nc, _, err := connectToBroker(context.Background(), cfg, "worker-a", logger) + if err != nil { + t.Fatalf("connectToBroker on an auth-off farm with no credential: %v", err) + } + defer nc.Close() + if !nc.IsConnected() { + t.Error("connection is not in the connected state") + } +} + +// TestConnectToBroker_AuthOnFarmWithNoCredentialNamesBothRemediations covers +// the case the boot-sequence exit message in cmd/sqi-worker/start.go exists +// for: a worker with neither a credential file nor a join token, against a +// broker that actually requires authentication. The failure must be fatal +// and must name both ways an operator can fix it. +func TestConnectToBroker_AuthOnFarmWithNoCredentialNamesBothRemediations(t *testing.T) { + b := startEmbeddedBroker(t, bus.BrokerAuthConfig{Enabled: true}) + + var cfg workerconfig.WorkerConfig + cfg.NATS.URL = b.ClientURL() + cfg.NATS.CredentialFile = t.TempDir() + "/worker.nk" + cfg.NATS.MaxReconnectAttempts = 0 + cfg.NATS.ReconnectWait = 10 * time.Millisecond + + logger := slog.New(slog.DiscardHandler) + _, _, err := connectToBroker(context.Background(), cfg, "worker-a", logger) + if err == nil { + t.Fatal("connectToBroker: want error against an auth-on broker with no credential, got nil") + } + if !strings.Contains(err.Error(), "no credential was found") { + t.Errorf("error %q does not say no credential was found", err.Error()) + } + if !strings.Contains(err.Error(), "sqi-worker keygen") { + t.Errorf("error %q does not mention pre-provisioning a key", err.Error()) + } + if !strings.Contains(err.Error(), "sqi-server worker token issue") { + t.Errorf("error %q does not mention obtaining a join token", err.Error()) + } +} diff --git a/cmd/sqi-worker/keygen.go b/cmd/sqi-worker/keygen.go new file mode 100644 index 00000000..45b30cb5 --- /dev/null +++ b/cmd/sqi-worker/keygen.go @@ -0,0 +1,175 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package main + +import ( + "errors" + "fmt" + "os" + + "github.com/spf13/cobra" + + "github.com/uberware/sqi/internal/brokerauth" + workerconfig "github.com/uberware/sqi/internal/worker/config" +) + +var keygenFlags struct { + DataDir string + Force bool +} + +var keygenCmd = &cobra.Command{ + Use: "keygen", + Short: "Generate a new nkey broker credential for this worker", + Long: `Generate a new Ed25519 nkey keypair for authenticating this worker to +sqi-server's NATS broker, and write the private seed to the path resolved +from nats.credential_file (which defaults to /worker.nk), mode +0600. + +Configuration is loaded the same way "sqi-worker start" loads it: the root +-c/--config file, SQI_WORKER_* environment variables, and built-in defaults. +Run this on the worker host with its normal config so the worker ID and +credential path match what that worker actually uses. --data-dir overrides +worker.data_dir explicitly, for a one-off run against a different directory. + +Refuses to overwrite an existing seed unless --force is given — generating a +new keypair for a worker that is already enrolled invalidates its current +credential; the broker will reject connections signed with the old key until +the worker is re-enrolled with the new public key. + +Prints the public key, whether the worker ID is an existing one loaded from +the data directory or one newly generated there, and the exact +"sqi-server worker enroll" command an operator must run on the server to +authorize this worker. This is the manual enrollment path — a worker that +can reach sqi-server's REST API can instead enroll itself automatically with +a join token issued by "sqi-server worker token issue".`, + RunE: runKeygen, +} + +func init() { + keygenCmd.Flags().StringVar( + &keygenFlags.DataDir, + "data-dir", "", + "override worker.data_dir (also moves the default credential path unless nats.credential_file is set explicitly)", + ) + keygenCmd.Flags().BoolVar( + &keygenFlags.Force, + "force", false, + "overwrite an existing seed file", + ) +} + +func runKeygen(cmd *cobra.Command, _ []string) error { + cfg, src, err := workerconfig.LoadWithSources(persistentFlags.ConfigFile, flagOverrides()) + if err != nil { + return fmt.Errorf("load config: %w", err) + } + + dataDir, seedPath := resolveKeygenPaths(cmd, cfg, src) + if dataDir == "" { + return errors.New("worker data directory is empty; use --data-dir or set worker.data_dir") + } + if seedPath == "" { + return errors.New("credential file path is empty; use --data-dir or set nats.credential_file") + } + + _, statErr := os.Stat(seedPath) + switch { + case statErr == nil: + if !keygenFlags.Force { + return fmt.Errorf( + "a credential already exists at %s; use --force to overwrite it (this invalidates the existing enrollment)", + seedPath, + ) + } + case errors.Is(statErr, os.ErrNotExist): + // No existing seed — nothing to warn about below. + default: + return fmt.Errorf("stat %s: %w", seedPath, statErr) + } + seedExisted := statErr == nil + + seed, publicKey, err := brokerauth.GenerateSeed() + if err != nil { + return fmt.Errorf("generate keypair: %w", err) + } + + if err := brokerauth.SaveSeed(seedPath, seed); err != nil { + return err + } + + idPath := workerconfig.WorkerIDFilePath(dataDir) + _, idStatErr := os.Stat(idPath) + workerIDExisted := idStatErr == nil + if idStatErr != nil && !errors.Is(idStatErr, os.ErrNotExist) { + return fmt.Errorf("stat %s: %w", idPath, idStatErr) + } + + workerID, err := workerconfig.LoadOrCreateWorkerID(dataDir) + if err != nil { + return fmt.Errorf("load or create worker id: %w", err) + } + + fmt.Fprintf(os.Stdout, "Public key: %s\n", publicKey) + if workerIDExisted { + fmt.Fprintf(os.Stdout, "Worker ID: %s (existing, loaded from %s)\n", workerID, idPath) + } else { + fmt.Fprintf( + os.Stdout, + "Worker ID: %s (newly generated; if you expected an existing worker id here, "+ + "--data-dir/worker.data_dir is probably pointed at the wrong directory)\n", + workerID, + ) + } + fmt.Fprintln(os.Stdout, "On the server, run:") + fmt.Fprintf(os.Stdout, " sqi-server worker enroll --worker-id %s --public-key %s\n", workerID, publicKey) + fmt.Fprintln(os.Stdout, "A RUNNING sqi-server will not accept this credential until it restarts;"+ + " to enroll against a running server, use POST /api/v1/workers/enroll with a join token instead.") + + if seedExisted { + // --force just overwrote a seed that may still be the credential the + // server has enrolled. The new local key means nothing to the broker + // until the server side is updated too — and today the server still + // has the OLD public key on file, so it must be revoked before the + // enroll command above can succeed (worker_id stays unique among + // active credentials). + fmt.Fprintf( + os.Stderr, + "Warning: this replaced an existing seed. The previous credential for worker %s is likely still"+ + " enrolled on the server; revoke it first or the enroll command above will fail:\n"+ + " sqi-server worker revoke %s\n", + workerID, workerID, + ) + } + return nil +} + +// resolveKeygenPaths returns the worker data directory and credential seed +// path to use, applying --data-dir as an explicit override of the loaded +// cfg.Worker.DataDir. +// +// When --data-dir is passed and src.CredentialFile is false — meaning +// nats.credential_file was NOT explicitly set by the config file or +// SQI_WORKER_NATS_CREDENTIAL_FILE, so cfg.NATS.CredentialFile is only +// [workerconfig.Load]'s own config-derived fill-in — the seed path is +// re-derived under the overridden directory too, so it still lands at +// /worker.nk instead of the pre-override location. This is a +// provenance check, not a value comparison: comparing seedPath against +// DefaultCredentialFile(dataDir) would wrongly treat an explicitly +// configured nats.credential_file that happens to equal that computed +// default as "unset" and relocate it anyway. An explicitly configured +// nats.credential_file is always left untouched. +func resolveKeygenPaths(cmd *cobra.Command, cfg workerconfig.WorkerConfig, src workerconfig.Sources) (dataDir, seedPath string) { + dataDir = cfg.Worker.DataDir + seedPath = cfg.NATS.CredentialFile + + if cmd == nil || !cmd.Flags().Changed("data-dir") { + return dataDir, seedPath + } + + if !src.CredentialFile { + seedPath = workerconfig.DefaultCredentialFile(keygenFlags.DataDir) + } + dataDir = keygenFlags.DataDir + return dataDir, seedPath +} diff --git a/cmd/sqi-worker/keygen_test.go b/cmd/sqi-worker/keygen_test.go new file mode 100644 index 00000000..3df1a67e --- /dev/null +++ b/cmd/sqi-worker/keygen_test.go @@ -0,0 +1,338 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package main + +import ( + "bytes" + "io" + "os" + "path/filepath" + "runtime" + "strings" + "testing" + + workerconfig "github.com/uberware/sqi/internal/worker/config" +) + +// writeWorkerConfigFile writes a minimal sqi-worker config file setting +// worker.data_dir and returns its path. +func writeWorkerConfigFile(t *testing.T, dataDir string) string { + t.Helper() + path := filepath.Join(t.TempDir(), "sqi-worker.yaml") + content := "worker:\n data_dir: " + dataDir + "\n" + if err := os.WriteFile(path, []byte(content), 0o600); err != nil { + t.Fatalf("write config file: %v", err) + } + return path +} + +// hermeticKeygenEnv makes a keygen test hermetic against the developer +// machine it happens to run on, and returns a config file path every test +// must pass via --config. +// +// runKeygen calls workerconfig.LoadWithSources, which applies SQI_WORKER_* +// environment variables on top of the config file and, when --config is +// empty, SEARCHES $HOME/.sqi/sqi-worker.yaml and /etc/sqi/sqi-worker.yaml. +// A developer running a local worker is exactly who is likely to have +// SQI_WORKER_NATS_CREDENTIAL_FILE exported or a real ~/.sqi/sqi-worker.yaml +// in place — and if either names a real credential path, keygen writes +// genuine Ed25519 key material there instead of under the test's +// t.TempDir(), independent of what --data-dir says (an explicitly +// configured nats.credential_file is deliberately left untouched by +// --data-dir; see resolveKeygenPaths). t.Setenv to "" is treated as unset by +// workerconfig's applyEnv (every branch is "if v := os.Getenv(key); v != +// \"\" { ... }"), so it neutralizes each variable for the duration of the +// test without needing to know or restore its ambient value. +func hermeticKeygenEnv(t *testing.T) (configPath string) { + t.Helper() + for _, v := range []string{ + "SQI_WORKER_DATA_DIR", + "SQI_WORKER_NATS_CREDENTIAL_FILE", + "SQI_WORKER_NATS_URL", + "SQI_WORKER_NATS_JOIN_TOKEN", + "SQI_WORKER_NATS_JOIN_TOKEN_FILE", + } { + t.Setenv(v, "") + } + // Every keygen test in this file passes --config explicitly (this + // function's return value, or its own), so persistentFlags.ConfigFile + // always points at a real path while the test runs — but that path + // lives under this test's t.TempDir() and is gone once it returns. + // Reset the package-level var so a later test that means to search the + // default paths (empty ConfigFile) does not inherit a now-deleted path. + t.Cleanup(func() { persistentFlags.ConfigFile = "" }) + + path := filepath.Join(t.TempDir(), "empty-sqi-worker.yaml") + if err := os.WriteFile(path, []byte("{}\n"), 0o600); err != nil { + t.Fatalf("write empty config file: %v", err) + } + return path +} + +// captureStderr redirects os.Stderr to a pipe for the duration of fn, then +// returns everything written to it. keygen's overwrite warning is +// deliberately written to stderr (see runKeygen), so asserting on it needs +// this alongside captureStdout — both streams are written within the same +// call and must be captured together, not sequentially. +// +// Must NOT be called from parallel sub-tests — the redirect is process-wide, +// same caveat as captureStdout in main_test.go. +func captureStderr(t *testing.T, fn func()) string { + t.Helper() + old := os.Stderr + r, w, err := os.Pipe() + if err != nil { + t.Fatalf("os.Pipe: %v", err) + } + os.Stderr = w + + fn() + + w.Close() + os.Stderr = old + + var buf bytes.Buffer + if _, err := io.Copy(&buf, r); err != nil { + t.Fatalf("io.Copy from stderr pipe: %v", err) + } + r.Close() + return buf.String() +} + +// TestKeygenCmd_WritesSeedAndPrintsEnrollCommand verifies that "keygen" +// writes a 0600 seed file, prints the public key, and prints the exact +// "sqi-server worker enroll" command to run — but never prints the seed. +func TestKeygenCmd_WritesSeedAndPrintsEnrollCommand(t *testing.T) { + cfgPath := hermeticKeygenEnv(t) + dataDir := filepath.Join(t.TempDir(), "worker-data") + + prepareRoot([]string{"keygen", "--config", cfgPath, "--data-dir", dataDir}) + out := captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("Execute(keygen) error = %v", err) + } + }) + + if !strings.Contains(out, "Public key: U") { + t.Errorf("output missing public key line; got:\n%s", out) + } + if !strings.Contains(out, "sqi-server worker enroll --worker-id") { + t.Errorf("output missing the enroll command; got:\n%s", out) + } + if !strings.Contains(out, "--public-key U") { + t.Errorf("output missing --public-key flag; got:\n%s", out) + } + if !strings.Contains(out, "will not accept this credential until it restarts") { + t.Errorf("output missing the running-server restart note; got:\n%s", out) + } + if !strings.Contains(out, "POST /api/v1/workers/enroll") { + t.Errorf("output missing the join-token REST enrollment alternative; got:\n%s", out) + } + + seedPath := filepath.Join(dataDir, "worker.nk") + info, err := os.Stat(seedPath) + if err != nil { + t.Fatalf("stat seed file: %v", err) + } + + seedBytes, err := os.ReadFile(seedPath) + if err != nil { + t.Fatalf("read seed file: %v", err) + } + if strings.Contains(out, string(seedBytes)) { + t.Error("stdout output must never contain the seed bytes") + } + + if runtime.GOOS != "windows" { + if perm := info.Mode().Perm(); perm != 0o600 { + t.Errorf("seed file mode = %o; want 0600", perm) + } + } +} + +// TestKeygenCmd_RefusesToOverwriteWithoutForce verifies that a second +// keygen invocation against the same data directory fails without --force, +// and succeeds (overwriting the seed) with it. +func TestKeygenCmd_RefusesToOverwriteWithoutForce(t *testing.T) { + cfgPath := hermeticKeygenEnv(t) + dataDir := filepath.Join(t.TempDir(), "worker-data") + + prepareRoot([]string{"keygen", "--config", cfgPath, "--data-dir", dataDir}) + _ = captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("first keygen: unexpected error: %v", err) + } + }) + + seedPath := filepath.Join(dataDir, "worker.nk") + before, err := os.ReadFile(seedPath) + if err != nil { + t.Fatalf("read seed after first keygen: %v", err) + } + + prepareRoot([]string{"keygen", "--config", cfgPath, "--data-dir", dataDir}) + var secondErr error + _ = captureStdout(t, func() { + secondErr = Execute() + }) + if secondErr == nil { + t.Fatal("expected keygen without --force to fail when a seed already exists") + } + + after, err := os.ReadFile(seedPath) + if err != nil { + t.Fatalf("read seed after refused overwrite: %v", err) + } + if string(before) != string(after) { + t.Error("seed file changed despite the overwrite being refused") + } + + prepareRoot([]string{"keygen", "--config", cfgPath, "--data-dir", dataDir, "--force"}) + var out string + errOut := captureStderr(t, func() { + out = captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("forced keygen: unexpected error: %v", err) + } + }) + }) + if !strings.Contains(out, "Public key: U") { + t.Errorf("forced keygen output missing public key line; got:\n%s", out) + } + + // The overwrite warning belongs where an operator who already passed + // --force will actually see it (stderr, on success) — not only in + // --help text or in the refusal message shown when --force is absent. + workerID, err := workerconfig.LoadOrCreateWorkerID(dataDir) + if err != nil { + t.Fatalf("LoadOrCreateWorkerID: %v", err) + } + if !strings.Contains(errOut, "Warning: this replaced an existing seed") { + t.Errorf("forced keygen stderr missing the overwrite warning; got:\n%s", errOut) + } + wantRevokeCmd := "sqi-server worker revoke " + workerID + if !strings.Contains(errOut, wantRevokeCmd) { + t.Errorf("forced keygen stderr missing the exact revoke command %q; got:\n%s", wantRevokeCmd, errOut) + } + + forced, err := os.ReadFile(seedPath) + if err != nil { + t.Fatalf("read seed after forced overwrite: %v", err) + } + if string(before) == string(forced) { + t.Error("seed file did not change after --force") + } + if runtime.GOOS != "windows" { + info, err := os.Stat(seedPath) + if err != nil { + t.Fatalf("stat seed after forced overwrite: %v", err) + } + if perm := info.Mode().Perm(); perm != 0o600 { + t.Errorf("seed file mode after --force = %o; want 0600", perm) + } + } +} + +// TestKeygenCmd_EmptyDataDir verifies that an explicitly-passed empty +// --data-dir (as opposed to the flag simply being omitted, which falls +// through to worker.data_dir) is rejected with a descriptive error. Routed +// through Execute() rather than calling runKeygen directly: with no cobra +// command, --data-dir cannot be told apart from "omitted", which would fall +// through to the real, config-resolved worker.data_dir instead of the empty +// value this test means to exercise. +func TestKeygenCmd_EmptyDataDir(t *testing.T) { + cfgPath := hermeticKeygenEnv(t) + prepareRoot([]string{"keygen", "--config", cfgPath, "--data-dir", ""}) + var runErr error + _ = captureStdout(t, func() { + runErr = Execute() + }) + if runErr == nil { + t.Fatal("expected an error for an empty --data-dir, got nil") + } + if !strings.Contains(runErr.Error(), "empty") { + t.Errorf("error should mention 'empty'; got: %v", runErr) + } +} + +// TestKeygenCmd_DataDir_ExplicitFlagBeatsConfig verifies that --data-dir +// wins even when a config file names a different worker.data_dir. +func TestKeygenCmd_DataDir_ExplicitFlagBeatsConfig(t *testing.T) { + hermeticKeygenEnv(t) + explicitDir := filepath.Join(t.TempDir(), "explicit-data") + configuredDir := filepath.Join(t.TempDir(), "configured-data") + cfgPath := writeWorkerConfigFile(t, configuredDir) + t.Cleanup(func() { persistentFlags.ConfigFile = "" }) + + prepareRoot([]string{"keygen", "--config", cfgPath, "--data-dir", explicitDir}) + _ = captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("keygen: unexpected error: %v", err) + } + }) + + if _, err := os.Stat(filepath.Join(explicitDir, "worker.nk")); err != nil { + t.Errorf("expected a seed under the explicit --data-dir %s: %v", explicitDir, err) + } + if _, err := os.Stat(configuredDir); err == nil { + t.Errorf("keygen wrote under the configured worker.data_dir %s despite an explicit --data-dir", configuredDir) + } +} + +// TestKeygenCmd_DataDir_ConfigFileHonoredWhenFlagOmitted verifies that +// omitting --data-dir resolves worker.data_dir through the config layer +// rather than the platform default under the real home directory. +func TestKeygenCmd_DataDir_ConfigFileHonoredWhenFlagOmitted(t *testing.T) { + withFlagUnchanged(t, keygenCmd.Flags(), "data-dir") + hermeticKeygenEnv(t) + + configuredDir := filepath.Join(t.TempDir(), "configured-data") + cfgPath := writeWorkerConfigFile(t, configuredDir) + t.Cleanup(func() { persistentFlags.ConfigFile = "" }) + + prepareRoot([]string{"keygen", "--config", cfgPath}) + _ = captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("keygen: unexpected error: %v", err) + } + }) + + if _, err := os.Stat(filepath.Join(configuredDir, "worker.nk")); err != nil { + t.Errorf("expected a seed under the configured worker.data_dir %s: %v", configuredDir, err) + } +} + +// TestKeygenCmd_ReportsNewVsExistingWorkerID verifies that keygen states +// plainly whether the worker ID it printed was loaded from an existing +// worker.id file or freshly generated — the one signal an operator rotating +// a key has that --data-dir/worker.data_dir points at the wrong directory. +func TestKeygenCmd_ReportsNewVsExistingWorkerID(t *testing.T) { + cfgPath := hermeticKeygenEnv(t) + dataDir := filepath.Join(t.TempDir(), "worker-data") + + prepareRoot([]string{"keygen", "--config", cfgPath, "--data-dir", dataDir}) + firstOut := captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("first keygen: unexpected error: %v", err) + } + }) + if !strings.Contains(firstOut, "newly generated") { + t.Errorf("first run against an empty data dir should report a newly generated worker id; got:\n%s", firstOut) + } + if strings.Contains(firstOut, "existing, loaded from") { + t.Errorf("first run must not claim an existing worker id; got:\n%s", firstOut) + } + + prepareRoot([]string{"keygen", "--config", cfgPath, "--data-dir", dataDir, "--force"}) + secondOut := captureStdout(t, func() { + if err := Execute(); err != nil { + t.Fatalf("second keygen: unexpected error: %v", err) + } + }) + if !strings.Contains(secondOut, "existing, loaded from") { + t.Errorf("second run against the same data dir should report the existing worker id; got:\n%s", secondOut) + } + if strings.Contains(secondOut, "newly generated") { + t.Errorf("second run must not claim a newly generated worker id; got:\n%s", secondOut) + } +} diff --git a/cmd/sqi-worker/lease_queueids_test.go b/cmd/sqi-worker/lease_queueids_test.go index 656f83f4..8a17d540 100644 --- a/cmd/sqi-worker/lease_queueids_test.go +++ b/cmd/sqi-worker/lease_queueids_test.go @@ -17,8 +17,11 @@ func TestLeaseQueueIDs(t *testing.T) { if len(got) != 1 || got[0] != bus.WildcardQueueToken { t.Fatalf("leaseQueueIDs(nil) = %v, want [%q]", got, bus.WildcardQueueToken) } - // The resulting subject must be valid (non-empty leaf). - if subj := bus.WorkLeaseSubject(got[0]); subj == bus.SubjectWorkLeasePrefix+"." { + // The resulting subject must be valid: a parseable worker → server + // lease subject with a non-empty queue token. + subj := bus.WorkLeaseSubject("w-1", got[0]) + workerID, queueID, ok := bus.ParseWorkerSubject(subj) + if !ok || workerID != "w-1" || queueID != bus.WildcardQueueToken { t.Fatalf("wildcard produced invalid subject %q", subj) } }) diff --git a/cmd/sqi-worker/main.go b/cmd/sqi-worker/main.go index b1400e26..7b65d94c 100644 --- a/cmd/sqi-worker/main.go +++ b/cmd/sqi-worker/main.go @@ -4,7 +4,7 @@ // // It discovers and connects to a running sqi-server, registers itself with // its capability tags and compute location, leases task assignments over core -// NATS (work.lease.), and executes bare-metal OS processes inside +// NATS (work.lease..), and executes bare-metal OS processes inside // OpenJD sessions. // // Run "sqi-worker --help" for usage. diff --git a/cmd/sqi-worker/main_test.go b/cmd/sqi-worker/main_test.go index 8f478642..6daf8ff0 100644 --- a/cmd/sqi-worker/main_test.go +++ b/cmd/sqi-worker/main_test.go @@ -9,6 +9,9 @@ import ( "path/filepath" "strings" "testing" + + "github.com/spf13/cobra" + "github.com/spf13/pflag" ) // captureStdout redirects os.Stdout to a pipe for the duration of fn, then @@ -39,10 +42,50 @@ func captureStdout(t *testing.T, fn func()) string { return buf.String() } +// withFlagUnchanged resets a flag's Changed state to false for the duration +// of the test, then restores it. +// +// pflag.Flag.Changed is sticky: it flips to true the first time a flag is +// parsed from the command line and the library never resets it. Flag +// objects in this binary are package-level singletons registered once in +// init(), so within a single test process a later test asserting "the flag +// was omitted" would otherwise observe a stale true left by an earlier test +// that passed it explicitly, purely as an artifact of test ordering. +func withFlagUnchanged(t *testing.T, fs *pflag.FlagSet, name string) { + t.Helper() + f := fs.Lookup(name) + if f == nil { + t.Fatalf("no such flag: %q", name) + } + f.Changed = false + // Reset to false, not the saved original: restoring a stale true would + // leave exactly the leftover this helper exists to prevent, for + // whichever test runs next. + t.Cleanup(func() { f.Changed = false }) +} + +// resetFlagsChanged clears pflag.Flag.Changed for every flag in cmd's own +// FlagSet and PersistentFlags, and recurses into every subcommand. Called +// from prepareRoot before every Execute() so a stale Changed=true left by an +// earlier test's cobra parse — Changed is sticky, pflag never resets it, and +// every command in this binary is a package-level singleton reused across +// the whole test process — cannot leak into a later test that means "this +// flag was not passed". This is the general form of withFlagUnchanged: every +// flag in the tree, before every Execute(), so no individual test needs to +// know which flags an earlier one touched. +func resetFlagsChanged(cmd *cobra.Command) { + cmd.Flags().VisitAll(func(f *pflag.Flag) { f.Changed = false }) + cmd.PersistentFlags().VisitAll(func(f *pflag.Flag) { f.Changed = false }) + for _, c := range cmd.Commands() { + resetFlagsChanged(c) + } +} + // prepareRoot sets the args that rootCmd will parse on the next Execute() call // and redirects cobra's own output writers (help, usage, error messages) to a // discard buffer so test output stays clean. func prepareRoot(args []string) { + resetFlagsChanged(rootCmd) rootCmd.SetArgs(args) var sink bytes.Buffer rootCmd.SetOut(&sink) diff --git a/cmd/sqi-worker/root.go b/cmd/sqi-worker/root.go index 46e5985f..971f59e1 100644 --- a/cmd/sqi-worker/root.go +++ b/cmd/sqi-worker/root.go @@ -17,7 +17,7 @@ var rootCmd = &cobra.Command{ It discovers and connects to a running sqi-server, registers itself with its capability tags and compute location, leases task assignments over core NATS -(work.lease.), and executes bare-metal OS processes inside OpenJD +(work.lease..), and executes bare-metal OS processes inside OpenJD sessions. Use "sqi-worker start" to start the worker agent. @@ -59,6 +59,7 @@ func init() { configCmd, capabilitiesCmd, isolationCmd, + keygenCmd, ) } diff --git a/cmd/sqi-worker/start.go b/cmd/sqi-worker/start.go index 5c6feaa3..c334eed7 100644 --- a/cmd/sqi-worker/start.go +++ b/cmd/sqi-worker/start.go @@ -26,6 +26,7 @@ import ( workerconfig "github.com/uberware/sqi/internal/worker/config" "github.com/uberware/sqi/internal/worker/diaglog" workerdiscovery "github.com/uberware/sqi/internal/worker/discovery" + "github.com/uberware/sqi/internal/worker/enroll" "github.com/uberware/sqi/internal/worker/executor" "github.com/uberware/sqi/internal/worker/heartbeat" "github.com/uberware/sqi/internal/worker/lease" @@ -146,15 +147,14 @@ func runStart(cmd *cobra.Command, _ []string) error { // all downstream log statements use the concrete URL. cfg.NATS.URL = natsURL - // ── NATS connection ───────────────────────────────────────── + // ── Broker credential + NATS connection ────────────────────── // // Connect failure at boot is fatal. closedCh is closed by the // NATS ClosedHandler when the connection permanently closes so we can // detect unexpected disconnects after the initial handshake. - nc, natsClosed, err := natsclient.Connect(ctx, cfg.NATS, logger) + nc, natsClosed, err := connectToBroker(ctx, cfg, workerID, logger) if err != nil { - // Boot-time connect failure is a fatal error. - return fmt.Errorf("nats connect: %w", err) + return err } // ── Diagnostic-log sink ───────────────────────────────────── @@ -295,7 +295,7 @@ func runStart(cmd *cobra.Command, _ []string) error { // SYNC: logstreamer.Config and workerconfig.LogStreamerConfig have matching // fields. If a field is added to one, it must be added to both and mapped // here. - logPub := logstreamer.New(nc, logstreamer.Config{ + logPub := logstreamer.New(nc, workerID, logstreamer.Config{ MaxLinesPerChunk: cfg.LogStreamer.MaxLinesPerChunk, MaxBytesPerChunk: cfg.LogStreamer.MaxBytesPerChunk, FlushInterval: cfg.LogStreamer.FlushInterval, @@ -352,7 +352,7 @@ func runStart(cmd *cobra.Command, _ []string) error { // ── Heartbeat ─────────────────────────────────────────────── // // The heartbeat Publisher ticks on cfg.Worker.HeartbeatInterval and - // publishes liveness + runtime-state messages to worker.heartbeat. + // publishes liveness + runtime-state messages to worker.heartbeat.. // The executor is wired in as the StateSource so each heartbeat carries // the current active-task count, active task IDs, and last-assignment time. // @@ -372,7 +372,7 @@ func runStart(cmd *cobra.Command, _ []string) error { // ── Work-lease loop ───────────────────────────────────────── // - // The worker asks the server for work on work.lease. and dispatches + // The worker asks the server for work on work.lease.. and dispatches // whatever the server leases it. The server gates capacity (CPU-core fit, // policy, usage pools), so the worker simply runs what it is given. leaseLoop := lease.New( @@ -467,6 +467,51 @@ func checkRootAndLoadWorkerID(cfg workerconfig.WorkerConfig, logger *slog.Logger return workerID, nil } +// connectToBroker loads or obtains this worker's nkey broker credential and +// dials the broker with it. Extracted from [runStart] to keep that +// function's cyclomatic complexity within the project limit. +// +// A missing credential (enroll.ErrNoCredential) is NOT immediately fatal: +// this may be a farm that does not require worker authentication at all, in +// which case the connect attempt below proceeds with no credential exactly +// as it does today. Only when the broker actually refuses that connection is +// authentication known to be required, and the operator gets a message +// naming both remediations (obtain a token, or pre-provision a key) instead +// of natsclient's generic rejection message, which assumes a credential +// existed to be rejected in the first place. +func connectToBroker( + ctx context.Context, + cfg workerconfig.WorkerConfig, + workerID string, + logger *slog.Logger, +) (*nats.Conn, <-chan struct{}, error) { + seed, publicKey, err := enroll.EnsureCredential(ctx, enroll.Config{ + WorkerID: workerID, + CredentialFile: cfg.NATS.CredentialFile, + JoinToken: cfg.NATS.JoinToken, + JoinTokenFile: cfg.NATS.JoinTokenFile, + ServerURL: cfg.NATS.ServerURL, + }, logger) + noCredential := errors.Is(err, enroll.ErrNoCredential) + if err != nil && !noCredential { + return nil, nil, fmt.Errorf("worker credential: %w", err) + } + + nc, natsClosed, err := natsclient.Connect(ctx, cfg.NATS, workerID, seed, publicKey, logger) + if err != nil { + if noCredential && (errors.Is(err, nats.ErrAuthorization) || errors.Is(err, nats.ErrAuthExpired)) { + //nolint:staticcheck // ST1005: wording is deliberately two sentences so the operator sees cause and remediation separately + return nil, nil, fmt.Errorf( + "sqi-worker: this server requires worker authentication, but no credential was found at %s and no join token is configured.\n"+ + "Obtain a token from an operator (`sqi-server worker token issue`) and set worker nats.join_token_file, or pre-provision a key with `sqi-worker keygen`.", + cfg.NATS.CredentialFile, + ) + } + return nil, nil, fmt.Errorf("nats connect: %w", err) + } + return nc, natsClosed, nil +} + // loadAndValidateConfig resolves CLI flag overrides, loads the layered // configuration, and runs validation — returning a ready-to-use [WorkerConfig] // or an error with an actionable message. Extracted from [runStart] to keep @@ -609,7 +654,7 @@ func withDiagnosticSink( } // leaseTransport adapts a raw *nats.Conn to [lease.Transport], issuing -// core-NATS request/reply work-lease requests on the work.lease. +// core-NATS request/reply work-lease requests on the work.lease.. // subject. The worker connects via natsclient (which yields a *nats.Conn), so // this thin wrapper provides the same RequestLease behavior as // [bus.Client.RequestLease] without a second client connection. @@ -617,13 +662,13 @@ type leaseTransport struct { nc *nats.Conn } -// RequestLease sends a work-lease request for queueID and waits up to timeout -// for the server's reply. It returns the raw reply bytes (a marshaled -// leaseReply) for the lease loop to decode. -func (t leaseTransport) RequestLease(ctx context.Context, queueID string, data []byte, timeout time.Duration) ([]byte, error) { +// RequestLease sends a work-lease request for queueID on behalf of workerID and +// waits up to timeout for the server's reply. It returns the raw reply bytes (a +// marshaled leaseReply) for the lease loop to decode. +func (t leaseTransport) RequestLease(ctx context.Context, workerID, queueID string, data []byte, timeout time.Duration) ([]byte, error) { reqCtx, cancelReq := context.WithTimeout(ctx, timeout) defer cancelReq() - msg, err := t.nc.RequestWithContext(reqCtx, bus.WorkLeaseSubject(queueID), data) + msg, err := t.nc.RequestWithContext(reqCtx, bus.WorkLeaseSubject(workerID, queueID), data) if err != nil { return nil, fmt.Errorf("worker: request lease for queue %q: %w", queueID, err) } @@ -633,7 +678,7 @@ func (t leaseTransport) RequestLease(ctx context.Context, queueID string, data [ // leaseQueueIDs maps the worker's configured queue list to the queues it // requests leases on. A worker with no configured queues serves any queue; it // must still request on a valid subject, so it uses [bus.WildcardQueueToken] -// (work.lease._any) rather than an empty leaf (work.lease., which routes to no +// (work.lease.._any) rather than an empty queue token, which routes to no // responder). The server selects tasks farm-wide and gates by eligibility, so a // queue-unaffiliated worker is matched to any queue's ready work. func leaseQueueIDs(configured []string) []string { diff --git a/config/sqi-server.example.yaml b/config/sqi-server.example.yaml index c1a54b6c..03deb61a 100644 --- a/config/sqi-server.example.yaml +++ b/config/sqi-server.example.yaml @@ -38,8 +38,10 @@ nats: # Defaults to all interfaces (0.0.0.0) so that workers which discover the # server via mDNS can connect to NATS at the advertised LAN host. Set this # to "127.0.0.1:4222" to restrict NATS to loopback (single-machine only). - # Note: broker authentication does not exist. Any host that can reach this - # port can register as a worker and receive assignments. Deferred to Phase 4. + # Broker authentication is opt-in and OFF by default (see nats.auth below). + # While it is off, any host that can reach this port can register as a + # worker and execute submitted job code. sqi-server warns at startup when + # this address is not loopback and nats.auth is disabled. # Type: string Env: SQI_NATS_ADDR addr: "0.0.0.0:4222" @@ -54,6 +56,33 @@ nats: # Type: int Env: SQI_NATS_MAX_STORE_MB max_store_mb: 1024 + # Broker authentication for worker connections. Opt-in and independent of + # the top-level auth block (auth.enabled) below — that block gates human + # users and API keys; this one gates workers connecting to the broker. + auth: + # Requires every NATS client to present a per-worker nkey credential. + # When false, the broker accepts any connection (the v0.3.0 behavior). + # Default: false + # Type: bool Env: SQI_NATS_AUTH_ENABLED + enabled: false + + # How long a newly issued worker join token remains valid. + # Default: 1h Min: 1m Max: 24h + # Type: duration Env: SQI_NATS_AUTH_JOIN_TOKEN_TTL + join_token_ttl: 1h + + # Consumes a join token on first successful enrollment. Leave true unless + # provisioning many identical machines from one image-baked token. + # Default: true + # Type: bool Env: SQI_NATS_AUTH_JOIN_TOKEN_SINGLE_USE + join_token_single_use: true + + # Mounts POST /api/v1/workers/enroll. Meaningful only when enabled above + # is true; set false at a site that provisions every credential by hand. + # Default: true + # Type: bool Env: SQI_NATS_AUTH_ENROLLMENT_ENDPOINT_ENABLED + enrollment_endpoint_enabled: true + # ── SQLite state store ──────────────────────────────────────────────────────── store: # Path to the SQLite database file. Created at startup if it does not exist. diff --git a/config/sqi-worker.example.yaml b/config/sqi-worker.example.yaml index d3ba4db8..9f42e84f 100644 --- a/config/sqi-worker.example.yaml +++ b/config/sqi-worker.example.yaml @@ -63,6 +63,41 @@ nats: # Type: duration Env: SQI_WORKER_NATS_RECONNECT_WAIT reconnect_wait: "2s" + # Path to this worker's nkey seed file. Only meaningful when the server's + # nats.auth.enabled is true. When empty, defaults to + # /worker.nk, and that default path is created by + # enrollment or by `sqi-worker keygen`. Set this to a non-default path and + # both paths honor it: `sqi-worker keygen` loads the same layered + # configuration `sqi-worker start` does (this file, SQI_WORKER_NATS_CREDENTIAL_FILE, + # and the root -c/--config flag) and writes wherever this field resolves, + # and self-service enrollment writes to the same resolved path as part of + # a live enrollment run. Must be mode 0600; writing it requires write + # permission on the containing directory, not just the file, since the + # write goes through a temp file plus rename. + # Type: string Env: SQI_WORKER_NATS_CREDENTIAL_FILE + credential_file: "" + + # A worker enrollment token, used exactly once on first start to obtain a + # credential. Ignored once credential_file already exists. Prefer + # join_token_file over this field — a token in a config file is a secret + # at rest. + # Type: string Env: SQI_WORKER_NATS_JOIN_TOKEN + join_token: "" + + # Path to a file containing a join token. Takes precedence over join_token. + # Type: string Env: SQI_WORKER_NATS_JOIN_TOKEN_FILE + join_token_file: "" + + # sqi-server HTTP base URL used for enrollment, e.g. + # "http://sqi-server.example:8080". Enrollment runs over REST, not NATS: + # the broker's job is to reject unauthenticated connections, so it cannot + # also be the channel a worker gets its first credential over. REQUIRED + # whenever a join token is configured, and NOT derived from mDNS + # discovery: enrollment fails fast naming this field rather than + # attempting a request with no host. + # Type: string Env: SQI_WORKER_NATS_SERVER_URL + server_url: "" + # ── Worker identity and runtime behavior ───────────────────────────────────── worker: # Human-readable name for this worker shown in the sqi-server web UI and logs. @@ -159,14 +194,20 @@ worker: keep_failed_sessions: false # Restrict this worker to serving specific queue IDs. The worker keeps one - # outstanding lease request per listed queue (work.lease.). When + # outstanding lease request per listed queue (work.lease..). When # empty (the default) it issues a single lease request on the reserved - # core-NATS subject work.lease._any — the server selects tasks farm-wide for + # core-NATS subject work.lease.._any — the server selects tasks farm-wide for # that token and gates by worker eligibility, so a queue-unaffiliated worker # is matched to any queue's ready work. Set this on heterogeneous farms where # some workers specialise in a subset of queues, e.g. separate GPU and CPU # queues. # + # Each entry becomes a literal token in that subject, so it must be usable + # as a single NATS subject token: an entry that is empty, or contains '.', + # whitespace, '*' or '>', is rejected at load (a startup failure) rather + # than producing lease requests the server silently refuses or never + # answers. + # # Type: []string Env: SQI_WORKER_QUEUE_IDS (comma-separated) queue_ids: [] # Examples: diff --git a/docs/api.md b/docs/api.md index dd0eb4e2..11aeba4d 100644 --- a/docs/api.md +++ b/docs/api.md @@ -124,9 +124,16 @@ Browser clients use the session cookie minted by `POST /api/v1/auth/login` instead. Three endpoints are always public because gating them would be circular: `GET /api/v1/openapi.yaml`, `POST /api/v1/auth/login`, and `GET /api/v1/auth/providers` (plus `GET /api/v1/auth/oidc/login` and -`GET /api/v1/auth/oidc/callback` when SSO is configured). - -Every other `/api/v1` REST endpoint requires a permission (`/healthz`, +`GET /api/v1/auth/oidc/callback` when SSO is configured). A fourth, +`POST /api/v1/workers/enroll`, is unauthenticated for the same reason — the +join token in its body is itself the credential — but it is not always +present: it is mounted only when the server's `nats.auth.enabled` and +`nats.auth.enrollment_endpoint_enabled` are both true. See +[Broker authentication](auth.md#broker-authentication-transport) for that +model in full. + +Every other `/api/v1` REST endpoint requires a permission, with the one +exception of `POST /api/v1/workers/enroll` noted above (`/healthz`, `/readyz` and `/metrics` sit outside the API prefix and are never gated; the `/ws` upgrade authenticates the same way but gates per-subject — see [WebSocket subscriptions](#websocket-subscriptions)): @@ -137,6 +144,7 @@ Every other `/api/v1` REST endpoint requires a permission (`/healthz`, | `POST /jobs`, `POST /products/{name}/jobs`, `PATCH/DELETE /jobs/{id}`, `POST /jobs/{id}/cancel`, `POST /jobs/{id}/retry`, `POST /tasks/{id}/retry`, `POST /tasks/{id}/cancel` | `jobs.write` | | `GET /workers`, `GET /workers/{id}` | `workers.read` | | `POST /workers/{id}/disable`, `POST /workers/{id}/enable`, `DELETE /workers/{id}` | `workers.manage` | +| `POST /workers/join-tokens`, `DELETE /workers/{id}/credential` | `workers.enroll` | | `GET` on farms, queues, storage-locations, compute-locations, usage-pools | `infra.read` | | `POST`/`PUT`/`DELETE` on farms, queues, storage-locations, compute-locations, usage-pools | `infra.manage` | | `GET /products`, `GET /products/{name}`, `GET /products/{name}/parameters`, `GET /presets`, `GET /presets/{name}` | `products.read` | @@ -148,6 +156,17 @@ Every other `/api/v1` REST endpoint requires a permission (`/healthz`, | `GET /users/{id}/api-keys`, `DELETE /users/{id}/api-keys/{keyId}` | `apikeys.admin` | | `POST /auth/logout`, `GET/PATCH /auth/me`, `PUT /auth/password`, `GET /version` | any authenticated principal | +`workers.enroll` is deliberately separate from `workers.manage`: minting a +join token or revoking a broker credential attaches or detaches arbitrary +compute, a different privilege in kind from enabling, disabling or deleting +a worker record. `POST /workers/join-tokens` is additionally mounted only +when `auth.enabled` — with authentication off there is no RBAC in front of +it to gate. `DELETE /workers/{id}` — removing the worker record itself — +also cascades a credential revoke through the same synchronous path as +`DELETE /workers/{id}/credential`, before deleting the worker row, so that +deleting a worker never leaves its broker access live with no permission +able to cut it. + Object routes (`/jobs/{id}…`, `/tasks/{id}…`) are additionally owner-scoped: a principal without `jobs.read.all` sees and acts on only its own jobs — the one permission governs both read and write scoping. A missing or rejected diff --git a/docs/architecture.md b/docs/architecture.md index c3829a45..384182b8 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -29,10 +29,10 @@ through scheduling, worker execution, and final state. │ ┌─────────────────────────────▼──────────────────────────────┐ │ │ │ embedded NATS (JetStream + core NATS) │ │ │ │ │ │ -│ │ work.lease. task.status. │ │ -│ │ task.logs. task.cancel. │ │ -│ │ worker.register worker.heartbeat │ │ -│ │ worker.deregister worker.diag. │ │ +│ │ work.lease.. task.status.. │ │ +│ │ task.logs.. task.cancel. │ │ +│ │ worker.register. worker.heartbeat. │ │ +│ │ worker.deregister. worker.diag. │ │ │ └────────┬────────────────────────────────────────────────┬─┘ │ │ │ │ │ │ ┌───────────▼──────────┐ ┌────────────────▼──┐ │ @@ -255,11 +255,11 @@ terminal state. ### 3. Assignment (lease-on-request) Ready tasks stay `ready` until a worker asks for work. Workers keep exactly one -outstanding lease request per queue on `work.lease.` (core-NATS +outstanding lease request per queue on `work.lease..` (core-NATS request/reply). When a request arrives the server: ``` -handleLeaseRequest(queueID, workerID) +handleLeaseRequest(workerID, queueID) │ ├─ store.CommittedCores(workerID, worker.CPUCount) → committed (Σ required_cores of assigned+running tasks) ├─ free = worker.CPUCount − committed @@ -322,12 +322,12 @@ for the config knob. ``` sqi-worker │ - ├─ Keeps one outstanding lease request per queue (work.lease.) + ├─ Keeps one outstanding lease request per queue (work.lease..) │ Long-poll (~35 s timeout); re-issues immediately on return ├─ Receives batch of AssignMsgs from server ├─ Executes each task (spawns child process, manages lifecycle) - ├─ Streams log chunks → NATS task.logs. - └─ Reports status changes → NATS task.status. + ├─ Streams log chunks → NATS task.logs.. + └─ Reports status changes → NATS task.status.. { task_id, attempt_id, status, exit_code, timestamp } ``` @@ -627,20 +627,26 @@ in the first place. | Subject pattern | Transport | Direction | Purpose | |---|---|---|---| -| `work.lease.` | Core NATS request/reply | worker → server (request); server → worker (reply) | Worker requests a batch of tasks; server replies with assignments or empty on timeout | -| `task.status.` | JetStream (`SQI_TASK`, MaxAge 24 h) | worker → server | Terminal and intermediate status updates | -| `task.logs.` | JetStream (`SQI_LOGS`, MaxAge 96 h) | worker → server | Log chunk delivery | +| `work.lease..` | Core NATS request/reply | worker → server (request); server → worker (reply) | Worker requests a batch of tasks; server replies with assignments or empty on timeout | +| `task.status..` | JetStream (`SQI_TASK`, MaxAge 24 h) | worker → server | Terminal and intermediate status updates | +| `task.logs..` | JetStream (`SQI_LOGS`, MaxAge 96 h) | worker → server | Log chunk delivery | | `task.cancel.` | JetStream (`SQI_CANCEL`, MaxAge 5 min) | server → worker | Cancellation signal; the worker holding the task interrupts the process | -| `worker.register` | JetStream (`SQI_WORKER`, MaxAge 2 min) | worker → server | Registration at startup and on reconnect | -| `worker.heartbeat` | JetStream (`SQI_WORKER`, MaxAge 2 min) | worker → server | Liveness heartbeat | -| `worker.deregister` | JetStream (`SQI_WORKER`, MaxAge 2 min) | worker → server | Graceful departure; marks the worker offline without waiting for heartbeat timeout | +| `worker.register.` | JetStream (`SQI_WORKER`, MaxAge 2 min) | worker → server | Registration at startup and on reconnect | +| `worker.heartbeat.` | JetStream (`SQI_WORKER`, MaxAge 2 min) | worker → server | Liveness heartbeat | +| `worker.deregister.` | JetStream (`SQI_WORKER`, MaxAge 2 min) | worker → server | Graceful departure; marks the worker offline without waiting for heartbeat timeout | | `worker.diag.` | Core NATS (best-effort) | worker → server | Diagnostic log records | -A queue-unaffiliated worker leases on the reserved leaf `work.lease._any` -(`bus.WildcardQueueToken`). +Every worker → server subject carries the publishing worker's ID directly after +its class prefix. NATS permissions are static per credential and JetStream does +not stamp publisher identity onto a message, so this placement is what lets the +broker restrict a worker to its own traffic, and what lets the server recover +who published a message it received (`bus.ParseWorkerSubject`). + +A queue-unaffiliated worker leases on the reserved queue token +`work.lease.._any` (`bus.WildcardQueueToken`). JetStream streams use file-backed storage with configurable size limits. -`work.lease.` uses core NATS request/reply — no stream is created for +`work.lease..` uses core NATS request/reply — no stream is created for it. The server holds an unfulfillable request in memory for up to 30 s before replying with an empty batch; the worker re-requests immediately. diff --git a/docs/auth.md b/docs/auth.md index 577b1b48..9f96f22e 100644 --- a/docs/auth.md +++ b/docs/auth.md @@ -18,6 +18,251 @@ Bearer API key (`auth.Chain(apikey, session)`); there is no more means in practice, and [First-admin bootstrap](#first-admin-bootstrap) for how to get your first credential. +## Broker authentication (transport) + +`auth.enabled` gates the HTTP REST API and the WebSocket upgrade. It does +**nothing** to the worker transport. A second, independent switch — +`nats.auth.enabled` (config file `nats.auth.enabled`, env +`SQI_NATS_AUTH_ENABLED`; default `false`) — gates the embedded NATS broker +that workers connect to. The two flags protect different surfaces and must +each be turned on deliberately; flipping one does not flip the other. + +### What it protects, and what it does not + +Broker authentication addresses two things: an unauthorized host attaching to +the broker and registering as a worker, and one enrolled worker forging +traffic that claims to come from another. **It does not encrypt the +transport.** Task payloads, assignment contents and log chunks all travel +over the connection in cleartext whether broker authentication is on or off. +Turning it on stops an attacker from attaching or impersonating another +worker; it does not stop someone who can already read the network from +reading what crosses it. Do not conclude "I turned on broker authentication" +means "the channel is private" — it does not. + +### The credential: an nkey per worker + +Each worker holds an Ed25519 **nkey** keypair. The worker generates it +locally — only the public key ever leaves the machine — and the private seed +is written to `nats.credential_file`, mode `0600`, which defaults to +`/worker.nk` but can be pointed anywhere. Both self-service +enrollment over REST and `sqi-worker keygen` honor a custom path: `keygen` +loads the worker's own configuration (the root `-c/--config` file, +`SQI_WORKER_*` environment variables, and built-in defaults, the same as +`sqi-worker start`) and writes the seed wherever `credential_file` resolves. +See +[`docs/worker-configuration.md`](worker-configuration.md#natscredential_file) +for the full field reference. Connecting to the broker is challenge-response +(the broker sends a nonce, the worker signs it with the seed), so nothing +replayable crosses the wire. + +### Getting a credential + +Two ways, both ultimately calling the same store write: + +- **Join token (self-service).** An operator mints a TTL-bounded token — + `sqi-server worker token issue` (see the database-path resolution rule below), + or `POST /api/v1/workers/join-tokens` when `auth.enabled` is on, gated on + the `workers.enroll` permission (admin-only by default) — and hands it to + the worker via `nats.join_token_file` (preferred) or `nats.join_token`. On + first boot with no credential file present, the worker generates its + keypair and calls `POST /api/v1/workers/enroll` with the token, its worker + ID and its public key. That endpoint is unauthenticated by design — the + token itself is the credential, so gating it on a session or API key would + be circular — and exists only while `nats.auth.enabled` and + `nats.auth.enrollment_endpoint_enabled` are both true. Tokens default to a + 1-hour TTL (bounded 1 minute to 24 hours) and are single-use by default + (`nats.auth.join_token_single_use`). +- **Manual pre-provisioning.** Run `sqi-worker keygen` on the worker host, + with its normal config in place — it writes the seed and prints the + public key and the exact `sqi-server worker enroll --worker-id … --public-key + …` command to run on the server. Pass `--data-dir` to override + `worker.data_dir` explicitly for a one-off run against a different + directory. No token, no REST call, no `POST /api/v1/workers/enroll` route + needs to exist at all. This is the path for an air-gapped worker, or a + site that sets `nats.auth.enrollment_endpoint_enabled: false` and wants no + self-service enrollment surface whatsoever. `sqi-server worker enroll` is + **offline**, exactly like `sqi-server worker revoke`: it writes the + database from a process with no broker handle, so a **running** server + keeps refusing the new key until it restarts. + +> **`sqi-server worker …` resolves its database path the same way `backup` +> and `migrate` do.** Highest priority first: an explicit `--db`; +> `store.sqlite_path` (the root `-c/--config` file and +> `SQI_STORE_SQLITE_PATH`); the legacy `SQI_SQLITE_PATH` environment +> variable; then `sqi.db`. Unlike `migrate`, these subcommands never create +> or migrate the database — pointing one at a path with no database there is +> a clear error naming the resolved path, not a silently created empty one. +> See [`docs/operations.md`](operations.md#worker-broker-credentials) for +> the full command reference. + +**Enroll every worker before flipping `nats.auth.enabled` on.** The enrolled +credential set is loaded once, at server `Start`, and only when broker +authentication is enabled — so a farm that flips the switch with nothing +enrolled yet does not fail closed gracefully, it takes every worker offline +at its next restart or reconnect: a rejected credential is fatal in the +worker, not a silent retry. A worker can hold a credential harmlessly while +`nats.auth.enabled` is still `false` — the broker does not check it — so the +safe order is: enroll every worker first (either path above), confirm each +one has a credential (`sqi-server worker list`), and only then set +`nats.auth.enabled: true` and restart the server. + +**Enrollment always runs over REST, never over NATS.** The broker's entire +job is to refuse unauthenticated connections, so it cannot also be the +channel a worker obtains its first credential over. This means +`nats.server_url` (the server's HTTP base URL) **must be set explicitly on +the worker** whenever a join token is configured — it is **not** derived +from mDNS discovery. A worker relying on mDNS with no `server_url` set fails +enrollment fast, naming that field, rather than attempting a request with no +host. + +See [`docs/worker-deployment.md`](worker-deployment.md) for both paths +walked end to end, and [`docs/worker-configuration.md`](worker-configuration.md) +for every `nats.*` field. + +### Revocation: two paths, deliberately distinct + +- `DELETE /api/v1/workers/{id}/credential` runs inside the server process, + which holds the live broker handle: it writes the store and reloads the + broker's authorized-key set in the same call, so the worker is + disconnected before the request returns. +- `sqi-server worker revoke ` is an **offline** CLI command — it opens + the SQLite file directly and writes the same row, but from a separate + process with no broker handle. It takes effect the next time `sqi-server` + starts, not immediately. + +Use the REST path when you need a worker off the farm right now (a +compromised host, a decommissioned machine). The CLI path is for offline +maintenance — revoking a credential before a server has ever started with +it, or as part of a startup script. + +`DELETE /api/v1/workers/{id}` — removing the worker record itself — revokes +its credential too, through the same synchronous path as the first bullet +above, and does so **before** deleting the worker row, not after. This +exists because `workers.manage` (what deleting a worker requires) does not +imply `workers.enroll` (what revoking a credential directly requires) — the +split is deliberate, so that the ability to delete a worker never doubles as +the ability to mint join tokens — and without the cascade, an operator who +can decommission a machine would have no way at all to cut its broker +access. A worker with no credential (broker authentication disabled, or a +worker that was never enrolled) is deleted exactly as before; a credential +that is already revoked is treated the same way. + +The ordering matters: `store.DeleteWorker` never rejects with a conflict — +removability was already decided by an earlier check — so revoking first +never wastes a revocation on a delete that was always going to be refused. +If the revoke fails, nothing has happened yet: the worker row is intact, the +request answers 500, and it is safe to retry. If the delete then fails after +a successful revoke, the worker row survives but its broker access is +already cut — the safe direction to fail in — and retrying `DELETE +/workers/{id}` simply re-revokes (a no-op the second time) and tries the +delete again. Deleting first and revoking after was tried and rejected: a +failure in the revoke's own store write, not just a broker-reload failure, +would leave the worker row gone, the credential never revoked, and nothing +left to reap it — the operator would be told 204 while the machine kept live +broker access permanently, with no other permission available to fix it. + +### Key rotation and re-enrollment + +Worker-ID uniqueness applies only to **active** credentials: a revoked +worker ID can be enrolled again with a brand-new key (`sqi-worker keygen +--force`, then a fresh `sqi-server worker enroll` or join token). Public +keys, by contrast, are **globally unique forever** — once a key has been +enrolled, even to a since-revoked credential, it can never be enrolled +again. To rotate a worker's key: + +1. `sqi-server worker revoke ` (or the REST path, if the old + credential should be disconnected immediately). +2. `sqi-worker keygen --force` on the worker host, with its normal config in + place — this overwrites the seed with a new keypair and prints the new + public key, plus the exact `sqi-server worker enroll` command for step 3. + `keygen` loads configuration the same way `sqi-worker start` does (the + root `-c/--config` file, `SQI_WORKER_*` environment variables, and + built-in defaults), so running it on the worker host with its own config + resolves `worker.data_dir` and `nats.credential_file` the same way the + worker itself does. `keygen` also states, in its output, whether the + worker ID it printed was loaded from an existing `worker.id` or freshly + generated — a freshly generated ID here means `--data-dir` or + `worker.data_dir` is pointed at the wrong directory, since it should be + rotating the key of the worker already named in step 1, not minting a + new one. Pass `--data-dir` to override the directory explicitly for a + one-off run. +3. `sqi-server worker enroll --worker-id --public-key ` + on the server host, using the public key `keygen` just printed. +4. **Restart `sqi-server`.** Like `worker revoke`, `worker enroll` is an + offline command: it writes the database from a process with no broker + handle, and the broker's authorized-key set is built once at startup. A + running server therefore still refuses the new key until it restarts. +5. Restart `sqi-worker`. It finds the new seed already on disk (`keygen` + wrote it directly) and connects with it — no REST enrollment call happens + on this path, since a credential file already exists. + +Running `keygen --force` before revoking the old credential is safe but +will not help: the worker ID stays bound to the *old* public key on the +server until it is revoked, so the enroll command in step 3 fails until +step 1 has run. To rotate via a fresh join token instead of `keygen`, +remove `worker.nk` before restarting the worker — with no credential file +present it re-enrolls exactly as it did on first boot. + +### What an enrolled worker's credential may do + +A worker's broker permissions are generated from the worker ID recorded at +enrollment, and cover only that worker's own subtree: + +- **Publish:** `task.status..*`, `task.logs..*`, + `worker.register.`, `worker.heartbeat.`, `worker.deregister.`, + `worker.diag.`, `work.lease..*`. +- **Subscribe:** `_INBOX_.>` — this worker's own reply inboxes — and + `task.cancel.>`. + +Everything else is denied. A worker cannot subscribe to `task.status.>` or +`task.logs.>`, so it cannot observe another worker's status or log traffic. + +The reply-inbox prefix is **per worker**, not nats.go's process-global +`_INBOX`. That matters because a work lease is core-NATS request/reply: the +assignment batch — command lines, embedded file contents, job and task +parameters, environment variables, the path map and the run-as-user account +— is delivered to the requester's reply inbox and nothing but the subscribe +permission guards it. A single `_INBOX.>` grant would let any enrolled +worker read every other worker's assignments (and `sqi-server`'s own +JetStream API replies) without leasing a task itself. Each worker therefore +connects with `nats.CustomInboxPrefix("_INBOX_")` and is granted +only that subtree. + +### The auth-off asymmetry + +With broker authentication **off — the permanent default** — the worker ID +is still present as a token in every worker-to-server subject (see +[NATS subjects and streams](architecture.md#nats-subjects-and-streams)), but +**nothing enforces it**. Any client that can reach the broker may publish +under any worker ID it chooses. The scheduler still parses that ID out of +the subject and compares it against the task attempt's recorded owner, and +still discards a mismatch — but with authentication off, that check is not a +security boundary. It catches **honest bugs**, not attackers: a stale +worker, a version mismatch, a client publishing to the wrong subject by +accident. Do not read those provenance checks as proof that cross-worker +forgery is prevented in the default configuration — it is not. Only +`nats.auth.enabled: true`, which authenticates the connection itself, closes +that gap. + +`sqi-server` emits a startup `WARN` when `nats.addr` binds to a non-loopback +address and broker authentication is off, naming both remediations (turn on +`nats.auth.enabled`, or bind `nats.addr` to `127.0.0.1`). + +### An accepted authorization gap: `task.cancel.>` + +This is the **only** subject a worker's credential can reach outside its own +subtree. Workers subscribe to `task.cancel.` per-task, at the moment a task is +assigned, not for the lifetime of the connection. NATS permissions are +static per credential, so a worker's subscribe permission cannot be +narrower than the whole `task.cancel.>` subtree — there is no way to grant +"only the cancel subjects for tasks I currently hold" without a permission +reload on every assignment, or a NATS auth callout. The practical effect: +**any enrolled worker can observe the cancel signal for any task**, not only +its own. This is accepted as low severity — a cancel message for a task the +worker does not hold is simply inert, it triggers no action — but it is a +real, deliberate gap rather than an oversight, and is recorded here rather +than left to be rediscovered. + ## Model Every request carries a `Principal` in its context. When auth is off, the @@ -133,6 +378,7 @@ built-in roles (no custom-role builder — YAGNI): | apikeys.self (own keys) | ✅ | ✅ | ✅ | ✅ | | apikeys.admin (anyone's keys) | ❌ | ❌ | ❌ | ✅ | | `isolation.manage` — set a queue's `run_as_user`/`run_as_group` | ❌ | ❌ | ❌ | ✅ | +| `workers.enroll` — mint worker join tokens; revoke a worker's broker credential | ❌ | ❌ | ❌ | ✅ | `apikeys.admin` is enforced by `GET /users/{id}/api-keys` and `DELETE /users/{id}/api-keys/{keyId}` — see [API keys](#api-keys). @@ -1454,15 +1700,22 @@ provider](development.md#testing-against-a-real-directory-or-identity-provider) ## Known gaps -- **Broker authentication remains absent.** `auth.enabled` gates the HTTP REST - API and the WebSocket upgrade **only**. It does nothing to the worker - transport: any host that can reach the embedded NATS broker's port (`4222` - by default) can register as a worker and receive task assignments, exactly - as if auth were off. There is no plan to change this before Phase 4 — see - the comment on `bus.BrokerConfig.Addr` (`internal/bus/broker.go`). An - operator reading "I flipped `auth.enabled` to `true`, so the server is now - locked down" should read that as "the HTTP/WebSocket surface is now locked - down" — the worker-registration surface is unaffected either way. +- **Broker authentication is off by default, and that default is permanent — + not a placeholder for a later phase.** See + [Broker authentication (transport)](#broker-authentication-transport) + above for the full model: the credential, both enrollment paths, + revocation, key rotation, and — stated plainly there — the asymmetry that + while `nats.auth.enabled` is off, the worker ID carried in every subject is + unenforced, so the scheduler's provenance checks catch bugs rather than + attackers, and any host that can reach the embedded NATS broker's port + (`4222` by default) can register as a worker and receive task assignments. + Turning broker authentication on does not encrypt the transport either — + task payloads, assignments and log chunks stay cleartext regardless. + `sqi-server` emits a startup WARN when the broker address is non-loopback + and broker auth is off, precisely because an operator reading "I flipped + `auth.enabled` to `true`, so the server is now locked down" should not + assume that also locked down the worker transport — the two flags are + independent and both must be set. - **Task isolation is implemented and integration-tested on both platforms.** See [Task isolation](#task-isolation) above for the current state: a queue's `run_as_user` runs job code as a distinct OS user on Linux/macOS workers, and @@ -1487,6 +1740,20 @@ provider](development.md#testing-against-a-real-directory-or-identity-provider) already run as the daemon's own account — so isolation being enabled on Windows is precisely what makes this reachable. Not yet fixed; tracked for a follow-up before this is considered hardened. +- **`DELETE /workers/{id}/credential` stays mounted with `auth.enabled=false`.** + Every other permission-gated route in this document is bypassed by the + anonymous superuser when `auth.enabled` is off — that is unchanged, + by-design behavior. This one route carries a real consequence from it: + with `auth.enabled=false` and `nats.auth.enabled=true` (a supported, + documented combination), an unauthenticated caller can revoke any worker's + broker credential, and repeated calls can take the whole farm off the + broker one worker at a time. It is an availability exposure, not an + escalation — revoke only removes access already granted, it cannot attach + new compute or obtain a credential — and every other destructive worker + route (`disable`, `DELETE /workers/{id}`) is exposed exactly the same way + in that configuration. Carving out this one route would break the rule + that auth-off behavior matches pre-auth sqi, so it is accepted rather than + special-cased. - **Per-user concurrent task caps.** A hard per-owner ceiling on running tasks was scoped for Phase 3 and deferred (2026-07-20) with no driver behind it. Nothing bounds a single user's farm consumption today: `max_concurrent_tasks` on farms and queues caps the container, not the diff --git a/docs/configuration.md b/docs/configuration.md index 22bf51ad..1afe872c 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -138,10 +138,11 @@ sqi-server serve --http-cors-origins=https://ui.example.com TCP address the embedded NATS server binds to. Defaults to all interfaces so that workers which discover the server over mDNS can connect to NATS at the advertised LAN host. Set this to `"127.0.0.1:4222"` to restrict NATS to loopback -(single-machine only). **Broker authentication does not exist**: any host -that can reach this port can register as a worker and receive task -assignments, regardless of `auth.enabled` — see -[Known gaps](auth.md#known-gaps). Deferred to Phase 4 hardening. +(single-machine only). **Broker authentication is a separate, opt-in gate from +`auth.enabled`** — see [`nats.auth.*`](#natsauthenabled) below. While +`nats.auth.enabled` is off (the default), any host that can reach this port +can register as a worker and receive task assignments, regardless of +`auth.enabled` — see [Known gaps](auth.md#known-gaps). ```yaml nats: @@ -189,6 +190,92 @@ nats: --- +### `nats.auth.enabled` + +| | | +|---|---| +| **Type** | `bool` | +| **Default** | `false` | +| **Env var** | `SQI_NATS_AUTH_ENABLED` | + +Requires every NATS client to present a per-worker nkey credential. When +`true`, the broker refuses any connection that does not present one. When +`false` (the default), the broker accepts any connection — see the warning +under [`nats.addr`](#natsaddr) above. Deliberately independent of +`auth.enabled`: that flag gates human users and API keys over the REST API, +this one gates workers over the broker transport. + +```yaml +nats: + auth: + enabled: true +``` + +--- + +### `nats.auth.join_token_ttl` + +| | | +|---|---| +| **Type** | `duration` | +| **Default** | `1h` | +| **Range** | `1m` – `24h` | +| **Env var** | `SQI_NATS_AUTH_JOIN_TOKEN_TTL` | + +How long a newly issued worker join token remains valid. The floor exists +because below a minute an operator cannot realistically get the token onto a +machine and boot it; the ceiling exists because a join token mints a worker +credential, so a token valid for weeks would be a standing secret wearing a +different name. Ignored while `nats.auth.enabled` is `false`. + +```yaml +nats: + auth: + join_token_ttl: 30m +``` + +--- + +### `nats.auth.join_token_single_use` + +| | | +|---|---| +| **Type** | `bool` | +| **Default** | `true` | +| **Env var** | `SQI_NATS_AUTH_JOIN_TOKEN_SINGLE_USE` | + +Consumes a join token on first successful enrollment. Leaving this `true` is +strongly recommended; set it `false` only for image-baked fleets that enroll +many identical machines from one token. + +```yaml +nats: + auth: + join_token_single_use: false +``` + +--- + +### `nats.auth.enrollment_endpoint_enabled` + +| | | +|---|---| +| **Type** | `bool` | +| **Default** | `true` | +| **Env var** | `SQI_NATS_AUTH_ENROLLMENT_ENDPOINT_ENABLED` | + +Mounts `POST /api/v1/workers/enroll`. Meaningful only when +`nats.auth.enabled` is `true`. Set `false` at a site that provisions every +worker credential by hand and wants no enrollment surface at all. + +```yaml +nats: + auth: + enrollment_endpoint_enabled: false +``` + +--- + ## `store` — SQLite state store ### `store.sqlite_path` @@ -208,11 +295,18 @@ store: sqlite_path: "/var/lib/sqi/sqi.db" ``` -> **The `migrate` and `backup` subcommands do not read this key.** Their -> `--db` flag defaults to `$SQI_SQLITE_PATH` (note: *not* -> `SQI_STORE_SQLITE_PATH`), falling back to `sqi.db` in the working -> directory. Pass `--db` explicitly, or export both variables, so schema -> migrations and backups operate on the database the server actually uses. +> **The `migrate`, `backup`, and `worker` subcommands resolve their database +> path the same way the server does.** Highest priority first: an explicit +> `--db` flag; this key (the root `-c/--config` file and +> `SQI_STORE_SQLITE_PATH`); the legacy `SQI_SQLITE_PATH` environment +> variable, which still works but prints a deprecation notice to stderr when +> it is what decided the path; and finally the built-in `"sqi.db"` default. +> `migrate up` creates the database when it does not already exist — that is +> its job. `backup` and `worker` (`sqi-server worker +> token issue|enroll|revoke|list` — see +> [`docs/operations.md`](operations.md#worker-broker-credentials)) never +> create one: pointing either at a database that does not exist is a clear +> error naming the resolved path, not a silently created empty database. --- @@ -1790,6 +1884,10 @@ for the detector schema reference. | `nats.addr` | string | `0.0.0.0:4222` | `SQI_NATS_ADDR` | — | | `nats.data_dir` | string | `data/nats` | `SQI_NATS_DATA_DIR` | — | | `nats.max_store_mb` | int | `1024` | `SQI_NATS_MAX_STORE_MB` | — | +| `nats.auth.enabled` | bool | `false` | `SQI_NATS_AUTH_ENABLED` | — | +| `nats.auth.join_token_ttl` | duration | `1h` | `SQI_NATS_AUTH_JOIN_TOKEN_TTL` | — | +| `nats.auth.join_token_single_use` | bool | `true` | `SQI_NATS_AUTH_JOIN_TOKEN_SINGLE_USE` | — | +| `nats.auth.enrollment_endpoint_enabled` | bool | `true` | `SQI_NATS_AUTH_ENROLLMENT_ENDPOINT_ENABLED` | — | | `store.sqlite_path` | string | `sqi.db` | `SQI_STORE_SQLITE_PATH` | — | | `store.checkpoint_interval` | duration | `5m` | `SQI_STORE_CHECKPOINT_INTERVAL` | — | | `log.level` | string | `info` | `SQI_LOG_LEVEL` | `--log-level` | diff --git a/docs/development.md b/docs/development.md index f0e9b91b..bbd8a789 100644 --- a/docs/development.md +++ b/docs/development.md @@ -14,7 +14,7 @@ guides for extending the worker. | Node.js ≥ 24 with npm ≥ 11 (see `.nvmrc` and `web/package.json` `engines`) | Build the web UI bundle embedded in `sqi-server` (`make build` runs it) | [nodejs.org](https://nodejs.org/) or `nvm use` | | `gofumpt` | Stricter formatter (superset of `gofmt`) | `go install mvdan.cc/gofumpt@latest` | | `goimports` | Import organizer | `go install golang.org/x/tools/cmd/goimports@latest` | -| `golangci-lint` ≥ 2.13.1 (CI pins this exact version; see below) | Linter suite | [golangci-lint.run/usage/install](https://golangci-lint.run/usage/install/) | +| `golangci-lint` ≥ 2.13.0 (see the version floor below) | Linter suite | [golangci-lint.run/usage/install](https://golangci-lint.run/usage/install/) | | `lefthook` | Git hook runner | `go install github.com/evilmartians/lefthook@latest` | | `pkgsite` | Local pkg.go.dev docs server | `go install golang.org/x/pkgsite/cmd/pkgsite@latest` | | Docker (optional) | Build and run the container image; also runs the real-directory LDAP tests (`make test-ldap`), the real-provider SSO tests (`make test-oidc`), and the real-root run-as-user isolation tests (`make test-isolation`), all of which skip cleanly without it | [docs.docker.com](https://docs.docker.com/get-docker/), or `brew install colima docker && colima start` | @@ -22,13 +22,17 @@ guides for extending the worker. `gofumpt`, `goimports`, and `golangci-lint` are required at commit time via pre-commit hooks. Install them before running `make hooks`. -**`golangci-lint` 2.13.1 is a hard floor, not a suggestion.** `.golangci.yml` +**`golangci-lint` 2.13.0 is a hard floor, not a suggestion.** `.golangci.yml` excludes `errors.AsType` from `errcheck` by function name, and `errcheck` before -2.13.1 cannot resolve a *generic* function's name — it reports `Error return +2.13.0 cannot resolve a *generic* function's name — it reports `Error return value is not checked` with no name at all, so the exclusion never matches and every `errors.AsType` call site in the repo is reported. On an older -golangci-lint `make lint` fails on code CI considers clean. CI pins the same -version in `.github/workflows/ci.yml`; keep the two in step when bumping. +golangci-lint `make lint` fails on code CI considers clean. + +CI pins v2.13.1 exactly (`.github/workflows/ci.yml`) rather than the floor, so +the runner never drifts; 2.13.0 is what Homebrew currently ships and lints this +repo clean, which is why the floor sits a patch below the pin. Raise both +together when bumping. --- @@ -476,8 +480,24 @@ ListStepsForJob(ctx context.Context, jobID string) ([]Step, error) ``` Then implement it in `internal/store/sqlite/step.go` using a prepared -statement, and add a corresponding stub to the in-memory fake in -`internal/store/fake/store.go` so existing tests keep compiling. +statement. + +**A wholly new aggregate — a new table, not just a new method on an +existing one — needs a migration.** Add a numbered SQL file to +`internal/store/migrations/` (e.g. `00031_my_feature.sql`, continuing the +existing sequence), with `+goose Up`/`+goose Down` sections following the +pattern of the surrounding files. A new column or index on an existing +table needs one too. A new method on an *existing* table (like +`ListStepsForJob` above, assuming `steps` already exists) needs no schema +change at all. + +Finally, add a corresponding stub to the in-memory fake. The fake is split +per aggregate, mirroring `internal/store/sqlite/`, not one +`internal/store/fake/store.go` file — a `step.go` under `internal/store/fake/` +for the `StepStore` methods above, or a new file named for a wholly new +aggregate (see `internal/store/fake/workercredential.go` for what a new +aggregate's fake looks like end to end). Stub it so existing tests keep +compiling. > **Not every new store method is REST-triggered.** The auto-retry + > failure-limit feature added `RecordTaskFailure`, `RequeueTaskForRetry` @@ -724,8 +744,9 @@ db := t.TempDir() + "/test.db" `go install golangci-lint` method is not supported by the project. **`make lint` reports `Error return value is not checked` on `errors.AsType`** -— your `golangci-lint` predates 2.13.1. Check with `golangci-lint version` and -upgrade; see the version floor noted under [Prerequisites](#prerequisites). +— your `golangci-lint` predates 2.13.0. Check with `golangci-lint version` and +upgrade (`brew upgrade golangci-lint`); see the version floor noted under +[Prerequisites](#prerequisites). **`gofumpt` or `goimports` not found after installing** — ensure `$(go env GOPATH)/bin` is on your `$PATH`: diff --git a/docs/observability.md b/docs/observability.md index 660d4c70..092079ac 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -16,7 +16,7 @@ and follow different flows. |---|---|---| | **What** | stdout/stderr of a task process running under a worker | sqi-server's and sqi-worker's own structured `slog` output | | **Who writes** | The task's child process | The sqi binaries themselves | -| **Transport** | NATS JetStream `task.logs.` | Core NATS `worker.diag.` (best-effort); server logs go in-process | +| **Transport** | NATS JetStream `task.logs..` | Core NATS `worker.diag.` (best-effort); server logs go in-process | | **Persistence** | Durable; retained ~96 h in JetStream, also written to SQLite | In-memory ring buffer on the server; lost on server restart | | **Where to read** | Task detail → log tab in the web UI | Worker detail panel, Admin → Server log, or REST/WS API | | **Disable** | Not configurable (always streamed for running tasks) | Server: `SQI_DIAGNOSTICS_BUFFER_SIZE=0`; worker: `SQI_DIAGNOSTICS_ENABLED=false` | @@ -25,7 +25,7 @@ and follow different flows. When a worker executes a task it captures the process's stdout and stderr through a `logstreamer` and publishes them in chunks to the JetStream subject -`task.logs.`. The server's consumer writes each chunk to SQLite and +`task.logs..`. The server's consumer writes each chunk to SQLite and fans it out over WebSocket so the task log page updates live. Because JetStream retains messages, you can reload the log page and still see all output. diff --git a/docs/operations.md b/docs/operations.md index 68f10aa2..6763c6b7 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -214,14 +214,19 @@ sqi-server backup \ --out /backups/sqi/sqi-$(date +%Y%m%d-%H%M%S).db ``` -The `--db` flag defaults to `$SQI_SQLITE_PATH` (or `sqi.db` if that variable is -unset). Note that this is a different environment variable from `SQI_STORE_SQLITE_PATH` -used by the running server; set `--db` explicitly or export `SQI_SQLITE_PATH` to -match your deployment. - -The command opens the source database read-only and writes an identical clean -copy to the destination path. It exits non-zero if the destination file already -exists — use a timestamped filename or a fresh directory each time. +`--db` is optional: omitted, the source path resolves through the same +config layer the server uses (the root `-c/--config` file and +`SQI_STORE_SQLITE_PATH`, i.e. `store.sqlite_path`), falling back to the +legacy `SQI_SQLITE_PATH` environment variable and then to `sqi.db` in the +working directory. Passing `--db` explicitly, as in the example above, +always wins. + +The command opens the source database without applying migrations (SQLite +still opens it read-write and may create `-wal`/`-shm` sidecar files) and +writes an identical clean copy to the destination path. It exits non-zero if +the source database does not exist (it never creates one) or if the +destination file already exists — use a timestamped filename or a fresh +directory each time. ### Automated daily backup (cron) @@ -270,6 +275,78 @@ To restore from a backup: --- +## Worker broker credentials + +`sqi-server worker` groups the offline CLI commands for NATS broker +authentication (`nats.auth.*` — see +[Broker authentication](auth.md#broker-authentication-transport) for the +full model). Like `migrate` and `backup`, these subcommands open the SQLite +database file directly and do not start an HTTP server or NATS broker, so +they work whether or not `sqi-server` is running — and, independently, +whether or not the user-facing `auth.enabled` is on. + +> **These commands resolve the database path the same way `backup` and +> `migrate` do** — an explicit `--db`, then `store.sqlite_path` (the root +> `-c/--config` file and `SQI_STORE_SQLITE_PATH`), then the legacy +> `SQI_SQLITE_PATH` environment variable, then `sqi.db`. Unlike `migrate`, +> they never create or migrate the database: pointing one at a path with no +> database there is a clear error naming the resolved path, not a silently +> created empty one. Run `sqi-server migrate up` against the target database +> first if it does not exist yet. + +### Issue a join token + +```sh +sqi-server worker token issue --db /data/sqi.db --ttl 1h +``` + +Prints the raw token to stdout exactly once — capture it +(`TOKEN=$(sqi-server worker token issue --db /data/sqi.db)`) or store it +securely; only its hash is kept in the database. `--ttl` defaults to `1h` +and is bounded 1 minute to 24 hours. Hand the token to a worker via +`nats.join_token_file` (preferred) or `nats.join_token`, or mint one over +REST instead with `POST /api/v1/workers/join-tokens` when `auth.enabled` is +also on. + +### Enroll a worker manually + +```sh +sqi-server worker enroll --db /data/sqi.db \ + --worker-id 3f2a... --public-key UABC...XYZ +``` + +Registers a worker's broker credential directly, by worker ID and public +key — the offline counterpart to self-service enrollment over REST. Run +`sqi-worker keygen` on the worker host first; it prints this exact command +with the worker's own ID and public key filled in. A **running** +`sqi-server` does not see the new credential until it restarts — the broker +builds its authorized-key set once at startup, and this command writes the +database from a separate process with no broker handle. + +### Revoke a worker's credential + +```sh +sqi-server worker revoke --db /data/sqi.db +``` + +Revokes a worker credential in the database. Takes effect the next time +`sqi-server` starts, not immediately — to disconnect a worker at once +against a running server, use `DELETE /api/v1/workers/{id}/credential` +instead. + +### List worker credentials + +```sh +sqi-server worker list --db /data/sqi.db +``` + +Lists every worker credential that has not been revoked: worker ID, name, +public key, enrollment time, and last-seen time. Last-seen is set on worker +registration (startup and reconnect) only, never by heartbeat — it answers +"when did this worker last (re)connect", not "is it up right now". + +--- + ## Log management See [`docs/observability.md`](observability.md) for the full observability diff --git a/docs/roadmap.md b/docs/roadmap.md index 72281231..55b14561 100644 --- a/docs/roadmap.md +++ b/docs/roadmap.md @@ -79,7 +79,7 @@ Configuration cascades: farm defaults → queue overrides, with retry policy (ma **Scheduling considers:** job priority, task dependencies, queue and farm policy (concurrency limits, scheduling mode), compute location affinity, worker capability tags (OS, GPU, installed software), and usage pool availability. - *Design:* ready tasks remain `ready` until a worker sends a core-NATS - request to `work.lease.`. The server computes free cores + request to `work.lease..`. The server computes free cores (`CPUCount − Σ committed`), selects a priority-ordered batch that fits, atomically transitions the batch `ready → assigned` (stamping `assigned_at` only now), and replies. The `SQI_WORK` JetStream stream, `work.assign.` @@ -212,7 +212,7 @@ NATS JetStream handles: - Heartbeats and worker registration Work leases use **core NATS** request/reply (not JetStream): the worker requests -work on `work.lease.` and the server replies with a batch it is +work on `work.lease..` and the server replies with a batch it is authorized to run (pull-based). Real-time UI updates reach web clients over WebSocket, fanned out by the server after it ingests the JetStream messages. diff --git a/docs/worker-configuration.md b/docs/worker-configuration.md index 2fdf9a25..79b1767a 100644 --- a/docs/worker-configuration.md +++ b/docs/worker-configuration.md @@ -140,6 +140,19 @@ nats: max_reconnect_attempts: -1 ``` +> **`0` weakens a diagnostic message, not just reconnect behavior.** When +> broker authentication is on and the server rejects this worker's +> credential mid-run (most often after `sqi-server worker revoke` or +> `DELETE /api/v1/workers/{id}/credential`), the worker names the cause and +> the remediation — but doing so relies on at least one reconnect attempt +> completing its handshake with the broker, which is where the specific +> rejection reason is confirmed. With `max_reconnect_attempts: 0`, no +> reconnect is attempted at all: the connection simply closes, and the +> operator sees only a generic "connection closed" with no named cause. The +> default of `-1` does not have this problem — reconnect indefinitely and +> the crafted diagnostic message is always reached before the worker exits. +> If you must cap reconnect attempts, keep it above `0`. + --- ### `nats.reconnect_wait` @@ -161,6 +174,104 @@ nats: --- +### `nats.credential_file` + +| | | +|---|---| +| **Type** | `string` | +| **Default** | `""` (resolves to `/worker.nk`) | +| **Env var** | `SQI_WORKER_NATS_CREDENTIAL_FILE` | + +Path to this worker's nkey seed file. Only meaningful when the server's +`nats.auth.enabled` is `true` — see +[`docs/configuration.md`](configuration.md#natsauthenabled). The worker +presents the credential in this file to authenticate to the broker. When +left empty, it defaults to `/worker.nk`, and that default +path is created by enrollment or by `sqi-worker keygen`. The file must be +mode `0600`, and `SaveSeed` requires **write permission on the containing +directory**, not just the file itself — the write goes through a temporary +file created alongside the target and then renamed into place, so a +directory locked down to e.g. `0500` fails the write even if the seed file +inside it is owner-writable. + +Set this field to a non-default path and both enrollment paths honor it: +`sqi-worker keygen` loads the same layered configuration `sqi-worker start` +does (the root `-c/--config` file, `SQI_WORKER_NATS_CREDENTIAL_FILE`, and +this field's config-file value) and writes wherever `credential_file` +resolves, and self-service enrollment writes to the same resolved path as +part of a live enrollment run. + +```yaml +nats: + credential_file: "/var/lib/sqi-worker/worker.nk" +``` + +--- + +### `nats.join_token` + +| | | +|---|---| +| **Type** | `string` | +| **Default** | `""` | +| **Env var** | `SQI_WORKER_NATS_JOIN_TOKEN` | + +A worker enrollment token, used exactly once on first start to obtain a +credential. Ignored once `nats.credential_file` already exists. Prefer +[`nats.join_token_file`](#natsjoin_token_file) over this field — a token in a +config file is a secret at rest. + +```yaml +nats: + join_token: "" +``` + +--- + +### `nats.join_token_file` + +| | | +|---|---| +| **Type** | `string` | +| **Default** | `""` | +| **Env var** | `SQI_WORKER_NATS_JOIN_TOKEN_FILE` | + +Path to a file containing a join token. Takes precedence over +[`nats.join_token`](#natsjoin_token). + +```yaml +nats: + join_token_file: "/run/secrets/sqi-join-token" +``` + +--- + +### `nats.server_url` + +| | | +|---|---| +| **Type** | `string` | +| **Default** | `""` | +| **Env var** | `SQI_WORKER_NATS_SERVER_URL` | + +`sqi-server` HTTP base URL used for enrollment, e.g. +`"http://sqi-server.example:8080"`. Enrollment runs over REST, not NATS: the +broker's job is to reject unauthenticated connections, so it cannot also be +the channel a worker gets its first credential over. + +This is **not** derived from mDNS discovery — it must be set explicitly +whenever [`nats.join_token`](#natsjoin_token) or +[`nats.join_token_file`](#natsjoin_token_file) is configured. A worker that +needs to enroll with no `server_url` set fails fast with an actionable error +naming this field, rather than attempting a request with no host. + +```yaml +nats: + server_url: "http://sqi-server.example:8080" +``` + +--- + ## `worker` — Identity and runtime behavior ### `worker.name` @@ -479,14 +590,23 @@ worker: | **Env var** | `SQI_WORKER_QUEUE_IDS` (comma-separated) | Restrict this worker to serving specific queue IDs. The worker keeps one -outstanding lease request per listed queue (`work.lease.`). When +outstanding lease request per listed queue (`work.lease..`). When empty (the default), the worker issues a single lease request on the reserved -subject `work.lease._any` — an empty leaf would produce the invalid subject -`work.lease.` with no responders. The server selects tasks farm-wide for that +subject `work.lease.._any` — an empty queue token would produce an +unroutable subject with no responders. The server selects tasks farm-wide for that token and gates by worker eligibility, so a queue-unaffiliated worker is matched to any queue's ready work. Set this on heterogeneous farms where some workers specialise in a subset of queues. +Each entry must be usable as a single NATS subject token, because it becomes +one: `` is a literal token in `work.lease..`, so +an entry containing `.` would silently split into extra tokens the +server-side parser and this worker's own broker-auth publish grant both +reject. An entry that is empty, or contains `.`, whitespace, `*` or `>`, is +therefore **rejected at load**, naming the offending entry: the worker +refuses to start rather than issuing lease requests that would be silently +refused or simply go unanswered with nothing in its logs saying why. + ```yaml worker: queue_ids: @@ -1463,6 +1583,10 @@ log_streamer: | `nats.insecure_skip_verify` | bool | `false` | `SQI_WORKER_NATS_INSECURE_SKIP_VERIFY` | `--nats-insecure-skip-verify` | | `nats.max_reconnect_attempts` | int | `-1` | `SQI_WORKER_NATS_MAX_RECONNECT_ATTEMPTS` | — | | `nats.reconnect_wait` | duration | `2s` | `SQI_WORKER_NATS_RECONNECT_WAIT` | — | +| `nats.credential_file` | string | `""` (`/worker.nk`) | `SQI_WORKER_NATS_CREDENTIAL_FILE` | — | +| `nats.join_token` | string | `""` | `SQI_WORKER_NATS_JOIN_TOKEN` | — | +| `nats.join_token_file` | string | `""` | `SQI_WORKER_NATS_JOIN_TOKEN_FILE` | — | +| `nats.server_url` | string | `""` | `SQI_WORKER_NATS_SERVER_URL` | — | | `worker.name` | string | hostname | `SQI_WORKER_NAME` | — | | `worker.farm_id` | string | `""` | `SQI_WORKER_FARM_ID` | — | | `worker.data_dir` | string | `~/.sqi/worker` (Linux/macOS); `%USERPROFILE%\.sqi\worker` (Windows) | `SQI_WORKER_DATA_DIR` | — | @@ -1554,12 +1678,15 @@ long as three things differ per instance: | Setting | Env var | Why it must differ | |---|---|---| -| [`worker.data_dir`](#workerdata_dir) | `SQI_WORKER_DATA_DIR` | Holds the persistent `worker.id` UUID; a shared dir means a duplicate identity on the server. | +| [`worker.data_dir`](#workerdata_dir) | `SQI_WORKER_DATA_DIR` | Holds the persistent `worker.id` UUID and (when `nats.credential_file` is left at its default) the worker's nkey seed; a shared dir means a duplicate identity on the server, or two instances fighting over one credential. | | [`metrics.addr`](#metricsaddr) | `SQI_WORKER_METRICS_ADDR` | The local health/metrics HTTP server; a second instance on the same port fails to bind. | | [`worker.name`](#workername) | `SQI_WORKER_NAME` | Cosmetic only — defaults to the hostname, so instances would otherwise share a label in the web UI. | Everything else (NATS URL, discovery, capability tags) can be shared or vary -as you like. +as you like. If broker authentication is on and each instance enrolls with a +join token, each also needs its own token: tokens are single-use by default +(`nats.auth.join_token_single_use`), so the second instance to redeem a +shared token fails enrollment. For local development the `make run-workers` target wires all of this up for you — see diff --git a/docs/worker-deployment.md b/docs/worker-deployment.md index 5cdd54c7..3fbe55aa 100644 --- a/docs/worker-deployment.md +++ b/docs/worker-deployment.md @@ -100,6 +100,162 @@ available option. --- +## Broker authentication + +Broker authentication (`nats.auth.enabled` on the server) is opt-in and off +by default; if the server you are joining does not have it enabled, skip +this section and connect as shown above with no credential. If it does, the +worker needs a credential before it can connect at all — see +[`docs/auth.md`](auth.md#broker-authentication-transport) for the full model +(what it protects, revocation, key rotation). This section walks the two +ways to obtain one. + +**Both paths need `nats.server_url` set on the worker to the server's HTTP +base URL** (e.g. `http://sqi-server.example:8080`). Enrollment always runs +over REST, never over NATS — the broker's whole job is to refuse +unauthenticated connections, so it cannot also hand out the first +credential. `nats.server_url` is **not** derived from mDNS discovery even +when `discovery.enable_mdns` is on for locating the NATS broker itself; a +worker with a join token configured and no `server_url` set fails fast, +naming the field, rather than guessing a host. + +### Path A — self-service enrollment with a join token + +1. On the server host, mint a token: + + ```sh + sqi-server worker token issue --ttl 1h + ``` + + The raw token prints to stdout exactly once — capture it + (`TOKEN=$(sqi-server worker token issue)`) or store it securely; only its + hash is kept in the database and it cannot be displayed again. If + `auth.enabled` is also on, an operator can mint one over REST instead — + `POST /api/v1/workers/join-tokens`, gated on the `workers.enroll` + permission — but the CLI command above always works, whether or not + `auth.enabled` is on. + +2. Hand the token to the worker via a file (preferred — a token in the + config file itself is a secret at rest) and set the server URL: + + ```yaml + nats: + join_token_file: "/run/secrets/sqi-join-token" + server_url: "http://sqi-server.example:8080" + ``` + +3. Start the worker normally. On first boot, finding no credential file at + `nats.credential_file` (default `/worker.nk`), it + generates an nkey keypair locally and calls `POST /api/v1/workers/enroll` + with the token, its worker ID and its new **public** key — the private + seed never leaves the machine, and nothing from the server's response is + persisted. Only after the server confirms enrollment does the worker + write the seed it generated itself to `worker.nk` (mode `0600`) and + connect with it. The token is single-use by default, so a second worker + needs its own token. + +Every failure past this point is fatal and explicit — an expired, unknown or +already-used token, or a worker ID already bound to a different key — the +worker logs the cause and exits rather than looping in the background with +no visible reason. An unreachable *server* is different: that still falls +back to the ordinary reconnect-with-backoff behavior once a credential +exists. + +### Path B — manual pre-provisioning + +For a worker that cannot reach the server's REST API (air-gapped hosts), or +a site that wants no self-service enrollment endpoint at all +(`nats.auth.enrollment_endpoint_enabled: false` on the server): + +1. On the worker host, generate a keypair without connecting anywhere: + + ```sh + sqi-worker keygen --data-dir /var/lib/sqi-worker + ``` + + `keygen` loads the worker's own configuration the same way `sqi-worker + start` does (the root `-c/--config` file, `SQI_WORKER_*` environment + variables, and built-in defaults), so `--data-dir` here is an explicit + override of `worker.data_dir` — pass it when generating a keypair before + the worker's own config is in place, or to target a different directory + for a one-off run. With the worker's normal config already set up, a bare + `sqi-worker keygen` resolves the same directory the worker itself uses. + + This writes `/var/lib/sqi-worker/worker.nk` (mode `0600`) and prints the + worker's public key and the exact command to run next: + + ``` + Public key: UABC...XYZ + Worker ID: 3f2a... (newly generated; if you expected an existing worker id here, --data-dir/worker.data_dir is probably pointed at the wrong directory) + On the server, run: + sqi-server worker enroll --worker-id 3f2a... --public-key UABC...XYZ + A RUNNING sqi-server will not accept this credential until it restarts; to enroll against a running server, use POST /api/v1/workers/enroll with a join token instead. + ``` + +2. Copy that command to the server host (out of band — SSH, a console, a + provisioning script) and run it there: + + ```sh + sqi-server worker enroll --worker-id 3f2a... --public-key UABC...XYZ + ``` + +3. **Restart `sqi-server`.** `sqi-server worker enroll` writes the database + from a separate process with no broker handle, and the broker builds its + authorized-key set once at startup — so a *running* server does not know + about the credential just written, and a worker presenting it is refused + with an authorization error (which the worker treats as fatal, by + design). If restarting the server is not acceptable, enroll over REST + with a join token instead (Path A), which reloads the running broker in + the same request. + +4. Start the worker. It finds the seed already on disk, skips enrollment + entirely, and connects directly with the credential. + +No join token and no REST call are involved in this path at any point. + +### Revoking a worker's credential + +- **Immediately**, against a running server: `DELETE + /api/v1/workers/{id}/credential`. This disconnects the worker inside the + call — its in-flight leases reclaim through the normal heartbeat-sweep + path. +- **Offline**, without a running server (or as part of a maintenance + script): `sqi-server worker revoke `. This writes the database + directly and takes effect the next time `sqi-server` starts, not before. + +### Rotating a compromised or lost key + +1. Revoke the current credential — `sqi-server worker revoke `, + or the REST path above for an immediate disconnect. +2. `sqi-worker keygen --force` on the worker host, with the worker's normal + config in place (pass `--data-dir` to override `worker.data_dir` + explicitly for a one-off run against a different directory). This + overwrites `worker.nk` with a new keypair and prints the new public key, + whether the worker ID is the existing one or newly generated — a freshly + generated one here is the signal that the resolved data directory does + not match the worker being rotated — and the `sqi-server worker enroll` + command for step 3. It must run *after* step 1, since the worker ID stays + bound to the old public key on the server until that credential is + revoked. +3. `sqi-server worker enroll --worker-id --public-key ` + on the server host, using the key `keygen` printed. +4. **Restart `sqi-server`**, for the reason given in Path B step 3: this + offline command cannot reload a running broker's authorized-key set, so + until the server restarts it still refuses the new key. To rotate + without a server restart, remove `worker.nk` and re-enroll over REST + with a fresh join token instead (see the note below). +5. Restart the worker. It finds the new seed `keygen` already wrote and + connects directly — no join token or REST enrollment call is needed on + this path. + +The worker ID can be reused once its old credential is revoked; the old +public key itself can never be enrolled again, on this worker ID or any +other. To rotate via a fresh join token instead of `keygen`, remove +`worker.nk` before restarting so the worker re-enrolls as it did on first +boot (Path A above). + +--- + ## Linux — systemd ### 1. Create a dedicated user @@ -481,6 +637,7 @@ accordingly. ## See also - [`docs/worker-configuration.md`](worker-configuration.md) — Every configuration option. +- [`docs/auth.md`](auth.md#broker-authentication-transport) — The broker authentication model in full: threat model, revocation, key rotation, and the auth-off asymmetry. - [`docs/worker-capabilities.md`](worker-capabilities.md) — Capability tag reference. - [`docs/worker-docker.md`](worker-docker.md) — Docker deployment details. - [`docs/observability.md`](observability.md) — In-UI diagnostic panels, REST/WS log API, and log forwarding to journald, Docker, Loki, and ELK. diff --git a/docs/worker-docker.md b/docs/worker-docker.md index b6f5e41b..ae532132 100644 --- a/docs/worker-docker.md +++ b/docs/worker-docker.md @@ -76,7 +76,17 @@ Strongly recommended when mounting a data volume: | Variable | Description | |---|---| -| `SQI_WORKER_DATA_DIR` | Set to the volume mount path (`/var/lib/sqi-worker`) so the worker ID is written to the persistent volume rather than the default `~/.sqi/worker` | +| `SQI_WORKER_DATA_DIR` | Set to the volume mount path (`/var/lib/sqi-worker`) so the worker ID is written to the persistent volume rather than the default `~/.sqi/worker`. This **must** be a persistent volume once broker authentication is in use: the worker's nkey seed lives under this directory too (`/worker.nk` by default), and losing the volume does not just lose the identity — the worker's already-enrolled credential still blocks re-enrollment under the same worker ID until an operator revokes it. | + +Required in addition to `SQI_WORKER_NATS_URL` when the server has broker +authentication on (`nats.auth.enabled`): + +| Variable | Description | +|---|---| +| `SQI_WORKER_NATS_SERVER_URL` | The server's HTTP base URL (e.g. `http://sqi-server:8080`), used for enrollment over REST. Required whenever a join token is configured — it is not derived from mDNS discovery. | +| `SQI_WORKER_NATS_JOIN_TOKEN_FILE` | Path to a file containing a join token, mounted as a secret. Preferred over `SQI_WORKER_NATS_JOIN_TOKEN` — a token in an environment variable is visible via the container's inspect output. | +| `SQI_WORKER_NATS_JOIN_TOKEN` | A join token supplied directly. Ignored once a credential file already exists at `SQI_WORKER_NATS_CREDENTIAL_FILE`. | +| `SQI_WORKER_NATS_CREDENTIAL_FILE` | Path to the worker's nkey seed file. Defaults to `/worker.nk`; only meaningful under a mounted, persistent `SQI_WORKER_DATA_DIR`. | Useful optional variables for container deployments: @@ -86,7 +96,7 @@ Useful optional variables for container deployments: | `SQI_WORKER_CAPABILITY_TAGS` | Comma-separated capability tags merged with auto-detected ones, e.g. `maya-2025,gpu` | | `SQI_WORKER_COMPUTE_LOCATION` | Compute-location name for multi-site farms (see [Compute locations](compute-locations.md)) | | `SQI_WORKER_FARM_ID` | Restrict the worker to a single farm (empty = accept tasks from any farm) | -| `SQI_WORKER_QUEUE_IDS` | Comma-separated queue IDs the worker serves (empty = all queues) | +| `SQI_WORKER_QUEUE_IDS` | Comma-separated queue IDs the worker serves (empty = all queues). Each ID becomes a token in the `work.lease..` subject, so an entry that is empty or contains `.`, whitespace, `*` or `>` is rejected at startup. | | `SQI_WORKER_LOG_FORMAT` | `json` (default) or `text` | | `SQI_DIAGNOSTICS_ENABLED` | `true` (default) mirrors the worker's own logs to the server's diagnostics view; set `false` to disable | @@ -161,11 +171,12 @@ config file (they have no environment-variable form). The worker container must be able to reach the **NATS port** of `sqi-server` (default `4222`). No inbound connections from the server to the worker are -required — all communication is worker-initiated over NATS. +required — all communication is worker-initiated. | Direction | Protocol | Port | Purpose | |---|---|---|---| | Worker → Server | TCP | `4222` | NATS — core-NATS work leases (task assignments) plus JetStream (task status, logs, heartbeat, registration) | +| Worker → Server | TCP | `8080` (default `http.addr`) | REST — only when the server's `nats.auth.enabled` is on and a join token is configured: a one-time `POST /api/v1/workers/enroll` before the worker ever connects to NATS | | (optional) Prometheus → Worker | TCP | `9091` | Metrics scraping — only if metrics are exposed | ### Docker networking diff --git a/go.mod b/go.mod index 20d8abbb..800f505e 100644 --- a/go.mod +++ b/go.mod @@ -14,9 +14,11 @@ require ( github.com/minio/minio-go/v7 v7.2.1 github.com/nats-io/nats-server/v2 v2.14.5 github.com/nats-io/nats.go v1.53.1 + github.com/nats-io/nkeys v0.4.16 github.com/pressly/goose/v3 v3.27.3 github.com/prometheus/client_golang v1.24.1 github.com/spf13/cobra v1.10.2 + github.com/spf13/pflag v1.0.10 golang.org/x/crypto v0.55.0 golang.org/x/oauth2 v0.36.0 golang.org/x/sys v0.47.0 @@ -48,7 +50,6 @@ require ( github.com/minio/md5-simd v1.1.2 // indirect github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect github.com/nats-io/jwt/v2 v2.8.2 // indirect - github.com/nats-io/nkeys v0.4.16 // indirect github.com/nats-io/nuid v1.0.1 // indirect github.com/ncruces/go-strftime v1.0.0 // indirect github.com/philhofer/fwd v1.2.0 // indirect @@ -59,7 +60,6 @@ require ( github.com/rs/xid v1.6.0 // indirect github.com/ryszard/goskiplist v0.0.0-20150312221310-2dfbae5fcf46 // indirect github.com/sethvargo/go-retry v0.4.0 // indirect - github.com/spf13/pflag v1.0.10 // indirect github.com/tinylib/msgp v1.6.1 // indirect github.com/zeebo/xxh3 v1.1.0 // indirect go.shabbyrobe.org/gocovmerge v0.0.0-20230507111327-fa4f82cfbf4d // indirect diff --git a/internal/README.md b/internal/README.md index 8f751cca..da2ee00e 100644 --- a/internal/README.md +++ b/internal/README.md @@ -32,7 +32,7 @@ expression leaves `internal/openjd/expr` and `internal/openjd/intrange` exist. |---|---| | `internal/api` | REST surface: chi router, handlers, request/response wire types, error shape | | `internal/auth` | Opt-in authentication and authorization: local passwords, sessions, API keys, RBAC policy, role mapping, LDAP/AD and OIDC/SSO | -| `internal/bus` | Embedded NATS broker (JetStream streams for task status/logs/cancel and worker registration/heartbeat/deregister; plain core NATS for `work.lease.` request/reply and `worker.diag.`) and the typed client wrapper over it | +| `internal/bus` | Embedded NATS broker (JetStream streams for task status/logs/cancel and worker registration/heartbeat/deregister; plain core NATS for `work.lease..` request/reply and `worker.diag.`) and the typed client wrapper over it | | `internal/config` | Layered runtime configuration (defaults → file → env → flags) and validation | | `internal/diag` | Bounded in-memory ring buffer of diagnostic (operational) log records from the server and connected workers | | `internal/discovery` | mDNS responder that advertises the running server on the local network | diff --git a/internal/api/authz_integration_test.go b/internal/api/authz_integration_test.go index 7696e5e9..09b1fa86 100644 --- a/internal/api/authz_integration_test.go +++ b/internal/api/authz_integration_test.go @@ -161,6 +161,14 @@ var expectedRoutes = []routeExpectation{ {method: http.MethodPost, pattern: "/api/v1/workers/{id}/enable", perm: policy.WorkersManage}, {method: http.MethodDelete, pattern: "/api/v1/workers/{id}", perm: policy.WorkersManage}, + // workers.enroll. POST /api/v1/workers/enroll is NOT listed here: it is + // unauthenticated by design (the join token is the credential) and, in + // this suite's router, never even mounted — it requires + // Deps.NATSAuthEnabled, which authRouterWith does not set. See + // workerenroll_test.go for its own coverage. + {method: http.MethodPost, pattern: "/api/v1/workers/join-tokens", perm: policy.WorkersEnroll}, + {method: http.MethodDelete, pattern: "/api/v1/workers/{id}/credential", perm: policy.WorkersEnroll}, + // infra.read / infra.manage (farms, queues, storage-locations, // compute-locations, usage-pools) {method: http.MethodGet, pattern: "/api/v1/farms", perm: policy.InfraRead}, @@ -225,12 +233,13 @@ type routeKey struct{ method, pattern string } // placeholder instead of panicking on a nil receiver. func authRouterWith(st store.Store, mutate func(*Deps)) chi.Router { deps := Deps{ - Store: st, - Products: product.NewCatalog(st), - Auth: session.New(st, "sqi_session", nil), - SessionTTL: time.Hour, - CookieName: "sqi_session", - CookieSecure: "false", + Store: st, + Products: product.NewCatalog(st), + Auth: session.New(st, "sqi_session", nil), + SessionTTL: time.Hour, + CookieName: "sqi_session", + CookieSecure: "false", + WorkerRevoker: storeRevoker{store: st}, } if mutate != nil { mutate(&deps) @@ -366,7 +375,7 @@ func TestAuthz_RouteSweep_DisabledAuthAllowsAll(t *testing.T) { st := fake.New() r := NewRouter( Config{DisableRateLimit: true}, - Deps{Store: st, Products: product.NewCatalog(st), Auth: auth.Anonymous(), CookieName: "sqi_session"}, + Deps{Store: st, Products: product.NewCatalog(st), Auth: auth.Anonymous(), CookieName: "sqi_session", WorkerRevoker: storeRevoker{store: st}}, newTestLogger(), metrics.New(), health.NewRegistry(), ) srv := httptest.NewServer(r) diff --git a/internal/api/openapi.yaml b/internal/api/openapi.yaml index 94e652f6..d325536b 100644 --- a/internal/api/openapi.yaml +++ b/internal/api/openapi.yaml @@ -29,9 +29,9 @@ info: Every route under `/api/v1` except `POST /auth/login`, `POST /auth/logout`, `GET /auth/me`, `GET /auth/providers`, `GET - /auth/oidc/login`, `GET /auth/oidc/callback`, `GET /version`, and `GET - /openapi.yaml` is gated on the caller's role holding a specific permission (`jobs.read`, - `jobs.write`, `workers.read`, `workers.manage`, `infra.read`, + /auth/oidc/login`, `GET /auth/oidc/callback`, `GET /version`, `GET + /openapi.yaml`, and `POST /workers/enroll` is gated on the caller's role holding a specific permission (`jobs.read`, + `jobs.write`, `workers.read`, `workers.manage`, `workers.enroll`, `infra.read`, `infra.manage`, `products.read`, `products.manage`, `diagnostics.read`, `users.read`, `users.manage`, `apikeys.self`, `apikeys.admin`) — see the `read-only`/`user`/`operator`/`admin` role matrix. A request from an @@ -44,6 +44,26 @@ info: not unauthenticated either: it runs its own upgrade hook (a valid principal is still required) plus a subject-level `diagnostics.read` check before allowing a subscription to the diagnostics feed. + + `POST /workers/enroll` is unauthenticated for a different reason than the + login-related routes above: the join token in its request body is itself + the credential that authorizes the call, so requiring a session or API + key would be circular — a worker enrolling for the first time has + neither. It is also the only route in this API mounted conditionally on + server configuration rather than always-present-but-gated: it exists + only when broker authentication (`nats.auth.enabled`) is on and the + operator has left `nats.auth.enrollment_endpoint_enabled` at its + default. When either is off, the route is not registered at all — but a + request to it does not receive `404`: `/workers/enroll` has the same + two-segment shape as `/workers/{id}` (`GET`/`DELETE /workers/{id}` are + always registered), so it is routed there with `id=enroll` and answers + `405 Method Not Allowed` instead — the path matches a registered + pattern, just not for `POST`. `POST /workers/join-tokens`, which mints + the tokens this endpoint consumes, is permission-gated as normal + (`workers.enroll`) but is itself mounted only when `auth.enabled` — with + auth off, every request is the anonymous superuser described above, and + minting a join token attaches arbitrary compute, a materially larger + blast radius than every other superuser-bypassed action. version: "0.3.0" license: name: AGPL-3.0-or-later @@ -811,6 +831,98 @@ components: type: string enum: [online, offline, disabled] + # ── Worker enrollment & broker credentials ─────────────────────────────── + WorkerEnrollRequest: + type: object + required: [join_token, worker_id, public_key] + properties: + join_token: + type: string + description: >- + The raw join token minted by POST /workers/join-tokens (or + `sqi-server worker token issue`). This IS the credential that + authorizes the call — there is no other authentication on this + route. + worker_id: + type: string + description: >- + The worker's stable ID. Must be a single NATS subject token: + non-empty, and containing no ".", whitespace, "*" or ">" — + this value becomes a subject pattern in the credential's broker + grants, so a wildcard would authorize publishing as any worker. + Rejected with 400 if it is not; rejected with 409 if it already + has an active credential (re-enrolling the same worker ID with + the same public key also conflicts — worker ID uniqueness is + enforced among active credentials regardless of which key is + given), or if the public key is already enrolled to another + worker. + public_key: + type: string + description: The worker's nkey public key, starting with "U". + name: + type: string + description: Optional human-readable label for the credential. + + WorkerCredential: + type: object + required: [id, worker_id, public_key, enrolled_at] + properties: + id: + type: string + worker_id: + type: string + public_key: + type: string + name: + type: string + enrolled_at: + type: string + format: date-time + last_seen_at: + type: string + format: date-time + nullable: true + description: >- + Set on worker registration (startup and reconnect) only — never + updated by heartbeats. Do not read this as a liveness signal for + a currently-connected worker; it answers "when did this worker + last (re)connect", not "is it up right now". + revoked_at: + type: string + format: date-time + nullable: true + + WorkerJoinTokenCreate: + type: object + properties: + name: + type: string + description: Optional human-readable label for the token. + + WorkerJoinTokenCreated: + type: object + required: [id, token, prefix, expires_at, created_at] + properties: + id: + type: string + token: + type: string + description: >- + The raw join token. Shown once, at creation, and never returned + again — only its hash is stored. Pass it as `join_token` to + POST /workers/enroll. + prefix: + type: string + description: Leading characters of the raw token, for identification in a list. Never enough to reconstruct the token. + name: + type: string + expires_at: + type: string + format: date-time + created_at: + type: string + format: date-time + # ── Diagnostics ─────────────────────────────────────────────────────────── DiagnosticRecord: type: object @@ -2609,6 +2721,17 @@ paths: live-disabled workers are rejected with `409`. Task and attempt history referencing the worker is preserved by ID. A removed worker that reconnects simply re-registers. + + Also revokes the worker's broker credential, if it has one, BEFORE + the worker record is deleted, so a decommissioned machine loses + broker access along with its worker record rather than keeping the + ability to connect, lease work and execute job code. A worker with + no credential (broker authentication disabled, or never enrolled) is + removed exactly as before. This does not require the + `workers.enroll` permission that `DELETE /workers/{id}/credential` + needs on its own. If the revoke fails, the worker record is NOT + deleted and this answers `500` — safe to retry, since nothing has + happened yet. responses: "204": description: Worker removed. @@ -2672,6 +2795,190 @@ paths: "500": $ref: "#/components/responses/InternalServerError" + /workers/enroll: + post: + operationId: enrollWorker + summary: Enroll a worker with a join token + tags: [workers] + description: | + Exchanges a valid join token for a broker credential. Unauthenticated + by design — the join token in the request body is itself the + credential that authorizes this call, so no session or API key is + required or accepted. + + Mounted only when broker authentication (`nats.auth.enabled`) is on + and `nats.auth.enrollment_endpoint_enabled` has not been turned off. + When not mounted, this does NOT return `404`: the path + "/workers/enroll" has the same two-segment shape as + "/workers/{id}", which is always registered (`GET`, `DELETE`), so + the request is routed there with `id=enroll` and answers `405 + Method Not Allowed` instead — the pattern matches, just not for + `POST`. + + An unknown join token, an expired one, and (when single-use tokens + are configured) an already-used one all return the same `401` with + the same body: this route is unauthenticated and may be reachable + from an untrusted network, so distinguishing those cases would let a + caller enumerate which tokens exist or whether one has been claimed. + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/WorkerEnrollRequest" + responses: + "201": + description: Worker enrolled; the credential now exists. + content: + application/json: + schema: + $ref: "#/components/schemas/WorkerCredential" + "400": + description: >- + Malformed request body, missing fields, an invalid public key, or + a worker_id that is not a single NATS subject token. + content: + application/problem+json: + schema: + $ref: "#/components/schemas/ProblemDetail" + "401": + description: >- + The join token is unknown, expired, or already used — or, when + single-use tokens are configured, the store failed while + claiming the token. That last case is deliberately folded into + this same 401 rather than answered as a 500: this route is + unauthenticated and reachable from an untrusted network, so + distinguishing "invalid token" from "transient server-side + failure" would let a caller probe which is which. A client + should treat 401 here as "get a new token from an operator", not + as "retry the same request" — this route never answers 500. + content: + application/problem+json: + schema: + $ref: "#/components/schemas/ProblemDetail" + "405": + description: >- + The enrollment endpoint is not mounted (see the operation + description for why this is 405, not 404). + content: + application/problem+json: + schema: + $ref: "#/components/schemas/ProblemDetail" + "409": + description: >- + The worker ID already has an active credential, or the public + key is already enrolled to another worker. The join token + supplied in the request is NOT consumed by this response: the + claim and the credential write happen in one transaction, so a + conflict here rolls the claim back and the same token remains + redeemable for a request that does not conflict. + content: + application/problem+json: + schema: + $ref: "#/components/schemas/ProblemDetail" + + /workers/join-tokens: + post: + operationId: createWorkerJoinToken + summary: Mint a worker join token + tags: [workers] + description: | + Mints a new join token for self-service worker enrollment over + POST /workers/enroll. The raw token is returned exactly once, in + this response; only its hash is stored, so it cannot be recovered or + displayed again. + + Requires `workers.enroll`, a permission deliberately separate from + `workers.manage`: minting a join token attaches arbitrary compute + that receives and executes job code, a different privilege in kind + from enabling, disabling, or deleting an existing worker. This route + is additionally mounted only when `auth.enabled` — see the + Authorization section above. + requestBody: + required: false + content: + application/json: + schema: + $ref: "#/components/schemas/WorkerJoinTokenCreate" + responses: + "201": + description: Join token created. + content: + application/json: + schema: + $ref: "#/components/schemas/WorkerJoinTokenCreated" + "400": + $ref: "#/components/responses/BadRequest" + "401": + $ref: "#/components/responses/Unauthorized" + "403": + $ref: "#/components/responses/Forbidden" + "500": + $ref: "#/components/responses/InternalServerError" + + /workers/{id}/credential: + parameters: + - $ref: "#/components/parameters/id" + + delete: + operationId: revokeWorkerCredential + summary: Revoke a worker's broker credential + tags: [workers] + description: | + Soft-revokes the active broker credential for the worker named by + `id`. Requires `workers.enroll` (see POST /workers/join-tokens). + + A worker that was never enrolled and a worker whose credential is + already revoked are indistinguishable at the store layer, so both + return `404` with a body that says so rather than claiming the + worker does not exist. + + Revocation is **synchronous**: this call writes the revocation to + the database and then reloads the running broker's authorized-key + set, and nats-server re-authorizes every connected client inside + that reload. An already-connected worker is therefore disconnected + before the `204` is returned, and its in-flight leases reclaim + through the ordinary heartbeat-sweep path. + + If the store write succeeds but the broker reload fails, the + response is `500`: the credential is gone from the store while the + broker may still be honoring it, which the caller has to know + about. Retry the request, or restart `sqi-server`. + + With broker authentication off (`nats.auth.enabled: false`) there is + no authorized-key set to reload, so this only writes the store — + there is nothing to disconnect a worker from, because the broker + never authenticated it in the first place. + + The offline CLI equivalent, `sqi-server worker revoke `, has no + broker handle and does NOT disconnect a running worker — it takes + effect at the next server start. + responses: + "204": + description: >- + Credential revoked and the worker disconnected from the broker. + "401": + $ref: "#/components/responses/Unauthorized" + "403": + $ref: "#/components/responses/Forbidden" + "404": + description: >- + No active credential for this worker — it may never have been + enrolled, or its credential may already be revoked. + content: + application/problem+json: + schema: + $ref: "#/components/schemas/ProblemDetail" + "500": + description: >- + The store write or the broker reload failed. The credential may + be revoked in the store while the broker still honors it; retry, + or restart sqi-server. + content: + application/problem+json: + schema: + $ref: "#/components/schemas/ProblemDetail" + # ── Diagnostics ──────────────────────────────────────────────────────────── /diagnostics/logs: diff --git a/internal/api/router.go b/internal/api/router.go index 44edeab0..30c0de49 100644 --- a/internal/api/router.go +++ b/internal/api/router.go @@ -214,6 +214,49 @@ type Deps struct { // cheaper than a key to configure, distribute, and rotate. Only read when // OIDCProvider is non-nil. OIDCStateKey []byte + + // NATSAuthEnabled mirrors config.NATSAuthConfig.Enabled: the broker + // requires a per-worker nkey credential. POST /api/v1/workers/enroll is + // mounted only when this AND NATSAuthEnrollmentEndpointEnabled are both + // true — with broker auth off there is no credential for it to issue, and + // an operator who wants no self-service enrollment surface at all can + // turn EnrollmentEndpointEnabled off independently. + NATSAuthEnabled bool + + // NATSAuthEnrollmentEndpointEnabled mirrors + // config.NATSAuthConfig.EnrollmentEndpointEnabled. See NATSAuthEnabled. + NATSAuthEnrollmentEndpointEnabled bool + + // JoinTokenTTL is how long a join token minted by + // POST /api/v1/workers/join-tokens remains valid. Mirrors + // config.NATSAuthConfig.JoinTokenTTL, which is bounds-checked at config + // load — this handler does not re-validate it. + JoinTokenTTL time.Duration + + // JoinTokenSingleUse consumes a join token on its first successful + // enrollment, rejecting a second attempt with the same token. Mirrors + // config.NATSAuthConfig.JoinTokenSingleUse. + JoinTokenSingleUse bool + + // WorkerRevoker handles DELETE /api/v1/workers/{id}/credential. Required + // on any router that mounts REST resource routes — that route is always + // registered, regardless of NATSAuthEnabled — so a nil value panics the + // first time it is called, the same contract every other required Deps + // field carries. Production supplies *server.Server, which revokes in + // the store and then reloads the broker's authorized-key set; this + // package depends only on the [WorkerRevoker] interface, never on + // internal/bus or internal/server. + WorkerRevoker WorkerRevoker + + // BrokerCredentialReloader handles the post-enrollment reload inside + // POST /api/v1/workers/enroll. Only actually invoked when that route is + // mounted (NATSAuthEnabled && NATSAuthEnrollmentEndpointEnabled both + // true), but production always supplies it unconditionally — same + // *server.Server as WorkerRevoker — so a nil value here indicates a + // wiring bug rather than an intentionally-disabled feature. See + // [BrokerCredentialReloader] (the interface) for why enroll's reload + // failure is handled differently from revoke's. + BrokerCredentialReloader BrokerCredentialReloader } // resolveCORSOrigins returns the CORS allow-list to configure, dropping the @@ -377,7 +420,7 @@ func NewRouter(cfg Config, deps Deps, logger *slog.Logger, m *metrics.Metrics, h jobs := newJobHandler(deps.Store, deps.Submitter, deps.Scheduler, notifier, logger, retryDefaults, cfg.ValidateJobOwner, cfg.ExprSubmissionDeadline) tasks := newTaskHandler(deps.Store, deps.Scheduler, logger) - workers := newWorkerHandler(deps.Store, notifier, cfg.WorkerOfflineThreshold, logger) + workers := newWorkerHandler(deps.Store, notifier, deps.WorkerRevoker, cfg.WorkerOfflineThreshold, logger) farms := newFarmHandler(deps.Store, logger) queues := newQueueHandler(deps.Store, logger) storageLocs := newStorageLocationHandler(deps.Store, logger) @@ -407,6 +450,7 @@ func NewRouter(cfg Config, deps Deps, logger *slog.Logger, m *metrics.Metrics, h } usersH := newUsersHandler(deps.Store, logger, deps.LDAPConfig.RoleSource, deps.OIDCConfig.RoleSource) apiKeysH := newAPIKeysHandler(deps.Store, logger) + workerEnroll := newWorkerEnrollHandler(deps.Store, deps.WorkerRevoker, deps.BrokerCredentialReloader, logger, deps.JoinTokenSingleUse, deps.JoinTokenTTL) az := newAuthz(deps.Store, logger) wsH := newWSHandler(logger, deps.Hub, deps.Store, deps.Auth, wsOriginConfig{ @@ -464,6 +508,26 @@ func NewRouter(cfg Config, deps Deps, logger *slog.Logger, m *metrics.Metrics, h api.Get("/auth/oidc/callback", authH.oidcCallback) } + // Public, unauthenticated by design: the join token carried in the + // request body is itself the credential that authorizes an + // enrollment, so gating this route on a session or API key would be + // circular — a worker enrolling for the first time has neither. + // Mounted only when broker authentication is on (otherwise there is + // no credential for it to issue) and the operator has left the + // enrollment endpoint enabled (some sites provision every credential + // by hand via "sqi-server worker enroll" and want no self-service + // surface at all). + // + // When NOT mounted, a request here does not 404: "/workers/enroll" + // has the same two-segment shape as "/workers/{id}" (GET and DELETE + // are always registered on that pattern, in the groups below), so + // chi routes it there with id="enroll" and answers 405 Method Not + // Allowed — the path matches a registered pattern, just not for + // POST. openapi.yaml documents 405 for this reason, not 404. + if deps.NATSAuthEnabled && deps.NATSAuthEnrollmentEndpointEnabled { + api.Post("/workers/enroll", workerEnroll.enroll) + } + // REST resource routes — gated by the auth middleware. api.Group(func(rest chi.Router) { // CSRF must run before Auth: it only cares about the presence of @@ -560,6 +624,47 @@ func NewRouter(cfg Config, deps Deps, logger *slog.Logger, m *metrics.Metrics, h g.Delete("/workers/{id}", workers.removeWorker) }) + // workers.enroll — deliberately separate from workers.manage: + // minting a join token attaches arbitrary compute that receives + // and executes job code, a different privilege in kind from + // enabling, disabling, or deleting an existing worker. + // + // Join-token minting is additionally mounted only when + // auth.enabled. Every other permission-gated route in this file + // stays mounted with auth off, because the anonymous Superuser + // principal middleware.Auth installs in that case is granted + // everything by design (auth-off behavior must be unchanged from + // pre-auth sqi). That equivalence is the wrong default here: with + // no RBAC actually enforced, anyone who can reach this server + // would be able to mint a credential for arbitrary compute, which + // is a materially different blast radius than every other + // Superuser-bypassed action. + // + // Revoke stays mounted unconditionally, like every other + // permission-gated route, and that choice is NOT risk-free: with + // auth.enabled false and nats.auth.enabled true — a supported, + // documented combination — every caller is the anonymous + // Superuser, so an unauthenticated request can disconnect a + // worker, and repeated requests can take the whole farm off the + // broker one worker at a time. The risk it does not carry is + // escalation: revoke only removes access already granted, it + // cannot attach compute or obtain a credential. That is an + // AVAILABILITY exposure, deliberately accepted here because + // carving out this one route would break the rule that auth-off + // behavior matches pre-auth sqi, and because every other + // destructive worker route (disable, delete) is exposed exactly + // the same way in that configuration. + if cfg.AuthEnabled { + rest.Group(func(g chi.Router) { + g.Use(az.require(policy.WorkersEnroll)) + g.Post("/workers/join-tokens", workerEnroll.createJoinToken) + }) + } + rest.Group(func(g chi.Router) { + g.Use(az.require(policy.WorkersEnroll)) + g.Delete("/workers/{id}/credential", workerEnroll.revokeCredential) + }) + // infra.read / infra.manage (farms, queues, storage, compute, usage-pools) rest.Group(func(g chi.Router) { g.Use(az.require(policy.InfraRead)) diff --git a/internal/api/workerenroll.go b/internal/api/workerenroll.go new file mode 100644 index 00000000..26d5c29b --- /dev/null +++ b/internal/api/workerenroll.go @@ -0,0 +1,421 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package api + +// Worker broker-credential REST handlers: self-service enrollment, join-token +// minting, and credential revocation. +// +// POST /api/v1/workers/enroll — unauthenticated; the join +// token itself is the credential +// POST /api/v1/workers/join-tokens — workers.enroll +// DELETE /api/v1/workers/{id}/credential — workers.enroll +// +// The revoke handler delegates to an injected [WorkerRevoker] rather than +// writing the store directly, and enroll delegates to an injected +// [BrokerCredentialReloader] after it writes the credential. internal/api +// never holds a live broker handle — the process that does (internal/server) +// supplies implementations that write the store and then reload the +// broker's authorized-key set, so a worker that loses (or gains) a +// credential is disconnected (or made connectable) synchronously, inside +// the same request. This package depends only on the two narrow interfaces, +// never on internal/bus or internal/server, so it cannot import either and +// stays testable without a live broker. + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "log/slog" + "net/http" + "strings" + "time" + + "github.com/go-chi/chi/v5" + "github.com/google/uuid" + + "github.com/uberware/sqi/internal/auth" + "github.com/uberware/sqi/internal/auth/jointoken" + "github.com/uberware/sqi/internal/brokerauth" + "github.com/uberware/sqi/internal/store" +) + +// WorkerRevoker revokes a worker's broker credential and, when the +// implementation holds a live broker handle, disconnects it and lets the +// existing heartbeat-sweep/reclaim path return its in-flight work to ready. +// internal/api depends only on this interface so that DELETE +// /api/v1/workers/{id}/credential can be synchronous where a broker handle +// is available (internal/server, which runs in the same process) without +// this package importing internal/bus or internal/server itself. +type WorkerRevoker interface { + // RevokeWorker revokes workerID's active credential. It returns + // [store.ErrNotFound] if the worker has no active credential — never + // enrolled, or already revoked; those two cases are indistinguishable + // here for the same reason store.RevokeWorkerCredential collapses them. + RevokeWorker(ctx context.Context, workerID string) error +} + +// BrokerCredentialReloader re-syncs a running broker's authorized-key set +// with the store's active worker_credentials rows. Separate from +// [WorkerRevoker] on purpose: it is the enrollment side of the same +// underlying operation, triggered by a different event (a credential +// created, not revoked) and with a different failure posture — see enroll's +// use of it below. internal/api depends only on this interface, never on +// internal/bus or internal/server, for the same reason WorkerRevoker does. +type BrokerCredentialReloader interface { + // ReloadBrokerCredentials re-reads the active credential set from the + // store and reloads it into the broker's authorized-key set, so a + // worker just enrolled can connect to a RUNNING broker without an + // operator restarting it. + ReloadBrokerCredentials(ctx context.Context) error +} + +// errInvalidJoinToken is returned for every way a join token can fail to +// authorize an enrollment — unknown, expired, or already used when single-use +// is on. The endpoint is unauthenticated and may be internet-reachable, so +// distinguishing those cases in the response would let a caller enumerate +// which join tokens exist, or whether one has already been claimed. +const errInvalidJoinToken = "invalid or expired join token" //nolint:gosec // G101: a static response string, not a credential + +// workerEnrollHandler implements the worker broker-credential REST surface. +type workerEnrollHandler struct { + store store.Store + revoker WorkerRevoker + reloader BrokerCredentialReloader + logger *slog.Logger + + // singleUse mirrors config.NATSAuthConfig.JoinTokenSingleUse: whether an + // already-used join token is rejected on a second enrollment attempt. + singleUse bool + + // joinTokenTTL mirrors config.NATSAuthConfig.JoinTokenTTL: how long a + // token minted by createJoinToken remains valid. Already bounds-checked + // at config load, so it is not re-validated here. + joinTokenTTL time.Duration +} + +// newWorkerEnrollHandler returns a workerEnrollHandler wired to the given +// store, revoker, and reloader. +func newWorkerEnrollHandler(st store.Store, revoker WorkerRevoker, reloader BrokerCredentialReloader, logger *slog.Logger, singleUse bool, joinTokenTTL time.Duration) *workerEnrollHandler { + return &workerEnrollHandler{ + store: st, + revoker: revoker, + reloader: reloader, + logger: logger, + singleUse: singleUse, + joinTokenTTL: joinTokenTTL, + } +} + +// ── Wire-format types ─────────────────────────────────────────────────────── + +// workerEnrollRequest is the body of POST /api/v1/workers/enroll. +type workerEnrollRequest struct { + JoinToken string `json:"join_token"` + WorkerID string `json:"worker_id"` + PublicKey string `json:"public_key"` + Name string `json:"name,omitempty"` +} + +// workerCredentialResponse is the JSON representation of a +// [store.WorkerCredential]. It never carries the seed or any other secret — +// only the public key the credential was enrolled with. +type workerCredentialResponse struct { + ID string `json:"id"` + WorkerID string `json:"worker_id"` + PublicKey string `json:"public_key"` + Name string `json:"name,omitempty"` + EnrolledAt time.Time `json:"enrolled_at"` + LastSeenAt *time.Time `json:"last_seen_at,omitempty"` + RevokedAt *time.Time `json:"revoked_at,omitempty"` +} + +func toWorkerCredentialResponse(c store.WorkerCredential) workerCredentialResponse { + return workerCredentialResponse{ + ID: c.ID, WorkerID: c.WorkerID, PublicKey: c.PublicKey, Name: c.Name, + EnrolledAt: c.EnrolledAt, LastSeenAt: c.LastSeenAt, RevokedAt: c.RevokedAt, + } +} + +// workerJoinTokenCreateRequest is the body of POST /api/v1/workers/join-tokens. +// The body itself is optional — an empty POST mints an unnamed token with the +// operator-configured default TTL. +type workerJoinTokenCreateRequest struct { + Name string `json:"name,omitempty"` +} + +// workerJoinTokenCreatedResponse is the create-only shape carrying the raw +// token. This is the only place the raw value is ever returned; only its hash +// is stored, so it cannot be recovered or displayed again. +type workerJoinTokenCreatedResponse struct { + ID string `json:"id"` + Token string `json:"token"` + Prefix string `json:"prefix"` + Name string `json:"name,omitempty"` + ExpiresAt time.Time `json:"expires_at"` + CreatedAt time.Time `json:"created_at"` +} + +// ── POST /api/v1/workers/enroll ───────────────────────────────────────────── + +// enroll exchanges a valid join token for a broker credential. Unauthenticated +// by design: the join token supplied in the body is itself the credential +// that authorizes this call. +func (h *workerEnrollHandler) enroll(w http.ResponseWriter, r *http.Request) { + ctx := r.Context() + // enroll is the only route in this API reachable by an unauthenticated, + // potentially internet-facing caller: the per-IP rate limiter bounds + // request RATE, not per-request body SIZE, so an oversized join_token or + // public_key would otherwise be read fully into memory before any + // validation runs. Matches the jobs.go/queues.go precedent. + body, err := io.ReadAll(io.LimitReader(r.Body, 4<<20)) // 4 MiB cap + if err != nil { + writeProblem(w, r, http.StatusBadRequest, "failed to read request body") + return + } + var req workerEnrollRequest + if err := json.Unmarshal(body, &req); err != nil { + writeProblem(w, r, http.StatusBadRequest, "invalid JSON body") + return + } + req.JoinToken = strings.TrimSpace(req.JoinToken) + req.WorkerID = strings.TrimSpace(req.WorkerID) + req.PublicKey = strings.TrimSpace(req.PublicKey) + req.Name = strings.TrimSpace(req.Name) + if req.JoinToken == "" || req.WorkerID == "" || req.PublicKey == "" { + writeProblem(w, r, http.StatusBadRequest, "join_token, worker_id, and public_key are required") + return + } + + // Validate the key and the worker ID BEFORE the token is claimed, so a + // malformed request costs a 400 and not the operator's single-use token. + if err := brokerauth.ValidatePublicKey(req.PublicKey); err != nil { + writeProblem(w, r, http.StatusBadRequest, err.Error()) + return + } + // The recorded worker ID is what this credential's broker grants are + // built from (brokerauth.WorkerPermissions), and those grants are NATS + // subject PATTERNS. A worker_id of "*" would mint a credential allowed + // to publish "task.status.*.*", "worker.deregister.*", "work.lease.*.*" + // and the rest — concrete subjects belonging to ANY worker, so it could + // forge status and logs, deregister the farm, and lease work as another + // worker and receive that worker's assignment batch. The scheduler's + // provenance checks cannot catch that: the subject NATS vouches for + // genuinely names the victim. A worker_id of ">" is worse-shaped still, + // putting the malformed "task.status.>.*" into the broker's key set, + // which nats-server rejects outright. + if !brokerauth.ValidWorkerIDToken(req.WorkerID) { + writeProblem(w, r, http.StatusBadRequest, + "worker_id must be a single NATS subject token: non-empty, and containing no '.', whitespace, '*' or '>'") + return + } + + now := time.Now().UTC() + cred := store.WorkerCredential{ + ID: uuid.NewString(), + WorkerID: req.WorkerID, + PublicKey: req.PublicKey, + Name: req.Name, + EnrolledAt: now, + } + + var created store.WorkerCredential + if h.singleUse { + created, err = h.redeemSingleUse(ctx, jointoken.Hash(req.JoinToken), now, cred) + } else { + created, err = h.redeemReusable(ctx, jointoken.Hash(req.JoinToken), now, cred) + } + if err != nil { + if errors.Is(err, store.ErrConflict) { + writeProblem(w, r, http.StatusConflict, + "worker already has an active credential, or this public key is already enrolled to another worker") + return + } + // Unknown, expired, already claimed, and a store failure all deny + // identically — see errInvalidJoinToken. + writeProblem(w, r, http.StatusUnauthorized, errInvalidJoinToken) + return + } + + h.finishEnrollment(ctx, req.WorkerID) + + writeJSON(w, http.StatusCreated, toWorkerCredentialResponse(created)) +} + +// redeemSingleUse claims the join token hashed to hash and creates cred, in +// ONE transaction ([store.WorkerCredentialStore.RedeemWorkerJoinToken]). +// +// The token claim and the credential creation must be atomic together, not +// merely atomic individually: claiming the token first and creating the +// credential second — as two separate store calls — would burn a single-use +// token on a request that fails afterwards with a conflicting worker ID or +// public key, leaving the caller with nothing and the operator issuing a +// new token for no reason. A single transaction makes that failure roll +// back the claim too, so the token survives a rejected enrollment attempt +// and remains redeemable by a later, non-conflicting one. +// +// [store.ErrConflict] is returned to the caller as-is, distinct from every +// other failure, so enroll can answer 409 rather than folding it into +// errInvalidJoinToken's 401 — the same distinction CreateWorkerCredential's +// own conflict used to draw when this was two separate calls. +func (h *workerEnrollHandler) redeemSingleUse(ctx context.Context, hash string, now time.Time, cred store.WorkerCredential) (store.WorkerCredential, error) { + created, err := h.store.RedeemWorkerJoinToken(ctx, hash, now, cred) + if err != nil && !errors.Is(err, store.ErrNotFound) && !errors.Is(err, store.ErrConflict) { + h.logger.ErrorContext(ctx, "workerenroll: redeem join token failed", slog.Any("error", err)) + } + return created, err +} + +// redeemReusable authorizes this enrollment against a non-single-use join +// token and then creates cred as a separate call. With single-use disabled +// the token stays redeemable by design — UsedAt is only a "last redeemed" +// marker, not a claim — so there is nothing that needs the two writes to be +// one transaction: a failure creating the credential leaves the token +// exactly as redeemable as it already was, which is the correct outcome for +// a token whose whole point is to be used more than once. +func (h *workerEnrollHandler) redeemReusable(ctx context.Context, hash string, now time.Time, cred store.WorkerCredential) (store.WorkerCredential, error) { + token, err := h.store.GetWorkerJoinTokenByHash(ctx, hash) + if err != nil { + if !errors.Is(err, store.ErrNotFound) { + h.logger.ErrorContext(ctx, "workerenroll: join token lookup failed", slog.Any("error", err)) + } + return store.WorkerCredential{}, err + } + // Strictly after, matching RedeemWorkerJoinToken's "expires_at > ?". + if !token.ExpiresAt.After(now) { + return store.WorkerCredential{}, store.ErrNotFound + } + if err := h.store.MarkWorkerJoinTokenUsed(ctx, token.ID, now); err != nil { + h.logger.ErrorContext(ctx, "workerenroll: mark join token used failed", + slog.String("token_id", token.ID), slog.Any("error", err)) + } + + created, err := h.store.CreateWorkerCredential(ctx, cred) + if err != nil { + if !errors.Is(err, store.ErrConflict) { + h.logger.ErrorContext(ctx, "workerenroll: create credential failed", slog.Any("error", err)) + } + return store.WorkerCredential{}, err + } + return created, nil +} + +// finishEnrollment performs the side effect that follows a successful +// credential creation: reloading the broker's authorized-key set. A failure +// is logged and swallowed rather than turned into an error response — the +// credential itself is already created and durable by the time this runs, +// so telling the caller enrollment failed would be false. Split out of +// enroll to keep that handler's own branching within this repo's complexity +// budget. +// +// Redeeming the join token is NOT done here: it has to happen before the +// credential is created, as one atomic claim — see [claimJoinToken]. +// +// The reload failure here is the opposite direction from a revoke's reload +// failure (which IS surfaced to ITS caller): a revoke failing to reload +// leaves the broker too PERMISSIVE (still trusting a credential the store +// says is gone), which its caller needs to know about; an enroll failing to +// reload leaves the broker too STRICT (a valid worker just can't connect +// yet), which is safe but not self-correcting — there is no background +// reconciliation, so recovery depends on some later enroll or revoke +// triggering another reload, and on an idle farm the new worker stays +// unable to connect until the server restarts. It is not recoverable by the +// worker simply retrying, either: a credential the broker rejects is fatal +// in the worker (it exits rather than looping on reconnect), so getting it +// connected after this failure relies on an external process supervisor +// restarting it, or an operator restarting sqi-server — see +// [BrokerCredentialReloader]. +func (h *workerEnrollHandler) finishEnrollment(ctx context.Context, workerID string) { + if err := h.reloader.ReloadBrokerCredentials(ctx); err != nil { + h.logger.ErrorContext(ctx, "workerenroll: reload broker credentials failed after enrollment", + slog.String("worker_id", workerID), slog.Any("error", err)) + } +} + +// ── POST /api/v1/workers/join-tokens ──────────────────────────────────────── + +// createJoinToken mints a new join token an operator hands to a worker for +// self-service enrollment over POST /api/v1/workers/enroll. The raw token is +// returned exactly once, in this response; only its hash is ever stored. +func (h *workerEnrollHandler) createJoinToken(w http.ResponseWriter, r *http.Request) { + ctx := r.Context() + // Same body-size cap as enroll, for consistency across the two handlers + // on this route group — see enroll's comment for why it matters there. + body, err := io.ReadAll(io.LimitReader(r.Body, 4<<20)) // 4 MiB cap + if err != nil { + writeProblem(w, r, http.StatusBadRequest, "failed to read request body") + return + } + var req workerJoinTokenCreateRequest + // The body is optional (an empty POST mints an unnamed token), so an + // empty body is not itself an error — only malformed JSON is. + if len(body) > 0 { + if err := json.Unmarshal(body, &req); err != nil { + writeProblem(w, r, http.StatusBadRequest, "invalid JSON body") + return + } + } + req.Name = strings.TrimSpace(req.Name) + + rawToken, hash, prefix, err := jointoken.Generate() + if err != nil { + h.logger.ErrorContext(ctx, "workerenroll: generate join token failed", slog.Any("error", err)) + writeProblem(w, r, http.StatusInternalServerError, "failed to create join token") + return + } + now := time.Now().UTC() + createdBy := "" + if p, ok := auth.FromContext(ctx); ok { + createdBy = p.Subject + } + created, err := h.store.CreateWorkerJoinToken(ctx, store.WorkerJoinToken{ + ID: uuid.NewString(), + TokenHash: hash, + Prefix: prefix, + Name: req.Name, + ExpiresAt: now.Add(h.joinTokenTTL), + CreatedBy: createdBy, + CreatedAt: now, + }) + if err != nil { + h.logger.ErrorContext(ctx, "workerenroll: create join token failed", slog.Any("error", err)) + writeProblem(w, r, http.StatusInternalServerError, "failed to create join token") + return + } + + writeJSON(w, http.StatusCreated, workerJoinTokenCreatedResponse{ + ID: created.ID, + Token: rawToken, + Prefix: created.Prefix, + Name: created.Name, + ExpiresAt: created.ExpiresAt, + CreatedAt: created.CreatedAt, + }) +} + +// ── DELETE /api/v1/workers/{id}/credential ────────────────────────────────── + +// revokeCredential soft-revokes the active credential for the worker named by +// {id} and, via the injected [WorkerRevoker], disconnects it from the broker +// where a live handle is available. The underlying store write only matches +// a row with a nil RevokedAt, so a worker that was never enrolled and a +// worker whose credential is already revoked are indistinguishable here — +// the response says so rather than claiming the worker does not exist. +func (h *workerEnrollHandler) revokeCredential(w http.ResponseWriter, r *http.Request) { + ctx := r.Context() + workerID := chi.URLParam(r, "id") + if err := h.revoker.RevokeWorker(ctx, workerID); err != nil { + if errors.Is(err, store.ErrNotFound) { + writeProblem(w, r, http.StatusNotFound, fmt.Sprintf( + "no active credential for worker %q — it may never have been enrolled, or its credential may already be revoked", + workerID, + )) + return + } + h.logger.ErrorContext(ctx, "workerenroll: revoke credential failed", slog.Any("error", err)) + writeProblem(w, r, http.StatusInternalServerError, "failed to revoke worker credential") + return + } + w.WriteHeader(http.StatusNoContent) +} diff --git a/internal/api/workerenroll_test.go b/internal/api/workerenroll_test.go new file mode 100644 index 00000000..2699e3b7 --- /dev/null +++ b/internal/api/workerenroll_test.go @@ -0,0 +1,902 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package api + +// Unit tests for the worker broker-credential REST handlers. +// +// Route coverage: +// +// POST /api/v1/workers/enroll — enroll +// POST /api/v1/workers/join-tokens — createJoinToken +// DELETE /api/v1/workers/{id}/credential — revokeCredential + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/go-chi/chi/v5" + "github.com/google/uuid" + + "github.com/uberware/sqi/internal/auth" + "github.com/uberware/sqi/internal/auth/jointoken" + "github.com/uberware/sqi/internal/brokerauth" + "github.com/uberware/sqi/internal/health" + "github.com/uberware/sqi/internal/metrics" + "github.com/uberware/sqi/internal/store" + "github.com/uberware/sqi/internal/store/fake" +) + +// ── router helper ──────────────────────────────────────────────────────────── + +// storeRevoker adapts a store.Store directly to [WorkerRevoker] by writing +// the credential row and nothing else — no broker reload. The real +// broker-reload behavior belongs to *server.Server (internal/server cannot +// be imported here: internal/server imports internal/api, not the other way +// around) and is covered by test/integration's broker-auth suite. Every test +// in this file that does not care about broker semantics gets this by +// default via newWorkerEnrollRouter. +type storeRevoker struct{ store store.WorkerCredentialStore } + +func (r storeRevoker) RevokeWorker(ctx context.Context, workerID string) error { + return r.store.RevokeWorkerCredential(ctx, workerID, time.Now().UTC()) +} + +// noopReloader is a [BrokerCredentialReloader] that does nothing and +// returns nil — the enroll-side default for tests that don't care about +// broker-reload semantics, matching storeRevoker's role on the revoke side. +type noopReloader struct{} + +func (noopReloader) ReloadBrokerCredentials(context.Context) error { return nil } + +// recordingReloader is a [BrokerCredentialReloader] stub that counts calls +// and returns a configurable error, for tests verifying enroll's delegation +// to the reloader and its log-and-continue failure handling. +type recordingReloader struct { + calls int + err error +} + +func (r *recordingReloader) ReloadBrokerCredentials(context.Context) error { + r.calls++ + return r.err +} + +func newWorkerEnrollRouter(st store.Store, singleUse bool, ttl time.Duration) chi.Router { + return newWorkerEnrollRouterWith(st, storeRevoker{store: st}, noopReloader{}, singleUse, ttl) +} + +func newWorkerEnrollRouterWithRevoker(st store.Store, revoker WorkerRevoker, singleUse bool, ttl time.Duration) chi.Router { + return newWorkerEnrollRouterWith(st, revoker, noopReloader{}, singleUse, ttl) +} + +func newWorkerEnrollRouterWith(st store.Store, revoker WorkerRevoker, reloader BrokerCredentialReloader, singleUse bool, ttl time.Duration) chi.Router { + h := newWorkerEnrollHandler(st, revoker, reloader, newTestLogger(), singleUse, ttl) + r := chi.NewRouter() + r.Post("/workers/enroll", h.enroll) + r.Post("/workers/join-tokens", h.createJoinToken) + r.Delete("/workers/{id}/credential", h.revokeCredential) + return r +} + +// ── seed helpers ───────────────────────────────────────────────────────────── + +// seedJoinToken creates a join token directly in the store (bypassing the +// mint handler) and returns the raw token alongside the stored record. +// mutate, when non-nil, runs after the default fields are set so a test can +// override ExpiresAt or UsedAt to exercise the reject paths. +func seedJoinToken(t *testing.T, st store.Store, mutate func(*store.WorkerJoinToken)) (raw string, tok store.WorkerJoinToken) { + t.Helper() + raw, hash, prefix, err := jointoken.Generate() + if err != nil { + t.Fatalf("jointoken.Generate: %v", err) + } + rec := store.WorkerJoinToken{ + ID: uuid.NewString(), + TokenHash: hash, + Prefix: prefix, + ExpiresAt: time.Now().UTC().Add(time.Hour), + CreatedAt: time.Now().UTC(), + } + if mutate != nil { + mutate(&rec) + } + created, err := st.CreateWorkerJoinToken(t.Context(), rec) + if err != nil { + t.Fatalf("CreateWorkerJoinToken: %v", err) + } + return raw, created +} + +// genPublicKey returns a fresh, valid nkey user public key. +func genPublicKey(t *testing.T) string { + t.Helper() + _, pub, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed: %v", err) + } + return pub +} + +// ── POST /workers/enroll ───────────────────────────────────────────────────── + +func TestWorkerEnroll_ValidTokenAndKey_Created(t *testing.T) { + st := fake.New() + r := newWorkerEnrollRouter(st, true, time.Hour) + raw, _ := seedJoinToken(t, st, nil) + pub := genPublicKey(t) + + req := newReq(t, http.MethodPost, "/workers/enroll", jsonBody(t, workerEnrollRequest{ + JoinToken: raw, WorkerID: "w1", PublicKey: pub, Name: "render-01", + })) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + + if rr.Code != http.StatusCreated { + t.Fatalf("status = %d, want 201 — body: %s", rr.Code, rr.Body) + } + var resp workerCredentialResponse + if err := json.Unmarshal(rr.Body.Bytes(), &resp); err != nil { + t.Fatalf("decode response: %v", err) + } + if resp.WorkerID != "w1" || resp.PublicKey != pub || resp.Name != "render-01" { + t.Errorf("response = %+v, want worker_id=w1 public_key=%s name=render-01", resp, pub) + } + + cred, err := st.GetActiveWorkerCredentialByWorkerID(t.Context(), "w1") + if err != nil { + t.Fatalf("GetActiveWorkerCredentialByWorkerID: %v", err) + } + if cred.PublicKey != pub { + t.Errorf("stored credential public key = %q, want %q", cred.PublicKey, pub) + } + + // The token must never be echoed back in the response. + if strings.Contains(rr.Body.String(), raw) { + t.Error("response body echoes the raw join token") + } +} + +// TestWorkerEnroll_ReloadsBrokerCredentialsAfterSuccess proves enroll +// delegates to the injected [BrokerCredentialReloader] — not just the store +// — after a successful credential creation. This is the property that makes +// a freshly-enrolled worker able to connect to a RUNNING broker without a +// restart; test/integration's broker-auth suite proves the real +// *server.Server implementation end to end against a live broker. +func TestWorkerEnroll_ReloadsBrokerCredentialsAfterSuccess(t *testing.T) { + st := fake.New() + reloader := &recordingReloader{} + r := newWorkerEnrollRouterWith(st, storeRevoker{store: st}, reloader, true, time.Hour) + raw, _ := seedJoinToken(t, st, nil) + + req := newReq(t, http.MethodPost, "/workers/enroll", jsonBody(t, workerEnrollRequest{ + JoinToken: raw, WorkerID: "w1", PublicKey: genPublicKey(t), + })) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + + if rr.Code != http.StatusCreated { + t.Fatalf("status = %d, want 201 — body: %s", rr.Code, rr.Body) + } + if reloader.calls != 1 { + t.Errorf("reloader.calls = %d, want 1 — enroll must reload the broker's authorized-key set after creating the credential", reloader.calls) + } +} + +// TestWorkerEnroll_ReloadFailure_StillCreated pins the deliberate asymmetry +// with revoke: the credential is genuinely created and durable regardless of +// whether the broker reload succeeds, so a reload failure here is logged and +// swallowed, not turned into an error response — telling the caller +// enrollment failed would be false, and the worker can simply retry +// connecting once a later reload or restart picks up the row that already +// exists in the store. +func TestWorkerEnroll_ReloadFailure_StillCreated(t *testing.T) { + st := fake.New() + reloader := &recordingReloader{err: errors.New("broker not started")} + r := newWorkerEnrollRouterWith(st, storeRevoker{store: st}, reloader, true, time.Hour) + raw, _ := seedJoinToken(t, st, nil) + + req := newReq(t, http.MethodPost, "/workers/enroll", jsonBody(t, workerEnrollRequest{ + JoinToken: raw, WorkerID: "w1", PublicKey: genPublicKey(t), + })) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + + if rr.Code != http.StatusCreated { + t.Fatalf("status = %d, want 201 even though the reload failed — body: %s", rr.Code, rr.Body) + } + if reloader.calls != 1 { + t.Errorf("reloader.calls = %d, want 1", reloader.calls) + } + if _, err := st.GetActiveWorkerCredentialByWorkerID(t.Context(), "w1"); err != nil { + t.Errorf("credential was not durably created despite the reload failure: %v", err) + } + if strings.Contains(rr.Body.String(), "broker not started") { + t.Error("response leaks the underlying reload error; it must not appear given the 201 above") + } +} + +func TestWorkerEnroll_UnknownToken_Unauthorized(t *testing.T) { + st := fake.New() + r := newWorkerEnrollRouter(st, true, time.Hour) + + req := newReq(t, http.MethodPost, "/workers/enroll", jsonBody(t, workerEnrollRequest{ + JoinToken: "sqiw_does-not-exist", WorkerID: "w1", PublicKey: genPublicKey(t), + })) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + + if rr.Code != http.StatusUnauthorized { + t.Fatalf("status = %d, want 401 — body: %s", rr.Code, rr.Body) + } +} + +func TestWorkerEnroll_ExpiredToken_Unauthorized(t *testing.T) { + st := fake.New() + r := newWorkerEnrollRouter(st, true, time.Hour) + raw, _ := seedJoinToken(t, st, func(tok *store.WorkerJoinToken) { + tok.ExpiresAt = time.Now().UTC().Add(-time.Minute) + }) + + req := newReq(t, http.MethodPost, "/workers/enroll", jsonBody(t, workerEnrollRequest{ + JoinToken: raw, WorkerID: "w1", PublicKey: genPublicKey(t), + })) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + + if rr.Code != http.StatusUnauthorized { + t.Fatalf("status = %d, want 401 — body: %s", rr.Code, rr.Body) + } +} + +func TestWorkerEnroll_UsedSingleUseToken_Unauthorized(t *testing.T) { + st := fake.New() + r := newWorkerEnrollRouter(st, true, time.Hour) + usedAt := time.Now().UTC().Add(-time.Minute) + raw, _ := seedJoinToken(t, st, func(tok *store.WorkerJoinToken) { + tok.UsedAt = &usedAt + }) + + req := newReq(t, http.MethodPost, "/workers/enroll", jsonBody(t, workerEnrollRequest{ + JoinToken: raw, WorkerID: "w1", PublicKey: genPublicKey(t), + })) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + + if rr.Code != http.StatusUnauthorized { + t.Fatalf("status = %d, want 401 — body: %s", rr.Code, rr.Body) + } +} + +func TestWorkerEnroll_UsedTokenAllowedWhenNotSingleUse(t *testing.T) { + st := fake.New() + r := newWorkerEnrollRouter(st, false, time.Hour) // single-use OFF + usedAt := time.Now().UTC().Add(-time.Minute) + raw, _ := seedJoinToken(t, st, func(tok *store.WorkerJoinToken) { + tok.UsedAt = &usedAt + }) + + req := newReq(t, http.MethodPost, "/workers/enroll", jsonBody(t, workerEnrollRequest{ + JoinToken: raw, WorkerID: "w1", PublicKey: genPublicKey(t), + })) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + + if rr.Code != http.StatusCreated { + t.Fatalf("status = %d, want 201 (single-use disabled, a used token is still accepted) — body: %s", rr.Code, rr.Body) + } +} + +func TestWorkerEnroll_WorkerIDAlreadyBoundToDifferentKey_Conflict(t *testing.T) { + st := fake.New() + r := newWorkerEnrollRouter(st, true, time.Hour) + + raw1, _ := seedJoinToken(t, st, nil) + req1 := newReq(t, http.MethodPost, "/workers/enroll", jsonBody(t, workerEnrollRequest{ + JoinToken: raw1, WorkerID: "w1", PublicKey: genPublicKey(t), + })) + rr1 := httptest.NewRecorder() + r.ServeHTTP(rr1, req1) + if rr1.Code != http.StatusCreated { + t.Fatalf("first enroll: status = %d, want 201 — body: %s", rr1.Code, rr1.Body) + } + + raw2, _ := seedJoinToken(t, st, nil) + req2 := newReq(t, http.MethodPost, "/workers/enroll", jsonBody(t, workerEnrollRequest{ + JoinToken: raw2, WorkerID: "w1", PublicKey: genPublicKey(t), // same worker, different key + })) + rr2 := httptest.NewRecorder() + r.ServeHTTP(rr2, req2) + + if rr2.Code != http.StatusConflict { + t.Fatalf("second enroll: status = %d, want 409 — body: %s", rr2.Code, rr2.Body) + } +} + +func TestWorkerEnroll_MalformedPublicKey_BadRequest(t *testing.T) { + st := fake.New() + r := newWorkerEnrollRouter(st, true, time.Hour) + raw, _ := seedJoinToken(t, st, nil) + + req := newReq(t, http.MethodPost, "/workers/enroll", jsonBody(t, workerEnrollRequest{ + JoinToken: raw, WorkerID: "w1", PublicKey: "not-a-valid-key", + })) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + + if rr.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want 400 — body: %s", rr.Code, rr.Body) + } +} + +func TestWorkerEnroll_MalformedJSON_BadRequest(t *testing.T) { + st := fake.New() + r := newWorkerEnrollRouter(st, true, time.Hour) + + req := newReq(t, http.MethodPost, "/workers/enroll", badJSON()) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + + if rr.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want 400 — body: %s", rr.Code, rr.Body) + } +} + +func TestWorkerEnroll_MissingFields_BadRequest(t *testing.T) { + st := fake.New() + r := newWorkerEnrollRouter(st, true, time.Hour) + + req := newReq(t, http.MethodPost, "/workers/enroll", jsonBody(t, workerEnrollRequest{WorkerID: "w1"})) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + + if rr.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want 400 — body: %s", rr.Code, rr.Body) + } +} + +// ── the non-negotiable: unknown/expired/used/store-failure are indistinguishable ── +// +// enroll is unauthenticated and may be internet-reachable, so every way a +// join token can fail to authorize a call must answer with the identical +// status AND body — never just the same status. TestWorkerEnroll_UnknownToken_ +// Unauthorized and friends above already prove each individual case returns +// 401; this test is what actually pins the property those don't: that +// nothing distinguishes them from each other. A later change that makes any +// one path "more helpful" (e.g. naming which case it was) would pass every +// existing per-case test and only fail here. + +// joinTokenLookupErrStore forces the join-token claim to fail with a +// non-ErrNotFound error, simulating a store outage while the token is being +// redeemed — the fourth way enroll can fail to authorize a request, +// alongside unknown, expired, and already-used. Both the single-use claim +// and the reusable-token lookup are overridden so the case holds whichever +// path a test exercises. +type joinTokenLookupErrStore struct { + store.Store +} + +func (joinTokenLookupErrStore) GetWorkerJoinTokenByHash(context.Context, string) (store.WorkerJoinToken, error) { + return store.WorkerJoinToken{}, errors.New("simulated store outage") +} + +func (joinTokenLookupErrStore) RedeemWorkerJoinToken(context.Context, string, time.Time, store.WorkerCredential) (store.WorkerCredential, error) { + return store.WorkerCredential{}, errors.New("simulated store outage") +} + +func TestWorkerEnroll_TokenFailureModesAreIndistinguishable(t *testing.T) { + pub := genPublicKey(t) + usedAt := time.Now().UTC().Add(-time.Minute) + + // dispatch builds a fresh router over st (singleUse always on, so the + // "already-used" case actually rejects) and returns the recorded + // response to one enroll attempt with joinToken. + dispatch := func(st store.Store, joinToken string) *httptest.ResponseRecorder { + r := newWorkerEnrollRouter(st, true, time.Hour) + req := newReq(t, http.MethodPost, "/workers/enroll", jsonBody(t, workerEnrollRequest{ + JoinToken: joinToken, WorkerID: "w1", PublicKey: pub, + })) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + return rr + } + + cases := []struct { + name string + run func() *httptest.ResponseRecorder + }{ + { + name: "unknown token", + run: func() *httptest.ResponseRecorder { + return dispatch(fake.New(), "sqiw_does-not-exist") + }, + }, + { + name: "expired token", + run: func() *httptest.ResponseRecorder { + st := fake.New() + raw, _ := seedJoinToken(t, st, func(tok *store.WorkerJoinToken) { + tok.ExpiresAt = time.Now().UTC().Add(-time.Minute) + }) + return dispatch(st, raw) + }, + }, + { + name: "already-used single-use token", + run: func() *httptest.ResponseRecorder { + st := fake.New() + raw, _ := seedJoinToken(t, st, func(tok *store.WorkerJoinToken) { + tok.UsedAt = &usedAt + }) + return dispatch(st, raw) + }, + }, + { + name: "store failure during lookup", + run: func() *httptest.ResponseRecorder { + return dispatch(joinTokenLookupErrStore{Store: fake.New()}, "sqiw_irrelevant-lookup-always-fails") + }, + }, + } + + type outcome struct { + name string + code int + body string + } + var got []outcome + for _, tc := range cases { + rr := tc.run() + if rr.Code != http.StatusUnauthorized { + t.Errorf("%s: status = %d, want 401", tc.name, rr.Code) + } + got = append(got, outcome{name: tc.name, code: rr.Code, body: rr.Body.String()}) + } + + want := got[0] + for _, g := range got[1:] { + if g.code != want.code { + t.Errorf("%s: status = %d, want %d (same as %q) — a caller could distinguish "+ + "join-token failure modes by status code, which the endpoint must never allow", + g.name, g.code, want.code, want.name) + } + if g.body != want.body { + t.Errorf("%s: body = %q, want %q (same as %q) — a caller could distinguish "+ + "join-token failure modes by response body, which the endpoint must never allow", + g.name, g.body, want.body, want.name) + } + } +} + +// ── router-level mounting gate ─────────────────────────────────────────────── +// +// These use chi.Walk (via the liveRoutes/routeKey helpers from +// authz_integration_test.go) rather than firing an HTTP request: POST +// /workers/enroll and GET/DELETE /workers/{id} share the same two-segment +// shape, so when enroll is NOT mounted, a request for it is routed as +// /workers/{id} with id="enroll" — chi answers that with 405 (the path +// matches a registered pattern, just not for POST), not 404. Walking the +// route table sidesteps that collision and asserts what actually matters: +// whether the route is registered at all. + +func TestWorkerEnroll_EndpointAbsentWhenEnrollmentEndpointDisabled(t *testing.T) { + deps := Deps{ + Store: fake.New(), + Auth: auth.Anonymous(), + NATSAuthEnabled: true, + // NATSAuthEnrollmentEndpointEnabled left false. + } + r := NewRouter(Config{DisableRateLimit: true}, deps, newTestLogger(), metrics.New(), health.NewRegistry()) + if liveRoutes(t, r)[routeKey{http.MethodPost, "/api/v1/workers/enroll"}] { + t.Error("POST /api/v1/workers/enroll is registered even though the enrollment endpoint is disabled") + } +} + +func TestWorkerEnroll_EndpointAbsentWhenNATSAuthDisabled(t *testing.T) { + deps := Deps{ + Store: fake.New(), + Auth: auth.Anonymous(), + NATSAuthEnabled: false, + NATSAuthEnrollmentEndpointEnabled: true, + } + r := NewRouter(Config{DisableRateLimit: true}, deps, newTestLogger(), metrics.New(), health.NewRegistry()) + if liveRoutes(t, r)[routeKey{http.MethodPost, "/api/v1/workers/enroll"}] { + t.Error("POST /api/v1/workers/enroll is registered even though broker (nats.auth) authentication is disabled") + } +} + +func TestWorkerEnroll_EndpointMountedWhenBothEnabled(t *testing.T) { + deps := Deps{ + Store: fake.New(), + Auth: auth.Anonymous(), + NATSAuthEnabled: true, + NATSAuthEnrollmentEndpointEnabled: true, + } + r := NewRouter(Config{DisableRateLimit: true}, deps, newTestLogger(), metrics.New(), health.NewRegistry()) + if !liveRoutes(t, r)[routeKey{http.MethodPost, "/api/v1/workers/enroll"}] { + t.Error("POST /api/v1/workers/enroll is not registered even though both gating flags are true") + } +} + +// ── POST /workers/join-tokens ──────────────────────────────────────────────── + +func TestWorkerJoinTokenCreate_Created(t *testing.T) { + st := fake.New() + r := newWorkerEnrollRouter(st, true, time.Hour) + + req := newReq(t, http.MethodPost, "/workers/join-tokens", jsonBody(t, workerJoinTokenCreateRequest{Name: "batch-1"})) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + + if rr.Code != http.StatusCreated { + t.Fatalf("status = %d, want 201 — body: %s", rr.Code, rr.Body) + } + var resp workerJoinTokenCreatedResponse + if err := json.Unmarshal(rr.Body.Bytes(), &resp); err != nil { + t.Fatalf("decode response: %v", err) + } + if resp.Token == "" { + t.Fatal("response did not carry the raw token") + } + if resp.Name != "batch-1" { + t.Errorf("name = %q, want batch-1", resp.Name) + } + + stored, err := st.GetWorkerJoinTokenByHash(t.Context(), jointoken.Hash(resp.Token)) + if err != nil { + t.Fatalf("GetWorkerJoinTokenByHash: %v", err) + } + if stored.ID != resp.ID { + t.Errorf("stored token id = %q, want %q", stored.ID, resp.ID) + } + if stored.TokenHash == resp.Token { + t.Error("stored TokenHash equals the raw token — the raw value must never be persisted") + } +} + +func TestWorkerJoinTokenCreate_EmptyBody_Created(t *testing.T) { + st := fake.New() + r := newWorkerEnrollRouter(st, true, time.Hour) + + req := newReq(t, http.MethodPost, "/workers/join-tokens", nil) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + + if rr.Code != http.StatusCreated { + t.Fatalf("status = %d, want 201 — body: %s", rr.Code, rr.Body) + } +} + +// ── DELETE /workers/{id}/credential ───────────────────────────────────────── + +func TestWorkerCredentialRevoke_EnrolledWorker_NoContent(t *testing.T) { + st := fake.New() + r := newWorkerEnrollRouter(st, true, time.Hour) + if _, err := st.CreateWorkerCredential(t.Context(), store.WorkerCredential{ + ID: uuid.NewString(), WorkerID: "w1", PublicKey: genPublicKey(t), EnrolledAt: time.Now().UTC(), + }); err != nil { + t.Fatalf("seed CreateWorkerCredential: %v", err) + } + + req := newReq(t, http.MethodDelete, "/workers/w1/credential", nil) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + + if rr.Code != http.StatusNoContent { + t.Fatalf("status = %d, want 204 — body: %s", rr.Code, rr.Body) + } + + if _, err := st.GetActiveWorkerCredentialByWorkerID(t.Context(), "w1"); err == nil { + t.Error("credential still active after revoke") + } else if !errors.Is(err, store.ErrNotFound) { + t.Errorf("GetActiveWorkerCredentialByWorkerID after revoke: %v, want store.ErrNotFound", err) + } +} + +func TestWorkerCredentialRevoke_UnknownWorker_NotFound(t *testing.T) { + st := fake.New() + r := newWorkerEnrollRouter(st, true, time.Hour) + + req := newReq(t, http.MethodDelete, "/workers/does-not-exist/credential", nil) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + + if rr.Code != http.StatusNotFound { + t.Fatalf("status = %d, want 404 — body: %s", rr.Code, rr.Body) + } + // The response must not claim the worker doesn't exist — that case is + // indistinguishable from "already revoked" at the store layer. + if strings.Contains(rr.Body.String(), "does not exist") { + t.Error("response claims the worker does not exist, which the store cannot actually distinguish from already-revoked") + } +} + +// ── revoke delegates to the injected WorkerRevoker ────────────────────────── + +// recordingRevoker is a [WorkerRevoker] stub that records the workerID it +// was called with and returns a configurable error, so a test can prove the +// handler actually calls through the injected interface — not the store +// directly — and that it maps a non-[store.ErrNotFound] failure to 500 +// rather than leaking the underlying error text (which, in production, would +// be a broker-internal detail like a reload failure). +type recordingRevoker struct { + calledWith string + err error +} + +func (r *recordingRevoker) RevokeWorker(_ context.Context, workerID string) error { + r.calledWith = workerID + return r.err +} + +func TestWorkerCredentialRevoke_DelegatesToInjectedRevoker(t *testing.T) { + st := fake.New() + rev := &recordingRevoker{} + r := newWorkerEnrollRouterWithRevoker(st, rev, true, time.Hour) + + req := newReq(t, http.MethodDelete, "/workers/w1/credential", nil) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + + if rr.Code != http.StatusNoContent { + t.Fatalf("status = %d, want 204 — body: %s", rr.Code, rr.Body) + } + if rev.calledWith != "w1" { + t.Errorf("revoker called with %q, want %q — the handler must delegate to the injected WorkerRevoker, not write the store directly", + rev.calledWith, "w1") + } +} + +func TestWorkerCredentialRevoke_RevokerErrorNotFound_NotFound(t *testing.T) { + st := fake.New() + rev := &recordingRevoker{err: store.ErrNotFound} + r := newWorkerEnrollRouterWithRevoker(st, rev, true, time.Hour) + + req := newReq(t, http.MethodDelete, "/workers/does-not-exist/credential", nil) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + + if rr.Code != http.StatusNotFound { + t.Fatalf("status = %d, want 404 — body: %s", rr.Code, rr.Body) + } +} + +func TestWorkerCredentialRevoke_RevokerOtherError_InternalServerError(t *testing.T) { + st := fake.New() + rev := &recordingRevoker{err: errors.New("broker reload failed")} + r := newWorkerEnrollRouterWithRevoker(st, rev, true, time.Hour) + + req := newReq(t, http.MethodDelete, "/workers/w1/credential", nil) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + + if rr.Code != http.StatusInternalServerError { + t.Fatalf("status = %d, want 500 — body: %s", rr.Code, rr.Body) + } + if strings.Contains(rr.Body.String(), "broker reload failed") { + t.Error("response leaks the underlying revoker error; expected the generic message") + } +} + +// ── single-use redemption is atomic ───────────────────────────────────────── + +// TestWorkerEnroll_ConcurrentRedemptionsOfOneSingleUseToken proves the +// redemption is a claim, not a check followed by a claim. Reading the token, +// inspecting UsedAt and marking it used as separate steps lets two +// simultaneous enrollments with one single-use token BOTH observe UsedAt as +// nil and both succeed — the token's whole purpose defeated by timing +// alone. store.RedeemWorkerJoinToken makes check and claim one statement +// (inside one transaction with the credential creation), so exactly one of +// these can win. +func TestWorkerEnroll_ConcurrentRedemptionsOfOneSingleUseToken(t *testing.T) { + // Repeated rounds, not one: the handler is fast enough that a single + // burst can happen to serialize even under -race, so one round proves + // nothing about a check-then-act implementation. Every round must yield + // exactly one 201. + const rounds = 50 + const attempts = 8 + + for round := range rounds { + st := fake.New() + r := newWorkerEnrollRouter(st, true, time.Hour) + raw, _ := seedJoinToken(t, st, nil) + + // Build every request up front so the goroutines do nothing but + // serve once released. + reqs := make([]*http.Request, attempts) + for i := range reqs { + reqs[i] = newReq(t, http.MethodPost, "/workers/enroll", jsonBody(t, workerEnrollRequest{ + JoinToken: raw, + WorkerID: fmt.Sprintf("w%d", i), + PublicKey: genPublicKey(t), + })) + } + + codes := make([]int, attempts) + var ready, done sync.WaitGroup + ready.Add(attempts) + done.Add(attempts) + start := make(chan struct{}) + for i := range attempts { + go func() { + defer done.Done() + rr := httptest.NewRecorder() + ready.Done() + <-start + r.ServeHTTP(rr, reqs[i]) + codes[i] = rr.Code + }() + } + ready.Wait() + close(start) + done.Wait() + + created := 0 + for i, code := range codes { + switch code { + case http.StatusCreated: + created++ + case http.StatusUnauthorized: + default: + t.Fatalf("round %d attempt %d: status = %d, want 201 or 401", round, i, code) + } + } + if created != 1 { + t.Fatalf("round %d: %d of %d concurrent enrollments succeeded with one single-use token, want exactly 1", + round, created, attempts) + } + } +} + +// TestWorkerEnroll_MalformedPublicKey_TokenNotSpent pins that a request +// rejected before the token is ever claimed costs nothing: a malformed +// public key is caught by validation, ahead of any store call, so the token +// stays exactly as redeemable as it was. +func TestWorkerEnroll_MalformedPublicKey_TokenNotSpent(t *testing.T) { + st := fake.New() + r := newWorkerEnrollRouter(st, true, time.Hour) + + rawBadKey, badKeyTok := seedJoinToken(t, st, nil) + req := newReq(t, http.MethodPost, "/workers/enroll", jsonBody(t, workerEnrollRequest{ + JoinToken: rawBadKey, WorkerID: "w2", PublicKey: "not-a-valid-key", + })) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + if rr.Code != http.StatusBadRequest { + t.Fatalf("malformed key: status = %d, want 400 — body: %s", rr.Code, rr.Body) + } + stored, err := st.GetWorkerJoinTokenByHash(t.Context(), badKeyTok.TokenHash) + if err != nil { + t.Fatalf("GetWorkerJoinTokenByHash: %v", err) + } + if stored.UsedAt != nil { + t.Error("a malformed public key spent the join token; it must be rejected before the token is claimed") + } +} + +// TestWorkerEnroll_ConflictingEnrollment_TokenSurvives proves that a +// single-use token presented against a request that conflicts (here, a +// worker ID that already has an active credential) must NOT be burned. +// The claim and the credential creation are one transaction +// (store.RedeemWorkerJoinToken), so a conflict rolls the claim back along +// with the credential write — the opposite of the token being spent +// unconditionally ahead of the write, which this test would have caught as +// a regression back to that shape. It proves survival the only way that +// actually matters to an operator: the SAME token is still redeemable +// afterwards, for a different worker ID. +func TestWorkerEnroll_ConflictingEnrollment_TokenSurvives(t *testing.T) { + st := fake.New() + r := newWorkerEnrollRouter(st, true, time.Hour) + raw, _ := seedJoinToken(t, st, nil) + + // First enrollment for w1 succeeds and consumes nothing it shouldn't — + // it is the LEGITIMATE use of this token. + req := newReq(t, http.MethodPost, "/workers/enroll", jsonBody(t, workerEnrollRequest{ + JoinToken: raw, WorkerID: "w1", PublicKey: genPublicKey(t), + })) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + if rr.Code != http.StatusCreated { + t.Fatalf("first enrollment: status = %d, want 201 — body: %s", rr.Code, rr.Body) + } + + // A single-use token is spent by its first successful redemption, so a + // second attempt with the SAME token — even one that would also + // conflict on worker ID w1 — reports the generic invalid-token 401 + // rather than a 409: the token itself is what "already used" refuses. + // The point under test is what happens when the credential write fails + // for a reason OTHER than the token already being spent — a worker ID + // bound to an active credential from a DIFFERENT token's redemption — + // so seed a second token for that attempt. + raw2, _ := seedJoinToken(t, st, nil) + req = newReq(t, http.MethodPost, "/workers/enroll", jsonBody(t, workerEnrollRequest{ + JoinToken: raw2, WorkerID: "w1", PublicKey: genPublicKey(t), + })) + rr = httptest.NewRecorder() + r.ServeHTTP(rr, req) + if rr.Code != http.StatusConflict { + t.Fatalf("conflicting worker id: status = %d, want 409 — body: %s", rr.Code, rr.Body) + } + + // raw2 must still be redeemable: the conflict rolled its claim back. + // Confirm by actually redeeming it, for a DIFFERENT worker ID. + req = newReq(t, http.MethodPost, "/workers/enroll", jsonBody(t, workerEnrollRequest{ + JoinToken: raw2, WorkerID: "w3", PublicKey: genPublicKey(t), + })) + rr = httptest.NewRecorder() + r.ServeHTTP(rr, req) + if rr.Code != http.StatusCreated { + t.Fatalf("re-redemption after a conflict: status = %d, want 201 — body: %s", rr.Code, rr.Body) + } +} + +// ── worker_id must be a single NATS subject token ─────────────────────────── + +// TestWorkerEnroll_InvalidWorkerID_BadRequest is the enrollment boundary's +// half of the branch's whole premise. The stored worker_id flows verbatim +// into brokerauth.WorkerPermissions when the broker builds its key set, and +// those grants are NATS subject PATTERNS. A worker_id of "*" therefore mints +// a credential granted "task.status.*.*", "worker.deregister.*", +// "work.lease.*.*" and the rest — one credential that may publish concrete +// subjects belonging to ANY worker: forge status and logs, deregister the +// farm, and lease work as another worker, receiving that worker's assignment +// batch in its own inbox. The scheduler's provenance checks cannot catch it, +// because the subject NATS vouches for genuinely names the victim. +// +// ">" is worse-shaped still: it yields the malformed "task.status.>.*" +// inside Options.Nkeys, which can make every later ReloadOptions fail — +// revocation permanently 500ing — or wedge the broker at boot. +func TestWorkerEnroll_InvalidWorkerID_BadRequest(t *testing.T) { + cases := []struct { + name string + workerID string + }{ + {"single-token wildcard", "*"}, + {"multi-token wildcard", ">"}, + {"contains a dot", "render.01"}, + {"contains whitespace", "render 01"}, + {"empty", ""}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + st := fake.New() + r := newWorkerEnrollRouter(st, true, time.Hour) + raw, tok := seedJoinToken(t, st, nil) + + req := newReq(t, http.MethodPost, "/workers/enroll", jsonBody(t, workerEnrollRequest{ + JoinToken: raw, WorkerID: tc.workerID, PublicKey: genPublicKey(t), + })) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + + if rr.Code != http.StatusBadRequest { + t.Fatalf("worker_id %q: status = %d, want 400 — body: %s", tc.workerID, rr.Code, rr.Body) + } + + // Rejected before the claim, so the operator's token survives. + stored, err := st.GetWorkerJoinTokenByHash(t.Context(), tok.TokenHash) + if err != nil { + t.Fatalf("GetWorkerJoinTokenByHash: %v", err) + } + if stored.UsedAt != nil { + t.Error("an invalid worker_id spent the join token; it must be rejected before the token is claimed") + } + + // And nothing was enrolled under it. + creds, err := st.ListActiveWorkerCredentials(t.Context()) + if err != nil { + t.Fatalf("ListActiveWorkerCredentials: %v", err) + } + if len(creds) != 0 { + t.Errorf("credential created for invalid worker_id %q: %+v", tc.workerID, creds) + } + }) + } +} diff --git a/internal/api/workers.go b/internal/api/workers.go index 83e6211c..dd26401b 100644 --- a/internal/api/workers.go +++ b/internal/api/workers.go @@ -30,6 +30,12 @@ type workerHandler struct { // notifier pushes a removed event when a worker is deleted so other // connected clients drop it live. May be nil (no push). notifier ws.Notifier + // revoker revokes a deleted worker's broker credential, through the same + // serialized store-write-then-broker-reload path DELETE + // /api/v1/workers/{id}/credential uses. internal/api depends only on + // this interface — never on internal/bus or internal/server — the same + // seam workerEnrollHandler uses for the same reason. + revoker WorkerRevoker // offlineThreshold is the heartbeat-timeout window used to decide whether a // disabled worker is dead — and therefore removable — from its last // heartbeat age. It mirrors the scheduler's WorkerTimeout. @@ -39,10 +45,11 @@ type workerHandler struct { // newWorkerHandler returns a workerHandler wired to the given store. notifier // may be nil in tests that do not exercise WebSocket push. -func newWorkerHandler(st store.Store, notifier ws.Notifier, offlineThreshold time.Duration, logger *slog.Logger) *workerHandler { +func newWorkerHandler(st store.Store, notifier ws.Notifier, revoker WorkerRevoker, offlineThreshold time.Duration, logger *slog.Logger) *workerHandler { return &workerHandler{ store: st, notifier: notifier, + revoker: revoker, offlineThreshold: offlineThreshold, logger: logger, } @@ -348,6 +355,31 @@ func (h *workerHandler) setWorkerStatus(w http.ResponseWriter, r *http.Request, // than the offline threshold (the machine is gone). Online and live-disabled // workers return 409 Conflict. Offline workers already had their in-flight tasks // reclaimed when they went offline, so no reclaim is needed here. +// +// Revokes the worker's broker credential, if it has one, through the +// injected [WorkerRevoker] — the same path DELETE +// /api/v1/workers/{id}/credential uses — BEFORE deleting the worker row. +// Without this, decommissioning a machine from the farm would leave it able +// to connect to the broker, lease work and execute job code: WorkersManage +// (which this route requires) does not imply WorkersEnroll, so an operator +// who can delete a worker is not otherwise able to revoke what it can still +// do. +// +// Revoke-then-delete, not the reverse: store.DeleteWorker never returns +// ErrConflict in either backend (removability was already decided above via +// GetWorker + workerRemovable), so revoking first can never waste a +// revocation on a delete that was going to be legitimately rejected. A +// revoke failure then means nothing happened at all — worker row intact, a +// clean 500, safely retryable. Deleting first would instead let a failure +// of the revoke's own store write (not just a broker-reload failure — a +// documented, recoverable degraded mode) leave the worker row gone, the +// credential never revoked, and nothing left to reap it: the caller who +// holds WorkersManage without WorkersEnroll has no other way to revoke it, +// and would be told 204 while the machine kept live broker access +// permanently. A delete failure after a successful revoke leaves access cut +// and the row present — the safe direction — and a retried removeWorker +// re-invokes RevokeWorker, which is a correct no-op the second time via +// ErrNotFound. func (h *workerHandler) removeWorker(w http.ResponseWriter, r *http.Request) { ctx := r.Context() id := chi.URLParam(r, "id") @@ -370,6 +402,18 @@ func (h *workerHandler) removeWorker(w http.ResponseWriter, r *http.Request) { return } + // store.ErrNotFound covers both "never enrolled" (the normal shape with + // broker auth off) and "already revoked", neither of which is an error + // here. Any other error means the credential is not known to be + // revoked — see the doc comment above for why this must block the + // delete rather than merely being logged. + if err := h.revoker.RevokeWorker(ctx, id); err != nil && !errors.Is(err, store.ErrNotFound) { + h.logger.ErrorContext(ctx, "workers: revoke credential before delete failed", + slog.String("id", id), slog.Any("error", err)) + writeProblem(w, r, http.StatusInternalServerError, "failed to remove worker") + return + } + if err := h.store.DeleteWorker(ctx, id); err != nil { if errors.Is(err, store.ErrNotFound) { writeProblem(w, r, http.StatusNotFound, "worker not found") diff --git a/internal/api/workers_error_test.go b/internal/api/workers_error_test.go index e9ab3082..87ab7175 100644 --- a/internal/api/workers_error_test.go +++ b/internal/api/workers_error_test.go @@ -9,6 +9,7 @@ package api import ( "context" + "errors" "net/http" "net/http/httptest" "testing" @@ -22,14 +23,15 @@ import ( // ── workerErrStore: thin wrapper for worker store errors ────────────────────── -// workerErrStore wraps a store.Store to inject errors into ListWorkers and -// GetWorker. We use a separate type to avoid conflicts with storeErr's method -// set for other store methods. +// workerErrStore wraps a store.Store to inject errors into ListWorkers, +// GetWorker, and DeleteWorker. We use a separate type to avoid conflicts +// with storeErr's method set for other store methods. type workerErrStore struct { store.Store - listWorkersErr error - getWorkerErr error + listWorkersErr error + getWorkerErr error + deleteWorkerErr error } func (e *workerErrStore) ListWorkers(ctx context.Context, opts store.ListWorkersOptions) (store.Page[store.Worker], error) { @@ -46,6 +48,13 @@ func (e *workerErrStore) GetWorker(ctx context.Context, id string) (store.Worker return e.Store.GetWorker(ctx, id) } +func (e *workerErrStore) DeleteWorker(ctx context.Context, id string) error { + if e.deleteWorkerErr != nil { + return e.deleteWorkerErr + } + return e.Store.DeleteWorker(ctx, id) +} + // ── listWorkers: additional filter and error paths ──────────────────────────── func TestListWorkers_QueueIDFilterAndErrors(t *testing.T) { @@ -119,3 +128,41 @@ func TestGetWorker_StoreError(t *testing.T) { } }) } + +// ── removeWorker: a delete failure AFTER a successful revoke ────────────────── + +// TestRemoveWorker_DeleteFailsAfterSuccessfulRevoke proves the safe half of +// removeWorker's revoke-then-delete ordering: when the credential revoke +// succeeds but the subsequent store delete fails, the worker row survives +// (the request answers 500, safe to retry) but its broker access is already +// cut. newWorkerRouter wires the injected WorkerRevoker to the SAME +// underlying fake store as the handler's own store.Store (storeRevoker +// wraps whatever is passed in), so wrapping only DeleteWorker with an error +// here is enough to reach this case: RevokeWorkerCredential runs for real +// against the shared fake, unaffected by the wrapper. +func TestRemoveWorker_DeleteFailsAfterSuccessfulRevoke(t *testing.T) { + inner := fake.New() + w := seedWorker(t, inner, store.WorkerStatusOffline) + if _, err := inner.CreateWorkerCredential(t.Context(), store.WorkerCredential{ + ID: uuid.NewString(), WorkerID: w.ID, PublicKey: genPublicKey(t), EnrolledAt: time.Now().UTC(), + }); err != nil { + t.Fatalf("seed CreateWorkerCredential: %v", err) + } + + est := &workerErrStore{Store: inner, deleteWorkerErr: errInjected} + r := newWorkerRouter(est) + + req := newReq(t, http.MethodDelete, "/api/v1/workers/"+w.ID, nil) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + if rr.Code != http.StatusInternalServerError { + t.Fatalf("expected 500, got %d — body: %s", rr.Code, rr.Body) + } + + if _, err := inner.GetWorker(t.Context(), w.ID); err != nil { + t.Errorf("worker row should survive a failed delete: GetWorker: %v", err) + } + if _, err := inner.GetActiveWorkerCredentialByWorkerID(t.Context(), w.ID); !errors.Is(err, store.ErrNotFound) { + t.Errorf("credential should already be revoked even though the delete failed: got %v, want store.ErrNotFound", err) + } +} diff --git a/internal/api/workers_test.go b/internal/api/workers_test.go index 6c5133c9..1a035d12 100644 --- a/internal/api/workers_test.go +++ b/internal/api/workers_test.go @@ -13,6 +13,7 @@ package api import ( "encoding/json" + "errors" "net/http" "net/http/httptest" "testing" @@ -47,7 +48,15 @@ func newWorkerRouter(st store.Store) chi.Router { } func newWorkerRouterWithNotifier(st store.Store, notifier ws.Notifier) chi.Router { - h := newWorkerHandler(st, notifier, testOfflineThreshold, newTestLogger()) + return newWorkerRouterWith(st, notifier, storeRevoker{store: st}) +} + +func newWorkerRouterWithRevoker(st store.Store, revoker WorkerRevoker) chi.Router { + return newWorkerRouterWith(st, nil, revoker) +} + +func newWorkerRouterWith(st store.Store, notifier ws.Notifier, revoker WorkerRevoker) chi.Router { + h := newWorkerHandler(st, notifier, revoker, testOfflineThreshold, newTestLogger()) r := chi.NewRouter() r.Get("/api/v1/workers", h.listWorkers) r.Get("/api/v1/workers/{id}", h.getWorker) @@ -676,6 +685,86 @@ func TestRemoveWorker(t *testing.T) { t.Fatalf("expected 404, got %d", rr.Code) } }) + + t.Run("active credential is revoked", func(t *testing.T) { + st := fake.New() + r := newWorkerRouter(st) + w := seedWorker(t, st, store.WorkerStatusOffline) + if _, err := st.CreateWorkerCredential(t.Context(), store.WorkerCredential{ + ID: uuid.NewString(), WorkerID: w.ID, PublicKey: genPublicKey(t), EnrolledAt: time.Now().UTC(), + }); err != nil { + t.Fatalf("seed CreateWorkerCredential: %v", err) + } + + req := newReq(t, http.MethodDelete, "/api/v1/workers/"+w.ID, nil) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + if rr.Code != http.StatusNoContent { + t.Fatalf("expected 204, got %d — body: %s", rr.Code, rr.Body) + } + if _, err := st.GetActiveWorkerCredentialByWorkerID(t.Context(), w.ID); !errors.Is(err, store.ErrNotFound) { + t.Errorf("GetActiveWorkerCredentialByWorkerID after delete = %v, want store.ErrNotFound (credential revoked)", err) + } + }) + + t.Run("no credential still succeeds", func(t *testing.T) { + st := fake.New() + rev := &recordingRevoker{err: store.ErrNotFound} + r := newWorkerRouterWithRevoker(st, rev) + w := seedWorker(t, st, store.WorkerStatusOffline) + + req := newReq(t, http.MethodDelete, "/api/v1/workers/"+w.ID, nil) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + if rr.Code != http.StatusNoContent { + t.Fatalf("expected 204, got %d — body: %s", rr.Code, rr.Body) + } + if rev.calledWith != w.ID { + t.Errorf("revoker called with %q, want %q — delete must still call through the revoker even when it has nothing to revoke", rev.calledWith, w.ID) + } + }) + + t.Run("already-revoked credential still succeeds", func(t *testing.T) { + st := fake.New() + r := newWorkerRouter(st) + w := seedWorker(t, st, store.WorkerStatusOffline) + if _, err := st.CreateWorkerCredential(t.Context(), store.WorkerCredential{ + ID: uuid.NewString(), WorkerID: w.ID, PublicKey: genPublicKey(t), EnrolledAt: time.Now().UTC(), + }); err != nil { + t.Fatalf("seed CreateWorkerCredential: %v", err) + } + if err := st.RevokeWorkerCredential(t.Context(), w.ID, time.Now().UTC()); err != nil { + t.Fatalf("seed RevokeWorkerCredential: %v", err) + } + + req := newReq(t, http.MethodDelete, "/api/v1/workers/"+w.ID, nil) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + if rr.Code != http.StatusNoContent { + t.Fatalf("expected 204, got %d — body: %s", rr.Code, rr.Body) + } + }) + + t.Run("revoker failure blocks the delete", func(t *testing.T) { + // The worker row must survive a revoke failure: revoking runs + // BEFORE deleting specifically so that a failure here never leaves + // a deleted worker whose credential nothing revoked and nothing + // will ever reap. + st := fake.New() + rev := &recordingRevoker{err: errors.New("store write failed")} + r := newWorkerRouterWithRevoker(st, rev) + w := seedWorker(t, st, store.WorkerStatusOffline) + + req := newReq(t, http.MethodDelete, "/api/v1/workers/"+w.ID, nil) + rr := httptest.NewRecorder() + r.ServeHTTP(rr, req) + if rr.Code != http.StatusInternalServerError { + t.Fatalf("expected 500, got %d — body: %s", rr.Code, rr.Body) + } + if _, err := st.GetWorker(t.Context(), w.ID); err != nil { + t.Errorf("worker should NOT be deleted when the revoke fails: GetWorker: %v", err) + } + }) } // TestWorkerResponse_Removable verifies the server-authoritative removable flag diff --git a/internal/auth/jointoken/jointoken.go b/internal/auth/jointoken/jointoken.go new file mode 100644 index 00000000..e10d382a --- /dev/null +++ b/internal/auth/jointoken/jointoken.go @@ -0,0 +1,44 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +// Package jointoken mints and hashes worker join tokens: the one-time or +// reusable secrets an operator issues (via the CLI or the REST API) and a +// worker redeems to enroll and receive its broker credential. It is +// server-side only — issuance and redemption both happen on sqi-server — and +// reuses the same random-token and hashing primitives as session tokens and +// API keys (internal/auth/password) rather than re-deriving them. +package jointoken + +import ( + "fmt" + + "github.com/uberware/sqi/internal/auth/password" +) + +// prefix marks a raw token as a worker join token, distinguishing it at a +// glance from an API key (which carries the "sqi_" prefix) in logs and +// operator tooling. +const prefix = "sqiw_" + +// prefixLen is how many leading characters of the raw token are returned for +// list identification only. It is never used to look up or authenticate a +// token (that goes through the full-token hash), so storing a short, +// non-secret slice of the raw token leaks nothing usable. +const prefixLen = 12 + +// Generate creates a new random worker join token. It returns the raw token +// (shown to the operator exactly once), its SHA-256 hash (the only form +// stored at rest), and a short display prefix for list identification. +func Generate() (token, hash, displayPrefix string, err error) { + raw, err := password.GenerateToken() + if err != nil { + return "", "", "", fmt.Errorf("jointoken: generate: %w", err) + } + token = prefix + raw + return token, Hash(token), token[:prefixLen], nil +} + +// Hash returns the hex SHA-256 of a join token, for at-rest storage and +// constant-length lookup. +func Hash(token string) string { + return password.HashToken(token) +} diff --git a/internal/auth/jointoken/jointoken_test.go b/internal/auth/jointoken/jointoken_test.go new file mode 100644 index 00000000..c563ca59 --- /dev/null +++ b/internal/auth/jointoken/jointoken_test.go @@ -0,0 +1,43 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package jointoken_test + +import ( + "strings" + "testing" + + "github.com/uberware/sqi/internal/auth/jointoken" +) + +func TestGenerate(t *testing.T) { + tok, hash, prefix, err := jointoken.Generate() + if err != nil { + t.Fatalf("Generate: %v", err) + } + if !strings.HasPrefix(tok, "sqiw_") { + t.Errorf("token %q lacks the sqiw_ prefix", tok) + } + if strings.Contains(hash, tok) { + t.Error("hash contains the raw token") + } + if !strings.HasPrefix(tok, prefix) { + t.Errorf("prefix %q is not a prefix of token %q", prefix, tok) + } + if got := jointoken.Hash(tok); got != hash { + t.Errorf("Hash = %q, want %q", got, hash) + } +} + +func TestGenerate_Unique(t *testing.T) { + seen := make(map[string]bool, 100) + for range 100 { + tok, _, _, err := jointoken.Generate() + if err != nil { + t.Fatalf("Generate: %v", err) + } + if seen[tok] { + t.Fatalf("duplicate token %q", tok) + } + seen[tok] = true + } +} diff --git a/internal/auth/policy/policy.go b/internal/auth/policy/policy.go index 2a4098ae..6f4141ae 100644 --- a/internal/auth/policy/policy.go +++ b/internal/auth/policy/policy.go @@ -62,6 +62,12 @@ const ( // from InfraManage — which operator holds — because it is an escalation // surface, not ordinary queue configuration. IsolationManage Permission = "isolation.manage" + // WorkersEnroll grants issuing worker join tokens — that is, attaching + // arbitrary compute that receives and executes job code. It is kept + // separate from WorkersManage (enable/disable/delete an existing worker) + // for the same reason IsolationManage is: the privilege is different in + // kind, not merely in degree. + WorkersEnroll Permission = "workers.enroll" ) // grants maps each built-in role to the set of permissions it holds. A role @@ -86,7 +92,7 @@ var grants = map[string]map[Permission]bool{ WorkersRead: true, WorkersManage: true, InfraRead: true, InfraManage: true, ProductsRead: true, ProductsManage: true, DiagnosticsRead: true, UsersRead: true, UsersManage: true, - APIKeysSelf: true, APIKeysAdmin: true, IsolationManage: true, + APIKeysSelf: true, APIKeysAdmin: true, IsolationManage: true, WorkersEnroll: true, }, } @@ -97,7 +103,7 @@ var All = []Permission{ JobsRead, JobsReadAll, JobsWrite, JobsSubmitAs, WorkersRead, WorkersManage, InfraRead, InfraManage, ProductsRead, ProductsManage, DiagnosticsRead, UsersRead, UsersManage, - APIKeysSelf, APIKeysAdmin, IsolationManage, + APIKeysSelf, APIKeysAdmin, IsolationManage, WorkersEnroll, } // Roles returns a read-only snapshot of the role → permission grants matrix, diff --git a/internal/brokerauth/keys.go b/internal/brokerauth/keys.go new file mode 100644 index 00000000..2ccb5ed6 --- /dev/null +++ b/internal/brokerauth/keys.go @@ -0,0 +1,168 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +// Package brokerauth holds the credential primitives shared by sqi-server and +// sqi-worker for NATS broker authentication. +// +// It is deliberately a LEAF package: it imports neither internal/store nor +// internal/openjd, so the worker binary — which can never import the latter — +// may use it directly. +package brokerauth + +import ( + "errors" + "fmt" + "os" + "path/filepath" + "runtime" + "strings" + + "github.com/nats-io/nkeys" +) + +// GenerateSeed creates a new Ed25519 user nkey and returns its seed and the +// corresponding public key. The public key is returned alongside so that no +// caller has to derive it, and so none can derive it wrongly. +func GenerateSeed() (seed []byte, publicKey string, err error) { + kp, err := nkeys.CreateUser() + if err != nil { + return nil, "", fmt.Errorf("brokerauth: create user key: %w", err) + } + seed, err = kp.Seed() + if err != nil { + return nil, "", fmt.Errorf("brokerauth: extract seed: %w", err) + } + publicKey, err = kp.PublicKey() + if err != nil { + return nil, "", fmt.Errorf("brokerauth: extract public key: %w", err) + } + return seed, publicKey, nil +} + +// PublicKeyFromSeed derives the public key for a seed. +func PublicKeyFromSeed(seed []byte) (string, error) { + kp, err := nkeys.FromSeed(seed) + if err != nil { + return "", fmt.Errorf("brokerauth: parse seed: %w", err) + } + pub, err := kp.PublicKey() + if err != nil { + return "", fmt.Errorf("brokerauth: extract public key: %w", err) + } + return pub, nil +} + +// SaveSeed writes seed to path with owner-only permissions, creating parent +// directories as needed. +// +// The write is atomic: seed is written to a temporary file in the same +// directory as path, synced, and moved into place with os.Rename. A plain +// os.WriteFile truncates the target before writing, so a crash, power loss, +// or full disk between the truncate and the write would otherwise leave a +// zero-length seed file — a state worse than a missing one, since it is +// indistinguishable from "enrolled" to every caller that only checks +// existence, and the worker can never recover without an operator deleting +// the file by hand. Renaming within one directory is atomic on every +// platform sqi supports, so a reader always observes either the previous +// seed or the new one, never a truncated one. The temp file is created in +// filepath.Dir(path) rather than the system temp directory because a rename +// across filesystems is not atomic and can fail outright. +// +// The temp file is synced before the rename so the new bytes are not left +// sitting only in the page cache: a rename is atomic with respect to +// ordering, but on a hard power loss an unsynced write can still vanish, +// silently reverting a saved seed to whatever was on disk before it. The +// file is small enough (an nkey seed is a few dozen bytes) that this costs +// nothing worth avoiding. +// +// The final mode is 0600 regardless of what, if anything, existed at path +// before: the temp file is created 0600 by os.CreateTemp and that mode is +// set explicitly rather than relied upon, and the rename replaces whatever +// was there — including a more permissive leftover from key rotation, an +// earlier bug, or an operator who chmod'd it to look at it — with a file +// that was never observable under any mode but 0600. +// +// Because the write goes through a temporary file created in the same +// directory as path, the caller needs write permission on that CONTAINING +// DIRECTORY, not just on the seed file itself: a directory an operator locked +// down to 0500 while leaving an owner-writable seed inside it will fail here +// even though a direct write to the existing file would have succeeded. +func SaveSeed(path string, seed []byte) error { + dir := filepath.Dir(path) + if err := os.MkdirAll(dir, 0o700); err != nil { + return fmt.Errorf("brokerauth: create seed dir: %w", err) + } + + tmp, err := os.CreateTemp(dir, "."+filepath.Base(path)+".tmp-*") + if err != nil { + return fmt.Errorf("brokerauth: create temp seed file: %w", err) + } + tmpPath := tmp.Name() + renamed := false + defer func() { + if !renamed { + _ = os.Remove(tmpPath) + } + }() + + if chmodErr := tmp.Chmod(0o600); chmodErr != nil { + _ = tmp.Close() + return fmt.Errorf("brokerauth: chmod temp seed file %s: %w", tmpPath, chmodErr) + } + if _, writeErr := tmp.Write(seed); writeErr != nil { + _ = tmp.Close() + return fmt.Errorf("brokerauth: write temp seed file %s: %w", tmpPath, writeErr) + } + if syncErr := tmp.Sync(); syncErr != nil { + _ = tmp.Close() + return fmt.Errorf("brokerauth: sync temp seed file %s: %w", tmpPath, syncErr) + } + if closeErr := tmp.Close(); closeErr != nil { + return fmt.Errorf("brokerauth: close temp seed file %s: %w", tmpPath, closeErr) + } + + if renameErr := os.Rename(tmpPath, path); renameErr != nil { + return fmt.Errorf("brokerauth: rename seed into place %s: %w", path, renameErr) + } + renamed = true + return nil +} + +// LoadSeed reads a seed file, refusing one that is readable beyond its owner. +// +// The check is a real one, not hygiene theater: this seed IS the worker's +// identity, and a group-readable seed on a shared render node hands that +// identity to every account on the box. The mode check is skipped on Windows, +// where POSIX bits do not carry the same meaning. +func LoadSeed(path string) ([]byte, error) { + info, err := os.Stat(path) + if err != nil { + return nil, fmt.Errorf("brokerauth: stat seed %s: %w", path, err) + } + if runtime.GOOS != "windows" { + if perm := info.Mode().Perm(); perm&0o077 != 0 { + return nil, fmt.Errorf( + "brokerauth: seed file %s is mode %o; it must be readable only by its owner — run: chmod 600 %s", + path, perm, path, + ) + } + } + seed, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("brokerauth: read seed %s: %w", path, err) + } + if _, err := nkeys.FromSeed(seed); err != nil { + return nil, fmt.Errorf("brokerauth: seed file %s is not a valid nkey seed: %w", path, err) + } + return seed, nil +} + +// ValidatePublicKey reports whether pk is a well-formed user nkey. +func ValidatePublicKey(pk string) error { + if !strings.HasPrefix(pk, "U") { + return errors.New("brokerauth: public key must be a user nkey (starts with 'U')") + } + if !nkeys.IsValidPublicUserKey(pk) { + return errors.New("brokerauth: public key is not a valid user nkey") + } + return nil +} diff --git a/internal/brokerauth/keys_test.go b/internal/brokerauth/keys_test.go new file mode 100644 index 00000000..8e4c6fe3 --- /dev/null +++ b/internal/brokerauth/keys_test.go @@ -0,0 +1,234 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package brokerauth_test + +import ( + "os" + "path/filepath" + "runtime" + "strings" + "testing" + + "github.com/nats-io/nkeys" + + "github.com/uberware/sqi/internal/brokerauth" +) + +func TestGenerateSeed_RoundTrips(t *testing.T) { + seed, pub, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed: %v", err) + } + if !strings.HasPrefix(pub, "U") { + t.Errorf("public key %q is not a user nkey", pub) + } + got, err := brokerauth.PublicKeyFromSeed(seed) + if err != nil { + t.Fatalf("PublicKeyFromSeed: %v", err) + } + if got != pub { + t.Errorf("PublicKeyFromSeed = %q, want %q", got, pub) + } +} + +func TestSaveSeed_WritesOwnerOnly(t *testing.T) { + seed, _, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed: %v", err) + } + path := filepath.Join(t.TempDir(), "worker.nk") + if err := brokerauth.SaveSeed(path, seed); err != nil { + t.Fatalf("SaveSeed: %v", err) + } + info, err := os.Stat(path) + if err != nil { + t.Fatalf("Stat: %v", err) + } + if perm := info.Mode().Perm(); perm != 0o600 { + t.Errorf("mode = %o, want 600", perm) + } +} + +func TestSaveSeed_FixesModeOfExistingFile(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("POSIX mode bits") + } + seed, _, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed: %v", err) + } + path := filepath.Join(t.TempDir(), "worker.nk") + + // Simulate a pre-existing, more permissive file: key rotation over an + // old seed, or a file an operator chmod'd to inspect. os.WriteFile only + // applies its perm argument when it creates the file, so writing over an + // existing 0644 file must not leave it 0644. + if err := os.WriteFile(path, []byte("stale"), 0o644); err != nil { + t.Fatalf("seed pre-write: %v", err) + } + + if err := brokerauth.SaveSeed(path, seed); err != nil { + t.Fatalf("SaveSeed: %v", err) + } + info, err := os.Stat(path) + if err != nil { + t.Fatalf("Stat: %v", err) + } + if perm := info.Mode().Perm(); perm != 0o600 { + t.Errorf("mode = %o, want 600", perm) + } +} + +func TestSaveSeed_LeavesNoTempFileBehind(t *testing.T) { + seed, _, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed: %v", err) + } + dir := t.TempDir() + path := filepath.Join(dir, "worker.nk") + if err := brokerauth.SaveSeed(path, seed); err != nil { + t.Fatalf("SaveSeed: %v", err) + } + + entries, err := os.ReadDir(dir) + if err != nil { + t.Fatalf("ReadDir: %v", err) + } + if len(entries) != 1 { + names := make([]string, len(entries)) + for i, e := range entries { + names[i] = e.Name() + } + t.Fatalf("directory contains %v, want exactly [%s]", names, filepath.Base(path)) + } + if got := entries[0].Name(); got != filepath.Base(path) { + t.Errorf("directory contains %q, want %q", got, filepath.Base(path)) + } +} + +func TestSaveSeed_NoTempFileSurvivesCreateFailure(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("POSIX permission semantics") + } + if os.Geteuid() == 0 { + t.Skip("root ignores directory write permission") + } + + seed, _, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed: %v", err) + } + dir := t.TempDir() + // Pre-create the seed path with a normal save, then make the directory + // read-only. A second SaveSeed then fails deterministically at + // os.CreateTemp (it needs write permission on the directory to create + // the temp file) before ever touching the existing seed. + path := filepath.Join(dir, "worker.nk") + if err := brokerauth.SaveSeed(path, seed); err != nil { + t.Fatalf("SaveSeed (setup): %v", err) + } + + if err := os.Chmod(dir, 0o500); err != nil { + t.Fatalf("Chmod dir: %v", err) + } + //nolint:errcheck // best-effort restore so t.TempDir() can clean up even if the test fails earlier + defer func() { _ = os.Chmod(dir, 0o700) }() + + seed2, _, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed: %v", err) + } + if err := brokerauth.SaveSeed(path, seed2); err == nil { + t.Fatal("SaveSeed succeeded against a read-only directory; want error") + } + + if err := os.Chmod(dir, 0o700); err != nil { + t.Fatalf("Chmod dir (restore): %v", err) + } + entries, err := os.ReadDir(dir) + if err != nil { + t.Fatalf("ReadDir: %v", err) + } + if len(entries) != 1 || entries[0].Name() != filepath.Base(path) { + names := make([]string, len(entries)) + for i, e := range entries { + names[i] = e.Name() + } + t.Fatalf("directory contains %v after a failed save, want exactly [%s] (the untouched original)", names, filepath.Base(path)) + } +} + +func TestLoadSeed_RejectsPermissiveMode(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("POSIX mode bits") + } + seed, _, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed: %v", err) + } + path := filepath.Join(t.TempDir(), "worker.nk") + if err := brokerauth.SaveSeed(path, seed); err != nil { + t.Fatalf("SaveSeed: %v", err) + } + if err := os.Chmod(path, 0o644); err != nil { + t.Fatalf("Chmod: %v", err) + } + if _, err := brokerauth.LoadSeed(path); err == nil { + t.Error("LoadSeed accepted a world-readable seed file; want error") + } +} + +func TestValidatePublicKey(t *testing.T) { + userKP, err := nkeys.CreateUser() + if err != nil { + t.Fatalf("CreateUser: %v", err) + } + validUser, err := userKP.PublicKey() + if err != nil { + t.Fatalf("PublicKey: %v", err) + } + seed, err := userKP.Seed() + if err != nil { + t.Fatalf("Seed: %v", err) + } + + accountKP, err := nkeys.CreateAccount() + if err != nil { + t.Fatalf("CreateAccount: %v", err) + } + accountKey, err := accountKP.PublicKey() + if err != nil { + t.Fatalf("PublicKey: %v", err) + } + + // Corrupt the last character of an otherwise-valid user key so its + // trailing CRC16 no longer checks out, without changing its length or + // its "U" prefix. + corrupted := []byte(validUser) + if last := corrupted[len(corrupted)-1]; last == 'A' { + corrupted[len(corrupted)-1] = 'B' + } else { + corrupted[len(corrupted)-1] = 'A' + } + + tests := []struct { + name string + pk string + wantErr bool + }{ + {"valid generated user key", validUser, false}, + {"seed instead of a public key", string(seed), true}, + {"account key instead of a user key", accountKey, true}, + {"user key with corrupted CRC", string(corrupted), true}, + {"empty string", "", true}, + {"U-prefixed but not valid base32", "U!!!not-base32!!!", true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := brokerauth.ValidatePublicKey(tt.pk) + if (err != nil) != tt.wantErr { + t.Errorf("ValidatePublicKey(%q) error = %v, wantErr %v", tt.pk, err, tt.wantErr) + } + }) + } +} diff --git a/internal/brokerauth/natsoption.go b/internal/brokerauth/natsoption.go new file mode 100644 index 00000000..43908309 --- /dev/null +++ b/internal/brokerauth/natsoption.go @@ -0,0 +1,23 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package brokerauth + +import ( + "github.com/nats-io/nats.go" + "github.com/nats-io/nkeys" +) + +// NkeyOption returns a [nats.Option] that authenticates a connection as the +// nkey identified by publicKey, signing the server's nonce with seed. Both +// sqi-server (its admin connection, internal/bus.Broker.adminOptions) and +// sqi-worker (internal/worker/natsclient.buildOptions) use it to present +// their own credential to the broker. +func NkeyOption(publicKey string, seed []byte) nats.Option { + return nats.Nkey(publicKey, func(nonce []byte) ([]byte, error) { + kp, err := nkeys.FromSeed(seed) + if err != nil { + return nil, err + } + return kp.Sign(nonce) + }) +} diff --git a/internal/brokerauth/natsoption_test.go b/internal/brokerauth/natsoption_test.go new file mode 100644 index 00000000..cc888e0c --- /dev/null +++ b/internal/brokerauth/natsoption_test.go @@ -0,0 +1,46 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package brokerauth_test + +import ( + "testing" + + "github.com/nats-io/nats.go" + "github.com/nats-io/nkeys" + + "github.com/uberware/sqi/internal/brokerauth" +) + +func TestNkeyOption_SetsPublicKeyAndSignsNonce(t *testing.T) { + seed, pub, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed: %v", err) + } + + opt := brokerauth.NkeyOption(pub, seed) + + var o nats.Options + if err := opt(&o); err != nil { + t.Fatalf("apply option: %v", err) + } + if o.Nkey != pub { + t.Errorf("Options.Nkey = %q, want %q", o.Nkey, pub) + } + if o.SignatureCB == nil { + t.Fatal("Options.SignatureCB not set") + } + + nonce := []byte("test-nonce") + sig, err := o.SignatureCB(nonce) + if err != nil { + t.Fatalf("SignatureCB: %v", err) + } + + kp, err := nkeys.FromPublicKey(pub) + if err != nil { + t.Fatalf("nkeys.FromPublicKey: %v", err) + } + if err := kp.Verify(nonce, sig); err != nil { + t.Errorf("signature does not verify against the public key: %v", err) + } +} diff --git a/internal/brokerauth/permissions.go b/internal/brokerauth/permissions.go new file mode 100644 index 00000000..ec4ab2c1 --- /dev/null +++ b/internal/brokerauth/permissions.go @@ -0,0 +1,118 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package brokerauth + +import ( + "strings" + + natsserver "github.com/nats-io/nats-server/v2/server" +) + +// inboxPrefixRoot is the first token of every per-worker reply-inbox +// prefix. It deliberately does NOT contain a "." so that the whole prefix is +// a single subject token and ".>" covers exactly one worker's +// inboxes. +const inboxPrefixRoot = "_INBOX_" + +// InboxPrefix returns the reply-inbox subject prefix reserved for workerID. +// +// nats.go's default prefix is the process-global "_INBOX", under which each +// connection takes "_INBOX..*". Granting a worker "_INBOX.>" +// would therefore let it subscribe to every other client's reply inbox on +// this broker. Work leases are core-NATS request/reply and the batch is +// delivered with msg.Respond (internal/bus/lease.go), so one enrolled worker +// could read every other worker's assignment — the OnRun command line and +// arguments, embedded file contents, job and task parameters, environment +// variables, the path map and the isolation account — plus the JetStream API +// replies to sqi-server's own client, without ever leasing a task itself. +// +// Giving each worker its own prefix and granting only that subtree closes +// it. Workers connect with nats.CustomInboxPrefix(InboxPrefix(workerID)) — +// see internal/worker/natsclient. The server needs no matching change: it +// publishes ">" (see [ServerPermissions]), so msg.Respond and JetStream +// PubAcks still reach whichever prefix a worker chose. +func InboxPrefix(workerID string) string { + return inboxPrefixRoot + workerID +} + +// ValidWorkerIDToken reports whether workerID may be used as a single NATS +// subject token — which every worker→server subject, [InboxPrefix] and +// [WorkerPermissions] require of it. +// +// Worker IDs are UUIDs generated by workerconfig.LoadOrCreateWorkerID, so in +// practice this always holds. It is checked rather than assumed at three +// points, for two different reasons: +// +// - At ENROLLMENT (internal/api's enroll handler, and sqi-server's worker +// enroll command) the ID is supplied by the caller and becomes a subject +// PATTERN in [WorkerPermissions]. "*" would mint a credential granted +// "task.status.*.*", "worker.deregister.*", "work.lease.*.*" and the +// rest — able to publish concrete subjects belonging to any worker, so +// it could forge status and logs, deregister the farm, and lease work as +// another worker and receive that worker's assignment batch. ">" instead +// yields the malformed "task.status.>.*", which nats-server rejects: it +// would fail every credential reload, or stop the broker booting. +// internal/bus's buildNkeys re-checks and skips, so a row that reached +// the database some other way cannot do either. +// - At the WORKER's own connect path (internal/worker/natsclient) the ID +// comes from a file an operator can edit, and one containing "." would +// silently widen the inbox subtree its grant covers. +func ValidWorkerIDToken(workerID string) bool { + if workerID == "" { + return false + } + return !strings.ContainsAny(workerID, ". \t\r\n*>") +} + +// WorkerPermissions returns the NATS permissions for one worker. +// +// Every worker→server subject carries the worker's ID as a token precisely so +// that these permissions can bind it. That is not a stylistic choice: NATS +// permissions are static per credential and JetStream does not stamp +// publisher identity onto a message, so a subject scheme keyed only by job +// and task would give no way to express "only this worker's own traffic" — +// the worker's identity has to live in the subject itself. +func WorkerPermissions(workerID string) *natsserver.Permissions { + return &natsserver.Permissions{ + Publish: &natsserver.SubjectPermission{ + Allow: []string{ + "task.status." + workerID + ".*", + "task.logs." + workerID + ".*", + "worker.register." + workerID, + "worker.heartbeat." + workerID, + "worker.deregister." + workerID, + "worker.diag." + workerID, + "work.lease." + workerID + ".*", + }, + }, + Subscribe: &natsserver.SubjectPermission{ + Allow: []string{ + // This worker's own reply inboxes, and no other client's — + // see [InboxPrefix] for why the process-global "_INBOX.>" + // cannot be granted here. A request/reply lease cannot work + // without this. + InboxPrefix(workerID) + ".>", + + // ACCEPTED GAP. Workers subscribe per-task at assignment time + // (internal/worker/cancel/cancel.go), so a static permission + // cannot be narrower than the whole subtree: any enrolled + // worker can observe any task's cancel signals. Narrowing + // this needs a permission reload per assignment, or a NATS + // auth callout. Low severity — a cancel message for a task + // the worker does not hold is inert. + "task.cancel.>", + }, + }, + } +} + +// ServerPermissions returns the permissions for sqi-server's own broker +// connections, which need the full subject space: the scheduler consumes every +// worker subject and publishes cancels, and the broker's admin connection +// provisions streams. +func ServerPermissions() *natsserver.Permissions { + return &natsserver.Permissions{ + Publish: &natsserver.SubjectPermission{Allow: []string{">"}}, + Subscribe: &natsserver.SubjectPermission{Allow: []string{">"}}, + } +} diff --git a/internal/brokerauth/permissions_test.go b/internal/brokerauth/permissions_test.go new file mode 100644 index 00000000..1293a195 --- /dev/null +++ b/internal/brokerauth/permissions_test.go @@ -0,0 +1,135 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package brokerauth_test + +import ( + "strings" + "testing" + + natsserver "github.com/nats-io/nats-server/v2/server" + + "github.com/uberware/sqi/internal/brokerauth" +) + +// subjectAllowed reports whether subject matches an entry in sp.Allow. It +// supports the two NATS subject wildcards: "*" matches exactly one token, and +// a trailing ">" matches one or more trailing tokens. A subject with no +// matching entry is denied. This exists to check the permission data returned +// by WorkerPermissions, not to reimplement NATS subject matching in general. +func subjectAllowed(sp *natsserver.SubjectPermission, subject string) bool { + if sp == nil { + return false + } + for _, pattern := range sp.Allow { + if subjectMatches(pattern, subject) { + return true + } + } + return false +} + +func subjectMatches(pattern, subject string) bool { + patternTokens := strings.Split(pattern, ".") + subjectTokens := strings.Split(subject, ".") + + for i, pt := range patternTokens { + if pt == ">" { + // ">" must be the final token and matches one or more remaining + // subject tokens. + return i == len(patternTokens)-1 && i < len(subjectTokens) + } + if i >= len(subjectTokens) { + return false + } + if pt == "*" { + continue + } + if pt != subjectTokens[i] { + return false + } + } + return len(patternTokens) == len(subjectTokens) +} + +func TestWorkerPermissions_ConfinesToOwnSubtree(t *testing.T) { + const me = "worker-a" + const other = "worker-b" + p := brokerauth.WorkerPermissions(me) + + allowedPublish := []string{ + "task.status." + me + ".job-1", + "task.logs." + me + ".task-1", + "worker.register." + me, + "worker.heartbeat." + me, + "worker.deregister." + me, + "worker.diag." + me, + "work.lease." + me + ".queue-1", + } + for _, s := range allowedPublish { + if !subjectAllowed(p.Publish, s) { + t.Errorf("publish %q should be allowed", s) + } + } + + deniedPublish := []string{ + "task.status." + other + ".job-1", + "task.logs." + other + ".task-1", + "worker.deregister." + other, + "work.lease." + other + ".queue-1", + "task.cancel.task-1", + "task.status.>", + } + for _, s := range deniedPublish { + if subjectAllowed(p.Publish, s) { + t.Errorf("publish %q must NOT be allowed", s) + } + } + + if subjectAllowed(p.Subscribe, "task.status.>") { + t.Error("a worker must not be able to subscribe to task.status.>") + } + if !subjectAllowed(p.Subscribe, "task.cancel.task-1") { + t.Error("a worker must be able to subscribe to task.cancel.") + } + + // Reply inboxes. A worker needs its own, because a work lease is + // core-NATS request/reply — and must not have anyone else's, because the + // lease reply carries the whole assignment batch. + if !subjectAllowed(p.Subscribe, brokerauth.InboxPrefix(me)+".x.1") { + t.Error("a worker must be able to subscribe to its own reply inbox") + } + if subjectAllowed(p.Subscribe, brokerauth.InboxPrefix(other)+".x.1") { + t.Error("a worker must NOT be able to subscribe to another worker's reply inbox") + } + if subjectAllowed(p.Subscribe, "_INBOX.x.1") { + t.Error("a worker must NOT be able to subscribe to the process-global _INBOX subtree") + } +} + +func TestValidWorkerIDToken(t *testing.T) { + valid := []string{ + "0f1d2c3b-4a59-6879-8a9b-0c1d2e3f4a5b", // what LoadOrCreateWorkerID writes + "worker-a", + "render_01", + } + for _, s := range valid { + if !brokerauth.ValidWorkerIDToken(s) { + t.Errorf("ValidWorkerIDToken(%q) = false, want true", s) + } + } + + invalid := []string{ + "", // no token at all + "a.b", // two tokens: widens every grant built from it + "a b", // whitespace is not legal in a subject + "*", // single-token wildcard + ">", // multi-token wildcard + "worker.>", // both of the above + "worker\tname", // other whitespace + } + for _, s := range invalid { + if brokerauth.ValidWorkerIDToken(s) { + t.Errorf("ValidWorkerIDToken(%q) = true, want false", s) + } + } +} diff --git a/internal/bus/broker.go b/internal/bus/broker.go index 3d7c755f..96290417 100644 --- a/internal/bus/broker.go +++ b/internal/bus/broker.go @@ -9,11 +9,14 @@ import ( "log/slog" "net" "strconv" + "sync" "time" natsserver "github.com/nats-io/nats-server/v2/server" nats "github.com/nats-io/nats.go" "github.com/nats-io/nats.go/jetstream" + + "github.com/uberware/sqi/internal/brokerauth" ) // BrokerConfig holds the parameters needed to start the embedded NATS server. @@ -21,9 +24,14 @@ type BrokerConfig struct { // Addr is the TCP address the embedded NATS server binds to, in // "host:port" form. Defaults to "0.0.0.0:4222" (all interfaces) so that // workers which discover the server via mDNS can reach the broker at the - // advertised LAN host. Broker authentication does not exist: any host - // that can reach this port can register as a worker and receive - // assignments. Deferred to Phase 4 hardening. + // advertised LAN host. + // + // Broker authentication is OPT-IN and off by default. With it off, any + // host that can reach this port can register as a worker and receive + // assignments — including on a LAN, since the default binds all + // interfaces. sqi-server emits a startup WARN in exactly that case. + // Set nats.auth.enabled to require a per-worker credential, or bind this + // to 127.0.0.1:4222 for single-machine use. See docs/auth.md. Addr string // DataDir is the directory JetStream uses for file-backed stream storage. @@ -33,6 +41,29 @@ type BrokerConfig struct { // MaxStoreMB is the maximum disk space JetStream may use, in megabytes. // A value of 0 means unlimited (not recommended for production). MaxStoreMB int + + // Auth configures per-worker nkey authorization on the broker. Zero value + // leaves authorization disabled. + Auth BrokerAuthConfig +} + +// BrokerAuthConfig controls per-worker nkey authorization on the broker. +type BrokerAuthConfig struct { + // Enabled requires every connection to present a credential for an + // enrolled nkey. When false, the broker accepts anonymous connections. + Enabled bool + + // Credentials is the initial enrolled set, loaded from the store at boot. + Credentials []WorkerCredentialRef +} + +// WorkerCredentialRef is bus's own minimal view of an enrolled worker +// credential, carrying only what the broker needs to authorize a connection. +// It exists so that internal/bus does not need to import internal/store; +// callers map from store.WorkerCredential to WorkerCredentialRef themselves. +type WorkerCredentialRef struct { + WorkerID string + PublicKey string } // Broker wraps an in-process NATS server with JetStream enabled and manages @@ -45,8 +76,31 @@ type Broker struct { cfg BrokerConfig logger *slog.Logger + // mu guards every field below. It exists because ReloadCredentials can + // now be called from an HTTP handler goroutine (worker credential + // revocation, see internal/server) concurrently with Shutdown on the + // process's exit path — a combination that did not exist before that + // caller, since previously only Start (single-threaded, before the + // broker is handed to anything else) touched these fields. + // Without it, Shutdown nilling ns/nc while ReloadCredentials or Check + // reads them is a data race and a nil-pointer hazard, not just a + // theoretical one. + mu sync.Mutex + ns *natsserver.Server // embedded NATS server process nc *nats.Conn // admin connection used for stream provisioning + + // serverSeed and serverPub are the broker's own nkey, generated at boot + // and held only in memory, so that the broker's own connections (stream + // provisioning, the scheduler's in-process client) can authenticate once + // authorization is enabled. Empty when authorization is disabled. + serverSeed []byte + serverPub string + + // bootOpts is a pristine copy of the options the server was started + // with, retained so ReloadCredentials can clone from it rather than + // reusing options nats-server has already consumed. + bootOpts *natsserver.Options } // New creates a [Broker] with the given configuration and logger. @@ -95,6 +149,23 @@ func (b *Broker) Start(ctx context.Context) error { NoLog: true, } + var serverSeed []byte + var serverPub string + if b.cfg.Auth.Enabled { + seed, pub, err := brokerauth.GenerateSeed() + if err != nil { + return fmt.Errorf("bus: generate server key: %w", err) + } + serverSeed, serverPub = seed, pub + opts.Nkeys = buildNkeys(serverPub, b.cfg.Auth.Credentials, b.logger) + } + + // Retain a pristine copy of the boot options for ReloadCredentials: + // ReloadOptions documents that the Options passed to it must not be + // reused, so credential reloads clone from this copy rather than the one + // nats-server consumes below. + bootOpts := opts.Clone() + ns, err := natsserver.NewServer(opts) if err != nil { return fmt.Errorf("bus: create nats server: %w", err) @@ -109,7 +180,14 @@ func (b *Broker) Start(ctx context.Context) error { return errors.New("bus: nats server did not become ready within 10s") } - b.ns = ns + // Commit every field the other methods read, in one critical section, so + // no concurrent caller can observe ns non-nil while serverSeed/ + // serverPub/bootOpts are still zero, or any other partially-updated + // combination. + b.mu.Lock() + b.serverSeed, b.serverPub, b.bootOpts, b.ns = serverSeed, serverPub, bootOpts, ns + b.mu.Unlock() + b.logger.InfoContext( ctx, "bus: nats server started", slog.String("addr", b.cfg.Addr), @@ -120,13 +198,15 @@ func (b *Broker) Start(ctx context.Context) error { // Establish an admin connection used only for stream provisioning. // This is a plain TCP connection to the loopback listener; the latency // is negligible and avoids importing the server package into callers. - nc, err := nats.Connect(ns.ClientURL()) + nc, err := nats.Connect(ns.ClientURL(), b.adminOptions()...) if err != nil { ns.Shutdown() ns.WaitForShutdown() return fmt.Errorf("bus: admin connect: %w", err) } + b.mu.Lock() b.nc = nc + b.mu.Unlock() js, err := jetstream.New(nc) if err != nil { @@ -153,15 +233,29 @@ func (b *Broker) Start(ctx context.Context) error { // Shutdown drains the admin connection and stops the embedded NATS server, // waiting until it has fully exited before returning. It is safe to call // Shutdown more than once; subsequent calls are no-ops. +// +// It is also safe to call concurrently with [Broker.ReloadCredentials], +// [Broker.Check], [Broker.ClientURL] and [Broker.NewClient]: the fields +// those methods read are captured under mu +// and nilled here under the same lock, so a concurrent reader always sees +// either the pre-shutdown values or nil, never a torn or dangling pointer. +// The (possibly slow) nats-server calls below run after the lock is +// released, so Shutdown does not hold up an in-flight reload or health +// check any longer than it takes to copy two pointers. func (b *Broker) Shutdown() { - if b.nc != nil { - b.nc.Close() - b.nc = nil - } - if b.ns != nil { - b.ns.Shutdown() - b.ns.WaitForShutdown() - b.ns = nil + b.mu.Lock() + nc := b.nc + ns := b.ns + b.nc = nil + b.ns = nil + b.mu.Unlock() + + if nc != nil { + nc.Close() + } + if ns != nil { + ns.Shutdown() + ns.WaitForShutdown() b.logger.InfoContext(context.Background(), "bus: nats server stopped") } } @@ -171,10 +265,13 @@ func (b *Broker) Shutdown() { // otherwise. Registered with the health registry during server startup so // that GET /readyz reflects broker health. func (b *Broker) Check(_ context.Context) error { - if b.ns == nil || !b.ns.Running() { + b.mu.Lock() + ns, nc := b.ns, b.nc + b.mu.Unlock() + if ns == nil || !ns.Running() { return errors.New("nats server not running") } - if b.nc == nil || b.nc.IsClosed() { + if nc == nil || nc.IsClosed() { return errors.New("nats admin connection closed") } return nil @@ -183,10 +280,13 @@ func (b *Broker) Check(_ context.Context) error { // ClientURL returns the nats:// URL that in-process clients should connect to. // Returns an empty string if the broker has not been started yet. func (b *Broker) ClientURL() string { - if b.ns == nil { + b.mu.Lock() + ns := b.ns + b.mu.Unlock() + if ns == nil { return "" } - return b.ns.ClientURL() + return ns.ClientURL() } // NewClient dials the embedded NATS server and returns a connected [Client] @@ -197,8 +297,118 @@ func (b *Broker) ClientURL() string { // component (or one shared instance for the whole server) rather than dialing // repeatedly. func (b *Broker) NewClient() (*Client, error) { - if b.ns == nil { + b.mu.Lock() + ns := b.ns + b.mu.Unlock() + if ns == nil { return nil, errors.New("bus: broker not started") } - return NewClient(b.ns.ClientURL(), b.logger) + return NewClient(ns.ClientURL(), b.logger, b.adminOptions()...) +} + +// buildNkeys converts the enrolled credential set into NATS nkey users, plus +// the server's own credential, identified by serverPub. It is a standalone +// function of its arguments rather than a *Broker method reading b.serverPub +// so that both Start (before it publishes serverPub to b) and +// ReloadCredentials (which reads it under b.mu itself) can call it without +// also needing to take the lock — and risk taking it twice on the same +// goroutine. +// A credential whose worker ID is not a single NATS subject token, or whose +// public key is not a valid nkey, is SKIPPED and logged rather than +// installed. Both enrollment boundaries reject such rows before they can be +// stored (internal/api's enroll handler and sqi-server's worker enroll +// command validate both the worker ID and the public key), so this is +// defense in depth for a row that arrived some other way — a hand-edited +// database, or a binary predating those checks. The worker-ID case matters +// because the worker ID becomes a subject PATTERN here: installing "*" would +// grant one credential "task.status.*.*", "worker.deregister.*", +// "work.lease.*.*" and the rest, letting it publish concrete subjects +// belonging to any worker on the farm; installing ">" (or an empty or +// whitespace-bearing ID) produces a malformed subject that nats-server +// refuses. The public-key case matters because an nkey user with a +// malformed key is itself an option natsserver.NewServer rejects outright — +// without this guard, one bad row takes the whole broker down rather than +// costing the one worker it belongs to; with it, every good credential keeps +// working. +func buildNkeys(serverPub string, creds []WorkerCredentialRef, logger *slog.Logger) []*natsserver.NkeyUser { + users := make([]*natsserver.NkeyUser, 0, len(creds)+1) + users = append(users, &natsserver.NkeyUser{ + Nkey: serverPub, + Permissions: brokerauth.ServerPermissions(), + }) + for _, c := range creds { + if !brokerauth.ValidWorkerIDToken(c.WorkerID) { + logger.WarnContext( + context.Background(), + "bus: skipping a worker credential whose worker id is not a valid NATS subject token", + slog.String("worker_id", c.WorkerID), + slog.String("impact", "this worker cannot connect; its grants would have been subject wildcards or a malformed subject"), + slog.String("remediation", "revoke the credential and re-enroll the worker with an id containing no '.', whitespace, '*' or '>'"), + ) + continue + } + if err := brokerauth.ValidatePublicKey(c.PublicKey); err != nil { + logger.WarnContext( + context.Background(), + "bus: skipping a worker credential whose public key is not a valid nkey", + slog.String("worker_id", c.WorkerID), + slog.Any("error", err), + slog.String("impact", "this worker cannot connect; an invalid key here would otherwise stop the whole broker from starting"), + slog.String("remediation", "revoke the credential and re-enroll the worker with a key generated by sqi-worker keygen"), + ) + continue + } + users = append(users, &natsserver.NkeyUser{ + Nkey: c.PublicKey, + Permissions: brokerauth.WorkerPermissions(c.WorkerID), + }) + } + return users +} + +// adminOptions returns the connect options for the broker's own admin +// connection, which provisions streams. Empty when auth is disabled. +func (b *Broker) adminOptions() []nats.Option { + if !b.cfg.Auth.Enabled { + return nil + } + b.mu.Lock() + pub, seed := b.serverPub, b.serverSeed + b.mu.Unlock() + return []nats.Option{brokerauth.NkeyOption(pub, seed)} +} + +// ReloadCredentials replaces the enrolled worker set on a running broker. +// +// Revocation is synchronous, not eventually-consistent: nats-server's +// reloadAuthorization re-runs isClientAuthorized over every connected client +// and calls authViolation() on any that no longer pass, so a worker removed +// from creds is disconnected inside this call. +// +// ReloadOptions rejects changes to options that cannot be hot-swapped and +// documents that the Options passed to it must not be reused, so this clones +// the pristine boot options and mutates only Nkeys. +// +// Safe to call concurrently with [Broker.Shutdown] — see its doc comment. +// nats-server's own ReloadOptions and Shutdown both take the embedded +// server's internal lock, so once this method has read a live *ns off b, the +// call into nats-server itself is safe even if Shutdown wins a concurrent +// race to nil out b.ns first: that only means this call observes "broker not +// started" instead, never a corrupted server. +func (b *Broker) ReloadCredentials(creds []WorkerCredentialRef) error { + if !b.cfg.Auth.Enabled { + return errors.New("bus: broker authentication is disabled") + } + b.mu.Lock() + ns, bootOpts, serverPub := b.ns, b.bootOpts, b.serverPub + b.mu.Unlock() + if ns == nil || bootOpts == nil { + return errors.New("bus: broker not started") + } + opts := bootOpts.Clone() + opts.Nkeys = buildNkeys(serverPub, creds, b.logger) + if err := ns.ReloadOptions(opts); err != nil { + return fmt.Errorf("bus: reload broker credentials: %w", err) + } + return nil } diff --git a/internal/bus/broker_auth_test.go b/internal/bus/broker_auth_test.go new file mode 100644 index 00000000..fa3637b0 --- /dev/null +++ b/internal/bus/broker_auth_test.go @@ -0,0 +1,450 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package bus + +import ( + "context" + "errors" + "log/slog" + "net" + "sync" + "testing" + "time" + + nats "github.com/nats-io/nats.go" + "github.com/nats-io/nkeys" + + "github.com/uberware/sqi/internal/brokerauth" +) + +// startBrokerAuth boots an embedded broker configured with auth, on a temp +// JetStream dir and an OS-assigned loopback port, waits for it to be ready, +// and registers cleanup. +func startBrokerAuth(t *testing.T, auth BrokerAuthConfig) *Broker { + t.Helper() + logger := slog.New(slog.DiscardHandler) + cfg := BrokerConfig{ + Addr: net.JoinHostPort("127.0.0.1", itoa(freePort(t))), + DataDir: t.TempDir() + "/nats", + MaxStoreMB: 64, + Auth: auth, + } + b := New(cfg, logger) + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + if err := b.Start(ctx); err != nil { + t.Fatalf("startBrokerAuth: Start: %v", err) + } + t.Cleanup(b.Shutdown) + return b +} + +// enrolledWorker generates a fresh nkey and returns a WorkerCredentialRef for +// it, along with the raw seed needed to sign connection challenges. +func enrolledWorker(t *testing.T, workerID string) (WorkerCredentialRef, []byte) { + t.Helper() + seed, pub, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("enrolledWorker: GenerateSeed: %v", err) + } + return WorkerCredentialRef{WorkerID: workerID, PublicKey: pub}, seed +} + +// nkeyOption builds a nats.Option that authenticates as the nkey pair +// identified by pub, signing server challenges with seed. +func nkeyOption(t *testing.T, seed []byte, pub string) nats.Option { + t.Helper() + return nats.Nkey(pub, func(nonce []byte) ([]byte, error) { + kp, err := nkeys.FromSeed(seed) + if err != nil { + return nil, err + } + return kp.Sign(nonce) + }) +} + +func TestBrokerAuth(t *testing.T) { + t.Run("auth disabled accepts an anonymous connection", func(t *testing.T) { + b := startBrokerAuth(t, BrokerAuthConfig{Enabled: false}) + nc, err := nats.Connect(b.ClientURL()) + if err != nil { + t.Fatalf("Connect: %v", err) + } + defer nc.Close() + }) + + t.Run("auth enabled refuses an anonymous connection", func(t *testing.T) { + b := startBrokerAuth(t, BrokerAuthConfig{Enabled: true}) + if _, err := nats.Connect(b.ClientURL()); err == nil { + t.Fatal("Connect: want error for anonymous connection, got nil") + } + }) + + t.Run("auth enabled accepts an enrolled nkey", func(t *testing.T) { + ref, seed := enrolledWorker(t, "worker-a") + b := startBrokerAuth(t, BrokerAuthConfig{ + Enabled: true, + Credentials: []WorkerCredentialRef{ref}, + }) + nc, err := nats.Connect(b.ClientURL(), nkeyOption(t, seed, ref.PublicKey)) + if err != nil { + t.Fatalf("Connect: %v", err) + } + defer nc.Close() + }) + + t.Run("auth enabled refuses an unenrolled nkey", func(t *testing.T) { + ref, _ := enrolledWorker(t, "worker-a") + b := startBrokerAuth(t, BrokerAuthConfig{ + Enabled: true, + Credentials: []WorkerCredentialRef{ref}, + }) + + // A freshly generated keypair that was never passed in Credentials. + strangerSeed, strangerPub, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed: %v", err) + } + if _, err := nats.Connect(b.ClientURL(), nkeyOption(t, strangerSeed, strangerPub)); err == nil { + t.Fatal("Connect: want error for unenrolled nkey, got nil") + } + }) + + t.Run("an enrolled worker cannot subscribe to another worker's traffic", func(t *testing.T) { + ref, seed := enrolledWorker(t, "worker-a") + b := startBrokerAuth(t, BrokerAuthConfig{ + Enabled: true, + Credentials: []WorkerCredentialRef{ref}, + }) + nc, err := nats.Connect( + b.ClientURL(), + nkeyOption(t, seed, ref.PublicKey), + nats.PermissionErrOnSubscribe(true), + ) + if err != nil { + t.Fatalf("Connect: %v", err) + } + defer nc.Close() + + sub, err := nc.SubscribeSync("task.status.>") + if err != nil { + t.Fatalf("SubscribeSync: %v", err) + } + if err := nc.Flush(); err != nil { + // A flush error already demonstrates the permission violation. + if !errors.Is(err, nats.ErrPermissionViolation) { + t.Fatalf("Flush: unexpected error: %v", err) + } + return + } + + if _, err := sub.NextMsg(2 * time.Second); err == nil { + t.Fatal("NextMsg: want a permissions violation, got nil error") + } else if !errors.Is(err, nats.ErrPermissionViolation) { + t.Fatalf("NextMsg: want permissions violation, got: %v", err) + } + }) +} + +func TestBrokerReloadCredentials(t *testing.T) { + t.Run("enrolling a second worker does not disturb the first", func(t *testing.T) { + refA, seedA := enrolledWorker(t, "worker-a") + b := startBrokerAuth(t, BrokerAuthConfig{ + Enabled: true, + Credentials: []WorkerCredentialRef{refA}, + }) + + ncA, err := nats.Connect(b.ClientURL(), nkeyOption(t, seedA, refA.PublicKey)) + if err != nil { + t.Fatalf("connect as A: %v", err) + } + defer ncA.Close() + + refB, seedB := enrolledWorker(t, "worker-b") + if err := b.ReloadCredentials([]WorkerCredentialRef{refA, refB}); err != nil { + t.Fatalf("ReloadCredentials: %v", err) + } + + if err := ncA.Flush(); err != nil { + t.Fatalf("A's connection unusable after reload: %v", err) + } + if !ncA.IsConnected() { + t.Fatal("A's connection was disconnected by an unrelated enrollment") + } + + ncB, err := nats.Connect(b.ClientURL(), nkeyOption(t, seedB, refB.PublicKey)) + if err != nil { + t.Fatalf("connect as newly enrolled B: %v", err) + } + defer ncB.Close() + }) + + t.Run("revoking a worker disconnects it in the reload call", func(t *testing.T) { + refA, seedA := enrolledWorker(t, "worker-a") + refB, seedB := enrolledWorker(t, "worker-b") + b := startBrokerAuth(t, BrokerAuthConfig{ + Enabled: true, + Credentials: []WorkerCredentialRef{refA, refB}, + }) + + // A does not reconnect: the point of this test is to observe the + // broker's own revocation promptly, not nats.go's reconnect/backoff + // behavior (which would otherwise mask a slow revocation behind a + // retry that happens to succeed once A is no longer welcome). + closedCh := make(chan struct{}) + ncA, err := nats.Connect( + b.ClientURL(), + nkeyOption(t, seedA, refA.PublicKey), + nats.NoReconnect(), + nats.ClosedHandler(func(*nats.Conn) { close(closedCh) }), + ) + if err != nil { + t.Fatalf("connect as A: %v", err) + } + defer ncA.Close() + + ncB, err := nats.Connect(b.ClientURL(), nkeyOption(t, seedB, refB.PublicKey)) + if err != nil { + t.Fatalf("connect as B: %v", err) + } + defer ncB.Close() + + // Revoke A by reloading with only B enrolled. Revocation is + // synchronous: nats-server re-authorizes every connected client + // inside ReloadOptions, so A is disconnected before this call + // returns — the short deadline below is to tolerate scheduling + // jitter in observing that, not because the disconnect itself is + // expected to be delayed. + if err := b.ReloadCredentials([]WorkerCredentialRef{refB}); err != nil { + t.Fatalf("ReloadCredentials: %v", err) + } + + select { + case <-closedCh: + case <-time.After(2 * time.Second): + t.Fatal("A's connection was not closed after revocation") + } + + if err := ncB.Flush(); err != nil { + t.Fatalf("B's connection unusable after A's revocation: %v", err) + } + if !ncB.IsConnected() { + t.Fatal("B's connection was disconnected by an unrelated revocation") + } + }) +} + +// TestBrokerReloadCredentials_ConcurrentWithShutdown exercises the hazard +// that ReloadCredentials being reachable from an HTTP handler goroutine +// (worker credential revocation) activates: nothing previously called +// ReloadCredentials outside a test, so it never ran concurrently with +// Shutdown, which nils ns/nc/bootOpts/serverSeed/serverPub. Run under +// -race, this fails without Broker.mu guarding those fields — either as a +// reported data race or as a nil-pointer panic when Shutdown wins the race +// and ReloadCredentials dereferences a nil *natsserver.Server. +// +// This does not assert anything about which of the two operations "wins" — +// either outcome (the reload completing before shutdown tears the server +// down, or ReloadCredentials observing "broker not started" because +// Shutdown got there first) is a correct, safe result. The property under +// test is only that neither goroutine corrupts Broker's own state or +// crashes the process while racing the other. +func TestBrokerReloadCredentials_ConcurrentWithShutdown(t *testing.T) { + ref, _ := enrolledWorker(t, "worker-a") + b := startBrokerAuth(t, BrokerAuthConfig{ + Enabled: true, + Credentials: []WorkerCredentialRef{ref}, + }) + + var wg sync.WaitGroup + wg.Add(2) + + go func() { + defer wg.Done() + for range 20 { + // Any outcome is acceptable here; "bus: broker not started" is + // expected once Shutdown has run. Only a panic or a race-detector + // report would fail this test. + _ = b.ReloadCredentials([]WorkerCredentialRef{ref}) //nolint:errcheck // outcome is intentionally unchecked; see comment above + } + }() + + go func() { + defer wg.Done() + b.Shutdown() + }() + + wg.Wait() + + // A second Shutdown (this one via t.Cleanup) must still be a safe no-op + // after the concurrent one above already ran. +} + +// TestBrokerAuth_WorkerCannotSeeAnotherWorkersLeaseReply pins the property +// the per-worker reply-inbox prefix exists for: an enrolled worker must not +// be able to read another worker's lease reply, which carries that worker's +// whole assignment batch. +// +// The reply travels over core-NATS request/reply (msg.Respond, lease.go), so +// its subject is the requester's own reply inbox and nothing else guards it +// but the subscribe permission. With a process-global "_INBOX" prefix and an +// "_INBOX.>" grant, any enrolled worker could subscribe to every other +// client's inbox on the broker and collect the OnRun command lines, embedded +// files, parameters, environment and isolation account of work it never +// leased. Each worker therefore gets its own prefix +// ([brokerauth.InboxPrefix]) and is granted only that subtree. +func TestBrokerAuth_WorkerCannotSeeAnotherWorkersLeaseReply(t *testing.T) { + refA, seedA := enrolledWorker(t, "worker-a") + refB, seedB := enrolledWorker(t, "worker-b") + b := startBrokerAuth(t, BrokerAuthConfig{ + Enabled: true, + Credentials: []WorkerCredentialRef{refA, refB}, + }) + + // Stands in for an AssignMsg batch: whatever A can read here, it could + // read of a real assignment. + const assignment = "SECRET-ASSIGNMENT-PAYLOAD" + + // The server side of the lease, on the server's own credential and the + // real SubscribeLease path. + srv, err := b.NewClient() + if err != nil { + t.Fatalf("broker NewClient: %v", err) + } + defer srv.Close() + leaseSub, err := srv.SubscribeLease(func(string, string, []byte) []byte { return []byte(assignment) }) + if err != nil { + t.Fatalf("SubscribeLease: %v", err) + } + defer leaseSub.Unsubscribe() //nolint:errcheck // best-effort test cleanup + + // Worker A, the eavesdropper, camps on both the process-global inbox + // subtree nats.go uses by default and B's own per-worker one. Neither + // subscription may be granted, so the Flush is expected to fail; the + // subscriptions are kept regardless so the assertion at the end holds + // even if a future nats.go stops reporting the violation here. + ncA, err := nats.Connect(b.ClientURL(), nkeyOption(t, seedA, refA.PublicKey)) + if err != nil { + t.Fatalf("connect as A: %v", err) + } + defer ncA.Close() + var spies []*nats.Subscription + for _, subject := range []string{"_INBOX.>", brokerauth.InboxPrefix(refB.WorkerID) + ".>"} { + spy, err := ncA.SubscribeSync(subject) + if err != nil { + t.Fatalf("A SubscribeSync %q: %v", subject, err) + } + spies = append(spies, spy) + } + if err := ncA.Flush(); err != nil && !errors.Is(err, nats.ErrPermissionViolation) { + t.Fatalf("A Flush: unexpected error: %v", err) + } + + // Worker B leases work over its own connection, with the per-worker + // inbox prefix internal/worker/natsclient gives every real worker. + clientB, err := NewClient(b.ClientURL(), slog.New(slog.DiscardHandler), + nkeyOption(t, seedB, refB.PublicKey), + nats.CustomInboxPrefix(brokerauth.InboxPrefix(refB.WorkerID))) + if err != nil { + t.Fatalf("connect as B: %v", err) + } + defer clientB.Close() + + reply, err := clientB.RequestLease(context.Background(), refB.WorkerID, "queue-1", nil, 5*time.Second) + if err != nil { + t.Fatalf("B RequestLease: %v", err) + } + if string(reply) != assignment { + t.Fatalf("B's lease reply = %q, want %q — the test cannot prove anything if B never got the payload", reply, assignment) + } + + for i, spy := range spies { + if msg, err := spy.NextMsg(time.Second); err == nil { + t.Fatalf("worker A read worker B's lease reply on spy %d (%s): %q", i, spy.Subject, msg.Data) + } else if !errors.Is(err, nats.ErrTimeout) && !errors.Is(err, nats.ErrPermissionViolation) { + t.Fatalf("spy %d (%s): unexpected error: %v", i, spy.Subject, err) + } + } +} + +// TestBrokerAuth_SkipsCredentialWithAnInvalidWorkerID is the last line of +// defense behind the two enrollment boundaries, which reject an invalid +// worker ID before it is ever stored (internal/api's enroll handler and +// sqi-server's worker enroll command). +// +// The worker ID becomes a NATS subject PATTERN in this credential's grants, +// so a stored "*" would mint one credential allowed to publish concrete +// subjects belonging to any worker, and a stored ">" would put the +// malformed "task.status.>.*" into Options.Nkeys — which nats-server may +// reject, failing every later ReloadOptions (revocation permanently 500ing) +// or refusing to boot at all. A row from a hand-edited database or an older +// binary must not be able to do either, so buildNkeys drops it and logs. +func TestBrokerAuth_SkipsCredentialWithAnInvalidWorkerID(t *testing.T) { + good, goodSeed := enrolledWorker(t, "worker-a") + + for _, workerID := range []string{"*", ">", "a.b", "a b", ""} { + t.Run("worker id "+workerID, func(t *testing.T) { + bad, badSeed := enrolledWorker(t, workerID) + b := startBrokerAuth(t, BrokerAuthConfig{ + Enabled: true, + Credentials: []WorkerCredentialRef{good, bad}, + }) + + // The valid credential is unaffected — the broker booted and + // still authorizes it. + nc, err := nats.Connect(b.ClientURL(), nkeyOption(t, goodSeed, good.PublicKey)) + if err != nil { + t.Fatalf("the valid credential was refused after an invalid one was skipped: %v", err) + } + defer nc.Close() + + // The invalid one was never installed, so its key is unknown. + if bnc, err := nats.Connect(b.ClientURL(), nkeyOption(t, badSeed, bad.PublicKey)); err == nil { + bnc.Close() + t.Errorf("a credential with worker id %q was installed; it must be skipped", workerID) + } + + // And a reload over the same set still succeeds rather than + // failing on a malformed subject pattern. + if err := b.ReloadCredentials([]WorkerCredentialRef{good, bad}); err != nil { + t.Errorf("ReloadCredentials with a skipped credential failed: %v", err) + } + }) + } +} + +// TestBrokerAuth_SkipsCredentialWithAnInvalidPublicKey is the public-key +// counterpart of TestBrokerAuth_SkipsCredentialWithAnInvalidWorkerID. +// +// Both write paths validate the public key with brokerauth.ValidatePublicKey +// before a credential is ever stored, so reaching buildNkeys with a +// malformed one needs a hand-edited database or an older binary — but +// without this guard the consequence is worse than a skipped worker: an +// invalid key handed straight to natsserver.NkeyUser makes +// natsserver.NewServer reject the options outright, which fails Start (and +// every later ReloadCredentials) for the WHOLE broker, not just the one bad +// row. One bad row must cost one worker, never the farm. +func TestBrokerAuth_SkipsCredentialWithAnInvalidPublicKey(t *testing.T) { + good, goodSeed := enrolledWorker(t, "worker-a") + bad := WorkerCredentialRef{WorkerID: "worker-b", PublicKey: "not-an-nkey"} + + b := startBrokerAuth(t, BrokerAuthConfig{ + Enabled: true, + Credentials: []WorkerCredentialRef{good, bad}, + }) + + // The valid credential is unaffected — the broker booted (it would not + // have, with the malformed key installed) and still authorizes it. + nc, err := nats.Connect(b.ClientURL(), nkeyOption(t, goodSeed, good.PublicKey)) + if err != nil { + t.Fatalf("the valid credential was refused after an invalid one was skipped: %v", err) + } + defer nc.Close() + + // And a reload over the same set still succeeds rather than failing on + // the malformed key. + if err := b.ReloadCredentials([]WorkerCredentialRef{good, bad}); err != nil { + t.Errorf("ReloadCredentials with a skipped credential failed: %v", err) + } +} diff --git a/internal/bus/broker_test.go b/internal/bus/broker_test.go index e94f89a4..cbddbcab 100644 --- a/internal/bus/broker_test.go +++ b/internal/bus/broker_test.go @@ -183,28 +183,28 @@ func TestPushConsumers(t *testing.T) { { name: "task status", publish: func(c *Client, ctx context.Context, data []byte) error { - return c.PublishTaskStatus(ctx, "job-1", data) + return c.PublishTaskStatus(ctx, "w-1", "job-1", data) }, consume: (*Client).ConsumeTaskStatus, }, { name: "task logs", publish: func(c *Client, ctx context.Context, data []byte) error { - return c.PublishTaskLog(ctx, "task-1", data) + return c.PublishTaskLog(ctx, "w-1", "task-1", data) }, consume: (*Client).ConsumeTaskLogs, }, { name: "worker heartbeat", publish: func(c *Client, ctx context.Context, data []byte) error { - return c.PublishWorkerHeartbeat(ctx, data) + return c.PublishWorkerHeartbeat(ctx, "w-1", data) }, consume: (*Client).ConsumeWorker, }, { name: "worker register", publish: func(c *Client, ctx context.Context, data []byte) error { - return c.PublishWorkerRegister(ctx, data) + return c.PublishWorkerRegister(ctx, "w-1", data) }, consume: (*Client).ConsumeWorker, }, @@ -312,7 +312,7 @@ func TestClientDrain(t *testing.T) { t.Fatalf("ConsumeTaskStatus: %v", err) } - if err := c.PublishTaskStatus(ctx, "job-d", mustJSON(t, payload{ID: "t", Body: "x"})); err != nil { + if err := c.PublishTaskStatus(ctx, "w-1", "job-d", mustJSON(t, payload{ID: "t", Body: "x"})); err != nil { t.Fatalf("PublishTaskStatus: %v", err) } select { @@ -339,7 +339,7 @@ func TestClientDrain(t *testing.T) { } // Publishing on a closed connection must fail. - if err := c.PublishTaskStatus(context.Background(), "job-d", []byte("{}")); err == nil { + if err := c.PublishTaskStatus(context.Background(), "w-1", "job-d", []byte("{}")); err == nil { t.Fatal("publish after Drain: want error, got nil") } } diff --git a/internal/bus/client.go b/internal/bus/client.go index 7cd9fdf1..4ee93fa7 100644 --- a/internal/bus/client.go +++ b/internal/bus/client.go @@ -39,18 +39,21 @@ type Client struct { // exponential backoff, so transient disruptions (e.g. a brief broker restart // during development) recover automatically. // +// extraOpts are appended after the default options, so a caller can supply +// authentication (e.g. an nkey credential) without disturbing the reconnect +// behavior above. +// // Prefer [Broker.NewClient] over calling NewClient directly; the Broker // supplies the correct loopback URL without the caller needing to know it. -func NewClient(url string, logger *slog.Logger) (*Client, error) { - nc, err := nats.Connect( - url, +func NewClient(url string, logger *slog.Logger, extraOpts ...nats.Option) (*Client, error) { + opts := []nats.Option{ // Reconnect indefinitely — this is an in-process loopback connection // and should always come back up if the embedded broker restarts. nats.MaxReconnects(-1), // Wait 2 s between reconnect attempts. The embedded broker is on // loopback so it recovers quickly; a short fixed interval is sufficient. - nats.ReconnectWait(2*time.Second), + nats.ReconnectWait(2 * time.Second), // Log reconnect events through the server's structured logger so // operators can correlate them with other activity. @@ -65,7 +68,10 @@ func NewClient(url string, logger *slog.Logger) (*Client, error) { nats.ClosedHandler(func(_ *nats.Conn) { logger.InfoContext(context.Background(), "bus: client connection closed") }), - ) + } + opts = append(opts, extraOpts...) + + nc, err := nats.Connect(url, opts...) if err != nil { return nil, fmt.Errorf("bus: client connect %q: %w", url, err) } @@ -82,31 +88,34 @@ func NewClient(url string, logger *slog.Logger) (*Client, error) { // ── Publish methods ────────────────────────────────────────────────────────── // PublishTaskStatus publishes a task-status transition to the -// task.status. subject. Workers call this when a task transitions -// to running, succeeded, failed, or canceled. -func (c *Client) PublishTaskStatus(ctx context.Context, jobID string, data []byte) error { - return c.publish(ctx, TaskStatusSubject(jobID), data) +// task.status.. subject. Workers call this when a task +// transitions to running, succeeded, failed, or canceled. +func (c *Client) PublishTaskStatus(ctx context.Context, workerID, jobID string, data []byte) error { + return c.publish(ctx, TaskStatusSubject(workerID, jobID), data) } -// PublishTaskLog publishes a log chunk to the task.logs. subject. -// Workers call this continuously as a task produces output; the server retains -// chunks in JetStream for later retrieval via the REST and WebSocket APIs. -func (c *Client) PublishTaskLog(ctx context.Context, taskID string, data []byte) error { - return c.publish(ctx, TaskLogsSubject(taskID), data) +// PublishTaskLog publishes a log chunk to the task.logs.. +// subject. Workers call this continuously as a task produces output; the +// server retains chunks in JetStream for later retrieval via the REST and +// WebSocket APIs. +func (c *Client) PublishTaskLog(ctx context.Context, workerID, taskID string, data []byte) error { + return c.publish(ctx, TaskLogsSubject(workerID, taskID), data) } -// PublishWorkerHeartbeat publishes a heartbeat ping on the worker.heartbeat -// subject. Workers call this on a regular interval; the server-side heartbeat -// sweep marks workers offline when pings stop. -func (c *Client) PublishWorkerHeartbeat(ctx context.Context, data []byte) error { - return c.publish(ctx, SubjectWorkerHeartbeat, data) +// PublishWorkerHeartbeat publishes a heartbeat ping on the +// worker.heartbeat. subject. Workers call this on a regular +// interval; the server-side heartbeat sweep marks workers offline when pings +// stop. +func (c *Client) PublishWorkerHeartbeat(ctx context.Context, workerID string, data []byte) error { + return c.publish(ctx, WorkerHeartbeatSubject(workerID), data) } // PublishWorkerRegister publishes a capability advertisement on the -// worker.register subject. Workers call this on first connect and after any -// reconnect so the server always has a current view of their capabilities. -func (c *Client) PublishWorkerRegister(ctx context.Context, data []byte) error { - return c.publish(ctx, SubjectWorkerRegister, data) +// worker.register. subject. Workers call this on first connect and +// after any reconnect so the server always has a current view of their +// capabilities. +func (c *Client) PublishWorkerRegister(ctx context.Context, workerID string, data []byte) error { + return c.publish(ctx, WorkerRegisterSubject(workerID), data) } // PublishTaskCancel publishes a task-cancellation signal to the diff --git a/internal/bus/lease.go b/internal/bus/lease.go index 9f9071bd..11c1c145 100644 --- a/internal/bus/lease.go +++ b/internal/bus/lease.go @@ -4,37 +4,45 @@ package bus import ( "context" - "strings" "time" nats "github.com/nats-io/nats.go" ) -// RequestLease sends a core-NATS work-lease request for queueID and waits up to -// timeout for the server's reply. Returns nats.ErrNoResponders immediately if -// no server is subscribed, or nats.ErrTimeout if no reply arrives in time. -func (c *Client) RequestLease(ctx context.Context, queueID string, data []byte, timeout time.Duration) ([]byte, error) { +// RequestLease sends a core-NATS work-lease request for queueID on behalf of +// workerID and waits up to timeout for the server's reply. Returns +// nats.ErrNoResponders immediately if no server is subscribed, or +// nats.ErrTimeout if no reply arrives in time. +func (c *Client) RequestLease(ctx context.Context, workerID, queueID string, data []byte, timeout time.Duration) ([]byte, error) { reqCtx, cancel := context.WithTimeout(ctx, timeout) defer cancel() - msg, err := c.nc.RequestWithContext(reqCtx, WorkLeaseSubject(queueID), data) + msg, err := c.nc.RequestWithContext(reqCtx, WorkLeaseSubject(workerID, queueID), data) if err != nil { return nil, err } return msg.Data, nil } -// SubscribeLease subscribes to work-lease requests for all queues -// (work.lease.>) and replies with the bytes handler returns. handler does its -// own long-poll blocking before returning; it must respect no work by returning -// an empty/again-marker payload of the caller's choosing. -func (c *Client) SubscribeLease(handler func(queueID string, data []byte) []byte) (*nats.Subscription, error) { +// SubscribeLease subscribes to work-lease requests from all workers for all +// queues (work.lease.>) and replies with the bytes handler returns. The +// requesting worker's ID and the queue it asked about are taken from the +// subject. handler does its own long-poll blocking before returning; it must +// respect no work by returning an empty/again-marker payload of the caller's +// choosing. +func (c *Client) SubscribeLease(handler func(workerID, queueID string, data []byte) []byte) (*nats.Subscription, error) { // Spawn a goroutine per message so a parked handler (a long-poll lease that // blocks up to leaseHoldTimeout) never stalls delivery of other workers' // requests — NATS delivers one-at-a-time per subscription callback. sub, err := c.nc.Subscribe(SubjectWorkLeasePrefix+".>", func(msg *nats.Msg) { go func() { - queueID := strings.TrimPrefix(msg.Subject, SubjectWorkLeasePrefix+".") - reply := handler(queueID, msg.Data) + workerID, queueID, ok := ParseWorkerSubject(msg.Subject) + if !ok { + // A subject shape this server does not speak. Do not reply: + // the requester times out and retries, which is the same path + // an unreachable server produces. + return + } + reply := handler(workerID, queueID, msg.Data) _ = msg.Respond(reply) //nolint:errcheck // best-effort reply; worker retries on timeout }() }) diff --git a/internal/bus/lease_test.go b/internal/bus/lease_test.go index cb0f4018..269b1cf1 100644 --- a/internal/bus/lease_test.go +++ b/internal/bus/lease_test.go @@ -13,7 +13,10 @@ func TestLeaseRequestReply(t *testing.T) { server := newClient(t, b) worker := newClient(t, b) - sub, err := server.SubscribeLease(func(queueID string, data []byte) []byte { + sub, err := server.SubscribeLease(func(workerID, queueID string, data []byte) []byte { + if workerID != "w1" { + t.Errorf("workerID = %q, want w1", workerID) + } if queueID != "q1" { t.Errorf("queueID = %q, want q1", queueID) } @@ -24,24 +27,25 @@ func TestLeaseRequestReply(t *testing.T) { } t.Cleanup(func() { _ = sub.Unsubscribe() }) //nolint:errcheck // test cleanup - reply, err := worker.RequestLease(context.Background(), "q1", []byte("w1"), 2*time.Second) + reply, err := worker.RequestLease(context.Background(), "w1", "q1", []byte("payload"), 2*time.Second) if err != nil { t.Fatalf("RequestLease: %v", err) } - if string(reply) != "reply-to-w1" { - t.Errorf("reply = %q, want reply-to-w1", reply) + if string(reply) != "reply-to-payload" { + t.Errorf("reply = %q, want reply-to-payload", reply) } } // TestLeaseRequestReply_WildcardTokenRoutes guards the queueless-worker // regression: a request on the wildcard token reaches the server's work.lease.> -// subscription, whereas an empty leaf ("work.lease.") routes to no responder. +// subscription, whereas an empty queue token ("work.lease.w1.") routes to no +// responder. func TestLeaseRequestReply_WildcardTokenRoutes(t *testing.T) { b := startBroker(t) server := newClient(t, b) worker := newClient(t, b) - sub, err := server.SubscribeLease(func(queueID string, _ []byte) []byte { + sub, err := server.SubscribeLease(func(_, queueID string, _ []byte) []byte { return []byte("got:" + queueID) }) if err != nil { @@ -50,7 +54,7 @@ func TestLeaseRequestReply_WildcardTokenRoutes(t *testing.T) { t.Cleanup(func() { _ = sub.Unsubscribe() }) //nolint:errcheck // test cleanup // Wildcard token routes to the server. - reply, err := worker.RequestLease(context.Background(), WildcardQueueToken, []byte("w1"), 2*time.Second) + reply, err := worker.RequestLease(context.Background(), "w1", WildcardQueueToken, []byte("payload"), 2*time.Second) if err != nil { t.Fatalf("RequestLease(wildcard): %v", err) } @@ -58,8 +62,33 @@ func TestLeaseRequestReply_WildcardTokenRoutes(t *testing.T) { t.Errorf("wildcard reply = %q, want got:%s", reply, WildcardQueueToken) } - // Empty leaf does NOT route — this is the bug the wildcard token fixes. - if _, err := worker.RequestLease(context.Background(), "", []byte("w1"), 300*time.Millisecond); err == nil { - t.Error("empty-queue RequestLease unexpectedly succeeded; an empty leaf must not route to the server") + // Empty queue token does NOT route — this is the bug the wildcard token fixes. + if _, err := worker.RequestLease(context.Background(), "w1", "", []byte("payload"), 300*time.Millisecond); err == nil { + t.Error("empty-queue RequestLease unexpectedly succeeded; an empty queue token must not route to the server") + } +} + +// TestSubscribeLease_IgnoresUnparsableSubject pins the no-reply path: a request +// arriving on a subject that carries no worker identity gets no response at +// all, so the requester's own timeout is what ends the exchange. +func TestSubscribeLease_IgnoresUnparsableSubject(t *testing.T) { + b := startBroker(t) + server := newClient(t, b) + worker := newClient(t, b) + + sub, err := server.SubscribeLease(func(workerID, queueID string, _ []byte) []byte { + t.Errorf("handler called for an identity-less subject: (%q, %q)", workerID, queueID) + return []byte("unexpected") + }) + if err != nil { + t.Fatalf("SubscribeLease: %v", err) + } + t.Cleanup(func() { _ = sub.Unsubscribe() }) //nolint:errcheck // test cleanup + + // The pre-identity subject shape: a queue token with no worker before it. + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if _, err := worker.nc.RequestWithContext(ctx, SubjectWorkLeasePrefix+".q1", []byte("payload")); err == nil { + t.Fatal("request on an identity-less subject was answered; want a timeout") } } diff --git a/internal/bus/streams.go b/internal/bus/streams.go index dc517ae0..13265a41 100644 --- a/internal/bus/streams.go +++ b/internal/bus/streams.go @@ -42,7 +42,7 @@ const ( func streamDefs() []jetstream.StreamConfig { return []jetstream.StreamConfig{ { - // SQI_TASK — task.status. + // SQI_TASK — task.status.. // // Workers publish task-state transitions (running, succeeded, failed, // canceled) to this stream keyed by job ID. The server's status @@ -58,7 +58,7 @@ func streamDefs() []jetstream.StreamConfig { Replicas: 1, }, { - // SQI_LOGS — task.logs. + // SQI_LOGS — task.logs.. // // Workers publish structured log chunks as tasks run. LimitsPolicy // retains all chunks so the REST and WebSocket log-tail endpoints can @@ -78,7 +78,8 @@ func streamDefs() []jetstream.StreamConfig { Replicas: 1, }, { - // SQI_WORKER — worker.register, worker.heartbeat, worker.deregister + // SQI_WORKER — worker.register., worker.heartbeat., + // worker.deregister. // // Workers publish registration payloads on connect and reconnect, // heartbeat pings on a configurable interval, and a departure message @@ -93,12 +94,16 @@ func streamDefs() []jetstream.StreamConfig { // once processed. Name: StreamWorker, Description: "Worker registration and heartbeat messages.", - Subjects: []string{SubjectWorkerRegister, SubjectWorkerHeartbeat, SubjectWorkerDeregister}, - Retention: jetstream.WorkQueuePolicy, - Storage: jetstream.FileStorage, - MaxAge: 2 * time.Minute, - Discard: jetstream.DiscardOld, - Replicas: 1, + Subjects: []string{ + SubjectWorkerRegisterPrefix + ".>", + SubjectWorkerHeartbeatPrefix + ".>", + SubjectWorkerDeregisterPrefix + ".>", + }, + Retention: jetstream.WorkQueuePolicy, + Storage: jetstream.FileStorage, + MaxAge: 2 * time.Minute, + Discard: jetstream.DiscardOld, + Replicas: 1, }, { // SQI_CANCEL — task.cancel. diff --git a/internal/bus/streams_test.go b/internal/bus/streams_test.go index dbc1c99f..aa806a33 100644 --- a/internal/bus/streams_test.go +++ b/internal/bus/streams_test.go @@ -23,9 +23,13 @@ func TestStreamDefs(t *testing.T) { subjects []string retention jetstream.RetentionPolicy }{ - StreamTask: {[]string{SubjectTaskStatusPrefix + ".>"}, jetstream.WorkQueuePolicy}, - StreamLogs: {[]string{SubjectTaskLogsPrefix + ".>"}, jetstream.LimitsPolicy}, - StreamWorker: {[]string{SubjectWorkerRegister, SubjectWorkerHeartbeat, SubjectWorkerDeregister}, jetstream.WorkQueuePolicy}, + StreamTask: {[]string{SubjectTaskStatusPrefix + ".>"}, jetstream.WorkQueuePolicy}, + StreamLogs: {[]string{SubjectTaskLogsPrefix + ".>"}, jetstream.LimitsPolicy}, + StreamWorker: {[]string{ + SubjectWorkerRegisterPrefix + ".>", + SubjectWorkerHeartbeatPrefix + ".>", + SubjectWorkerDeregisterPrefix + ".>", + }, jetstream.WorkQueuePolicy}, StreamCancel: {[]string{SubjectTaskCancelPrefix + ".>"}, jetstream.WorkQueuePolicy}, } diff --git a/internal/bus/subjectgrants_test.go b/internal/bus/subjectgrants_test.go new file mode 100644 index 00000000..7b5b6fa1 --- /dev/null +++ b/internal/bus/subjectgrants_test.go @@ -0,0 +1,110 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package bus + +import ( + "strings" + "testing" + + natsserver "github.com/nats-io/nats-server/v2/server" + + "github.com/uberware/sqi/internal/brokerauth" +) + +// permAllows reports whether subject matches an entry in sp.Allow, honoring the +// two NATS subject wildcards: "*" matches exactly one token, and a trailing ">" +// matches one or more trailing tokens. +// +// This is a small local reimplementation rather than a shared helper: the point +// of the test below is to check one package's output against another package's +// data, so borrowing either package's own matcher would let a shared mistake +// agree with itself. +func permAllows(sp *natsserver.SubjectPermission, subject string) bool { + if sp == nil { + return false + } + for _, pattern := range sp.Allow { + if permMatches(pattern, subject) { + return true + } + } + return false +} + +func permMatches(pattern, subject string) bool { + patternTokens := strings.Split(pattern, ".") + subjectTokens := strings.Split(subject, ".") + + for i, pt := range patternTokens { + if pt == ">" { + // ">" is only valid as the final token and matches one or more + // remaining subject tokens. + return i == len(patternTokens)-1 && i < len(subjectTokens) + } + if i >= len(subjectTokens) { + return false + } + if pt == "*" { + continue + } + if pt != subjectTokens[i] { + return false + } + } + return len(patternTokens) == len(subjectTokens) +} + +// TestSubjectHelpers_MatchWorkerPermissions binds this package's subject +// construction to the broker's per-worker publish grants. +// +// The two are written independently — [brokerauth.WorkerPermissions] builds its +// patterns by string concatenation and knows nothing about the helpers here — +// so nothing but this test stops a reordered or renamed token in subjects.go +// from silently drifting out of the grant it is supposed to fall under. Such a +// drift compiles, passes every other test, and surfaces only as a runtime +// authorization denial on a farm that has broker auth switched on, which is the +// exact failure the identity-carrying subject scheme exists to make impossible. +func TestSubjectHelpers_MatchWorkerPermissions(t *testing.T) { + const me = "worker-a" + const other = "worker-b" + + perms := brokerauth.WorkerPermissions(me) + + tests := []struct { + name string + // mine is the subject this worker publishes; theirs is the same + // subject class published by a different worker. + mine string + theirs string + }{ + {"task status", TaskStatusSubject(me, "job-1"), TaskStatusSubject(other, "job-1")}, + {"task logs", TaskLogsSubject(me, "task-1"), TaskLogsSubject(other, "task-1")}, + {"worker register", WorkerRegisterSubject(me), WorkerRegisterSubject(other)}, + {"worker heartbeat", WorkerHeartbeatSubject(me), WorkerHeartbeatSubject(other)}, + {"worker deregister", WorkerDeregisterSubject(me), WorkerDeregisterSubject(other)}, + {"worker diag", WorkerDiagSubject(me), WorkerDiagSubject(other)}, + {"work lease", WorkLeaseSubject(me, "queue-1"), WorkLeaseSubject(other, "queue-1")}, + {"work lease wildcard queue", WorkLeaseSubject(me, WildcardQueueToken), WorkLeaseSubject(other, WildcardQueueToken)}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if !permAllows(perms.Publish, tt.mine) { + t.Errorf( + "subject %q is not granted by WorkerPermissions(%q) — the helper and the grant have drifted apart; allow list: %v", + tt.mine, me, perms.Publish.Allow, + ) + } + // The negative half is what makes the positive half mean + // something: a permission set of ">" would satisfy the check + // above for every subject, and confining a worker to its own + // traffic is the whole purpose of the scheme. + if permAllows(perms.Publish, tt.theirs) { + t.Errorf( + "subject %q is granted by WorkerPermissions(%q) — a worker can publish as another worker", + tt.theirs, me, + ) + } + }) + } +} diff --git a/internal/bus/subjects.go b/internal/bus/subjects.go index d86a0e74..ec3ecbe3 100644 --- a/internal/bus/subjects.go +++ b/internal/bus/subjects.go @@ -8,74 +8,92 @@ // Each subject class follows a hierarchical naming scheme so that JetStream // stream subject-filters and consumer subject-bindings are intuitive: // -// work.lease. — worker → server: request/reply work-lease batch -// task.status. — worker → server: task state transitions -// task.logs. — worker → server: log chunk ingestion -// worker.heartbeat — worker → server: periodic liveness pings -// worker.register — worker → server: capability advertisement -// worker.deregister — worker → server: graceful departure notification +// work.lease.. — worker → server: request/reply work-lease batch +// task.status.. — worker → server: task state transitions +// task.logs.. — worker → server: log chunk ingestion +// worker.heartbeat. — worker → server: periodic liveness pings +// worker.register. — worker → server: capability advertisement +// worker.deregister. — worker → server: graceful departure notification +// task.cancel. — server → worker: interrupt a running task // -// The leaf token (, , ) is the opaque string identifier of -// the corresponding entity in the SQLite store. Callers build full subject -// strings using the helper functions below rather than constructing them by -// hand. +// Every worker → server subject carries the publishing worker's ID directly +// after its class prefix. That placement is what makes the broker's per-worker +// publish permissions expressible: NATS permissions are static per credential +// and JetStream does not stamp publisher identity onto a message, so a scheme +// keyed only by job, task and queue would offer no way to say "only this +// worker's own traffic" — and no way for the server to learn who published a +// message it received. [ParseWorkerSubject] recovers the identity on the +// receiving side. +// +// The remaining tokens (, , ) are the opaque string +// identifiers of the corresponding entity in the SQLite store. Callers build +// full subject strings using the helper functions below rather than +// constructing them by hand. package bus -// Subject prefix and fixed-subject constants. +import "strings" + +// Subject prefix constants. const ( // SubjectTaskStatusPrefix is the prefix for task-status subjects. - // Full subject: SubjectTaskStatusPrefix + "." + jobID. + // Full subject: SubjectTaskStatusPrefix + "." + workerID + "." + jobID. SubjectTaskStatusPrefix = "task.status" // SubjectTaskLogsPrefix is the prefix for task-log subjects. - // Full subject: SubjectTaskLogsPrefix + "." + taskID. + // Full subject: SubjectTaskLogsPrefix + "." + workerID + "." + taskID. SubjectTaskLogsPrefix = "task.logs" // SubjectTaskCancelPrefix is the prefix for task-cancellation subjects. // Full subject: SubjectTaskCancelPrefix + "." + taskID. // The server publishes to this subject; the worker assigned to the task - // consumes it and interrupts the running process. + // consumes it and interrupts the running process. This is the one subject + // class that travels server → worker, so it carries no worker token. SubjectTaskCancelPrefix = "task.cancel" - // SubjectWorkerHeartbeat is the subject workers publish liveness pings to. - SubjectWorkerHeartbeat = "worker.heartbeat" + // SubjectWorkerHeartbeatPrefix is the prefix for worker liveness pings. + // Full subject: SubjectWorkerHeartbeatPrefix + "." + workerID. + SubjectWorkerHeartbeatPrefix = "worker.heartbeat" - // SubjectWorkerRegister is the subject workers publish registration - // messages to when they first connect or reconnect. - SubjectWorkerRegister = "worker.register" + // SubjectWorkerRegisterPrefix is the prefix workers publish registration + // messages under when they first connect or reconnect. + // Full subject: SubjectWorkerRegisterPrefix + "." + workerID. + SubjectWorkerRegisterPrefix = "worker.register" - // SubjectWorkerDeregister is the subject workers publish to on graceful - // shutdown so the server can mark the worker offline immediately rather - // than waiting for heartbeat timeout. The server handler for this subject - // calls [store.WorkerStore.UpdateWorkerStatus] with WorkerStatusOffline. - SubjectWorkerDeregister = "worker.deregister" + // SubjectWorkerDeregisterPrefix is the prefix workers publish under on + // graceful shutdown so the server can mark the worker offline immediately + // rather than waiting for heartbeat timeout. The server handler for these + // subjects calls [store.WorkerStore.UpdateWorkerStatus] with + // WorkerStatusOffline. + // Full subject: SubjectWorkerDeregisterPrefix + "." + workerID. + SubjectWorkerDeregisterPrefix = "worker.deregister" // SubjectWorkLeasePrefix is the prefix for worker work-lease requests. - // Full subject: SubjectWorkLeasePrefix + "." + queueID. Core NATS - // request/reply — workers ask for work; the server replies with a batch. + // Full subject: SubjectWorkLeasePrefix + "." + workerID + "." + queueID. + // Core NATS request/reply — workers ask for work; the server replies with + // a batch. SubjectWorkLeasePrefix = "work.lease" - // WildcardQueueToken is the lease-subject leaf a queue-unaffiliated worker - // (empty QueueIDs — "serve any queue") uses in place of a real queue ID, so - // it requests on a valid subject (work.lease._any) that the server's - // work.lease.> subscription actually receives. An empty leaf would produce - // the invalid subject "work.lease." (no responders). The leaf is reserved - // (underscore prefix) so it cannot collide with a real UUID queue ID; the - // server selects tasks farm-wide and gates by worker eligibility, so the - // token's only role is subject routing and wake-up bucketing. + // WildcardQueueToken is the lease-subject queue token a queue-unaffiliated + // worker (empty QueueIDs — "serve any queue") uses in place of a real queue + // ID, so it requests on a valid subject (work.lease.._any) that the + // server's work.lease.> subscription actually receives. An empty token would + // produce an unroutable subject. The token is reserved (underscore prefix) + // so it cannot collide with a real UUID queue ID; the server selects tasks + // farm-wide and gates by worker eligibility, so the token's only role is + // subject routing and wake-up bucketing. WildcardQueueToken = "_any" ) // TaskStatusSubject returns the full NATS subject for task-status messages -// belonging to the given job. -func TaskStatusSubject(jobID string) string { - return SubjectTaskStatusPrefix + "." + jobID +// published by the given worker about the given job. +func TaskStatusSubject(workerID, jobID string) string { + return SubjectTaskStatusPrefix + "." + workerID + "." + jobID } // TaskLogsSubject returns the full NATS subject for log-chunk messages -// belonging to the given task. -func TaskLogsSubject(taskID string) string { - return SubjectTaskLogsPrefix + "." + taskID +// published by the given worker for the given task. +func TaskLogsSubject(workerID, taskID string) string { + return SubjectTaskLogsPrefix + "." + workerID + "." + taskID } // TaskCancelSubject returns the full NATS subject for a task-cancellation @@ -84,8 +102,78 @@ func TaskCancelSubject(taskID string) string { return SubjectTaskCancelPrefix + "." + taskID } -// WorkLeaseSubject returns the full NATS subject a worker requests work on for -// the given queue. -func WorkLeaseSubject(queueID string) string { - return SubjectWorkLeasePrefix + "." + queueID +// WorkerRegisterSubject returns the full NATS subject the given worker +// publishes its capability advertisement to. +func WorkerRegisterSubject(workerID string) string { + return SubjectWorkerRegisterPrefix + "." + workerID +} + +// WorkerHeartbeatSubject returns the full NATS subject the given worker +// publishes its liveness pings to. +func WorkerHeartbeatSubject(workerID string) string { + return SubjectWorkerHeartbeatPrefix + "." + workerID +} + +// WorkerDeregisterSubject returns the full NATS subject the given worker +// publishes its graceful-departure notification to. +func WorkerDeregisterSubject(workerID string) string { + return SubjectWorkerDeregisterPrefix + "." + workerID +} + +// WorkLeaseSubject returns the full NATS subject the given worker requests +// work on for the given queue. +func WorkLeaseSubject(workerID, queueID string) string { + return SubjectWorkLeasePrefix + "." + workerID + "." + queueID +} + +// ParseWorkerSubject splits a worker → server subject into the ID of the worker +// that published it and the trailing token identifying the entity it concerns +// (the job for task status, the task for logs, the queue for a lease request; +// empty for the three worker-lifecycle subjects, which have no trailing token). +// +// ok is false for anything that is not one of the six worker → server subject +// shapes, including subjects with an empty token and subjects whose worker or +// trailing token carries a NATS wildcard. A message whose subject does not +// carry one concrete identity is one the server cannot attribute to a worker, +// and callers must reject it rather than act on a partial parse. +func ParseWorkerSubject(subject string) (workerID, leaf string, ok bool) { + tokens := strings.Split(subject, ".") + if len(tokens) < 3 || len(tokens) > 4 { + return "", "", false + } + prefix := tokens[0] + "." + tokens[1] + + if len(tokens) == 3 { + switch prefix { + case SubjectWorkerRegisterPrefix, SubjectWorkerHeartbeatPrefix, SubjectWorkerDeregisterPrefix: + if !concreteToken(tokens[2]) { + return "", "", false + } + return tokens[2], "", true + default: + return "", "", false + } + } + + switch prefix { + case SubjectTaskStatusPrefix, SubjectTaskLogsPrefix, SubjectWorkLeasePrefix: + if !concreteToken(tokens[2]) || !concreteToken(tokens[3]) { + return "", "", false + } + return tokens[2], tokens[3], true + default: + return "", "", false + } +} + +// concreteToken reports whether tok is a single, literal subject token: not +// empty, and free of the NATS wildcards "*" and ">". +// +// A wildcard names a set of workers rather than one, so a subject carrying one +// identifies nobody. nats-server refuses to publish on a wildcard subject, but +// a parser that hands back "*" as a worker ID is relying on that refusal for +// its own soundness — and callers treat this return value as an authorization +// input, which must not depend on a check happening somewhere else. +func concreteToken(tok string) bool { + return tok != "" && !strings.ContainsAny(tok, "*>") } diff --git a/internal/bus/subjects_test.go b/internal/bus/subjects_test.go index cce59857..e63666f7 100644 --- a/internal/bus/subjects_test.go +++ b/internal/bus/subjects_test.go @@ -2,7 +2,10 @@ package bus -import "testing" +import ( + "strings" + "testing" +) func TestSubjectHelpers(t *testing.T) { tests := []struct { @@ -10,10 +13,14 @@ func TestSubjectHelpers(t *testing.T) { got string want string }{ - {"work lease", WorkLeaseSubject("q-123"), "work.lease.q-123"}, - {"task status", TaskStatusSubject("job-abc"), "task.status.job-abc"}, - {"task logs", TaskLogsSubject("task-xyz"), "task.logs.task-xyz"}, + {"work lease", WorkLeaseSubject("w-1", "q-123"), "work.lease.w-1.q-123"}, + {"work lease wildcard queue", WorkLeaseSubject("w-1", WildcardQueueToken), "work.lease.w-1._any"}, + {"task status", TaskStatusSubject("w-1", "job-abc"), "task.status.w-1.job-abc"}, + {"task logs", TaskLogsSubject("w-1", "task-xyz"), "task.logs.w-1.task-xyz"}, {"task cancel", TaskCancelSubject("task-xyz"), "task.cancel.task-xyz"}, + {"worker register", WorkerRegisterSubject("w-1"), "worker.register.w-1"}, + {"worker heartbeat", WorkerHeartbeatSubject("w-1"), "worker.heartbeat.w-1"}, + {"worker deregister", WorkerDeregisterSubject("w-1"), "worker.deregister.w-1"}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { @@ -32,17 +39,73 @@ func TestSubjectPrefixConstants(t *testing.T) { prefix string full string }{ - {"work lease", SubjectWorkLeasePrefix, WorkLeaseSubject("x")}, - {"task status", SubjectTaskStatusPrefix, TaskStatusSubject("x")}, - {"task logs", SubjectTaskLogsPrefix, TaskLogsSubject("x")}, + {"work lease", SubjectWorkLeasePrefix, WorkLeaseSubject("x", "y")}, + {"task status", SubjectTaskStatusPrefix, TaskStatusSubject("x", "y")}, + {"task logs", SubjectTaskLogsPrefix, TaskLogsSubject("x", "y")}, {"task cancel", SubjectTaskCancelPrefix, TaskCancelSubject("x")}, + {"worker register", SubjectWorkerRegisterPrefix, WorkerRegisterSubject("x")}, + {"worker heartbeat", SubjectWorkerHeartbeatPrefix, WorkerHeartbeatSubject("x")}, + {"worker deregister", SubjectWorkerDeregisterPrefix, WorkerDeregisterSubject("x")}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - want := tt.prefix + ".x" - if tt.full != want { + want := tt.prefix + "." + if !strings.HasPrefix(tt.full, want) || tt.full == want { t.Fatalf("full subject %q does not compose from prefix %q", tt.full, tt.prefix) } }) } } + +func TestParseWorkerSubject(t *testing.T) { + tests := []struct { + subject string + wantWorker string + wantLeaf string + wantOK bool + }{ + {"task.status.w1.j1", "w1", "j1", true}, + {"task.logs.w1.t1", "w1", "t1", true}, + {"worker.register.w1", "w1", "", true}, + {"worker.heartbeat.w1", "w1", "", true}, + {"worker.deregister.w1", "w1", "", true}, + {"work.lease.w1.q1", "w1", "q1", true}, + {"work.lease.w1._any", "w1", "_any", true}, + + // Subject shapes that do not carry a worker identity. A publisher using + // one of these is not a worker this server can attribute a message to, + // so it must be rejected rather than parsed into a partial answer. + {"task.status.j1", "", "", false}, + {"worker.register", "", "", false}, + {"task.cancel.t1", "", "", false}, // server → worker + {"worker.diag.w1", "", "", false}, // core NATS, not a stream subject + {"task.status.w1.j1.extra", "", "", false}, + {"task.status..j1", "", "", false}, + {"task.status.w1.", "", "", false}, + {"worker.register.", "", "", false}, + {"", "", "", false}, + + // A wildcard token names a set of workers, not one, so it identifies + // nobody. Callers feed this worker ID to an authorization decision, so + // the parser must reject it here rather than lean on the broker's + // refusal to publish on a wildcard subject. + {"task.status.*.j1", "", "", false}, + {"task.status.w1.*", "", "", false}, + {"task.status.>.j1", "", "", false}, + {"task.status.w1.>", "", "", false}, + {"worker.register.*", "", "", false}, + {"worker.heartbeat.>", "", "", false}, + {"work.lease.*.q1", "", "", false}, + {"work.lease.w1.>", "", "", false}, + {"task.logs.w*1.t1", "", "", false}, + {"worker.deregister.w>1", "", "", false}, + } + for _, tt := range tests { + t.Run(tt.subject, func(t *testing.T) { + w, leaf, ok := ParseWorkerSubject(tt.subject) + if w != tt.wantWorker || leaf != tt.wantLeaf || ok != tt.wantOK { + t.Errorf("= (%q, %q, %v), want (%q, %q, %v)", w, leaf, ok, tt.wantWorker, tt.wantLeaf, tt.wantOK) + } + }) + } +} diff --git a/internal/config/config.go b/internal/config/config.go index 4b259f3d..68b4a3b9 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -80,8 +80,52 @@ type NATSConfig struct { // MaxStoreMB is the JetStream file-storage cap in megabytes. // Env: SQI_NATS_MAX_STORE_MB MaxStoreMB int `yaml:"max_store_mb"` + + // Auth configures per-worker broker authentication. Off by default. + // Deliberately INDEPENDENT of the top-level auth block: the two protect + // different surfaces, and coupling them would force an operator who + // wants worker authentication into user accounts they did not ask for. + Auth NATSAuthConfig `yaml:"auth"` +} + +// NATSAuthConfig configures broker authentication and worker enrollment. +type NATSAuthConfig struct { + // Enabled requires every NATS client to present a per-worker nkey + // credential. When false the broker accepts any connection, which is the + // v0.3.0 behavior and the default. + // Env: SQI_NATS_AUTH_ENABLED + Enabled bool `yaml:"enabled"` + + // JoinTokenTTL is how long a newly issued worker join token remains + // valid. Bounded by MinNATSAuthJoinTokenTTL and MaxNATSAuthJoinTokenTTL. + // Env: SQI_NATS_AUTH_JOIN_TOKEN_TTL + JoinTokenTTL time.Duration `yaml:"join_token_ttl"` + + // JoinTokenSingleUse consumes a join token on first successful + // enrollment. Leaving it true is strongly recommended; false exists for + // image-baked fleets that enroll many identical machines from one token. + // Env: SQI_NATS_AUTH_JOIN_TOKEN_SINGLE_USE + JoinTokenSingleUse bool `yaml:"join_token_single_use"` + + // EnrollmentEndpointEnabled mounts POST /api/v1/workers/enroll. Set it + // false at a site that provisions every credential by hand and wants no + // enrollment surface at all. Meaningful only when Enabled is true. + // Env: SQI_NATS_AUTH_ENROLLMENT_ENDPOINT_ENABLED + EnrollmentEndpointEnabled bool `yaml:"enrollment_endpoint_enabled"` } +const ( + // MinNATSAuthJoinTokenTTL is the floor: below a minute an operator + // cannot realistically get the token onto a machine and boot it. + MinNATSAuthJoinTokenTTL = 1 * time.Minute + + // MaxNATSAuthJoinTokenTTL is the ceiling. A join token mints a worker + // credential, so it is a bootstrap secret whose whole value is a short + // blast radius; a token valid for weeks is a standing credential wearing + // a different name. + MaxNATSAuthJoinTokenTTL = 24 * time.Hour +) + // StoreConfig controls the embedded SQLite state store. type StoreConfig struct { // SQLitePath is the path to the SQLite database file. @@ -698,6 +742,12 @@ func DefaultConfig() Config { Addr: "0.0.0.0:4222", DataDir: "data/nats", MaxStoreMB: 1024, + Auth: NATSAuthConfig{ + Enabled: false, + JoinTokenTTL: 1 * time.Hour, + JoinTokenSingleUse: true, + EnrollmentEndpointEnabled: true, + }, }, Store: StoreConfig{ SQLitePath: "sqi.db", diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 50021ab5..1d696cea 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -2407,3 +2407,39 @@ func TestOIDCConfig_MarshalYAMLRedactsSecret(t *testing.T) { t.Fatalf("client_secret leaked into YAML output:\n%s", out) } } + +func TestValidate_NATSAuthJoinTokenTTL(t *testing.T) { + tests := []struct { + name string + enabled bool + ttl time.Duration + wantErr bool + }{ + {"disabled block never errors", false, 0, false}, + {"disabled block ignores absurd ttl", false, 400 * 24 * time.Hour, false}, + {"below floor", true, config.MinNATSAuthJoinTokenTTL - time.Second, true}, + {"at floor", true, config.MinNATSAuthJoinTokenTTL, false}, + {"at ceiling", true, config.MaxNATSAuthJoinTokenTTL, false}, + {"above ceiling", true, config.MaxNATSAuthJoinTokenTTL + time.Second, true}, + {"zero when enabled", true, 0, true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg := config.DefaultConfig() + cfg.NATS.Auth.Enabled = tt.enabled + cfg.NATS.Auth.JoinTokenTTL = tt.ttl + + errs := config.Validate(cfg) + + var found bool + for _, e := range errs { + if e.Field == "nats.auth.join_token_ttl" { + found = true + } + } + if found != tt.wantErr { + t.Errorf("error on nats.auth.join_token_ttl = %v, want %v (errs: %v)", found, tt.wantErr, errs) + } + }) + } +} diff --git a/internal/config/loader.go b/internal/config/loader.go index 027720ac..a2867522 100644 --- a/internal/config/loader.go +++ b/internal/config/loader.go @@ -81,22 +81,58 @@ func defaultSearchPaths() []string { // // A missing config file is not an error unless filePath was set explicitly. func Load(filePath string, flags FlagOverrides) (Config, error) { + cfg, _, err := LoadWithSources(filePath, flags) + return cfg, err +} + +// Sources reports, for a subset of settings, whether the config file or +// environment layer explicitly decided the value — as opposed to it being +// the built-in default that nothing overrode. +// +// Comparing a resolved value against [DefaultConfig]'s is not a sound way to +// answer this: an operator's config file (or the shipped +// config/sqi-server.example.yaml) can restate a default value's exact text +// while editing other keys, in which case value comparison cannot tell +// "explicitly configured to the default" apart from "never touched, so it +// is still the default". Sources is built from the file/env layers +// themselves, before any default-fill happens, so it is not fooled by that. +type Sources struct { + // StoreSQLitePath is true when store.sqlite_path was set by the config + // file or by SQI_STORE_SQLITE_PATH. + StoreSQLitePath bool +} + +// LoadWithSources is [Load], additionally reporting which of a subset of +// settings (see [Sources]) were explicitly decided by the config file or +// environment layer. +func LoadWithSources(filePath string, flags FlagOverrides) (Config, Sources, error) { cfg := DefaultConfig() + var src Sources // ── Layer 2: config file ────────────────────────────────────────────── - if err := applyFile(&cfg, filePath); err != nil { - return Config{}, err + fc, err := loadFileConfig(filePath) + if err != nil { + return Config{}, Sources{}, err + } + if fc != nil { + mergeFileConfig(&cfg, *fc) + if fc.Store != nil && fc.Store.SQLitePath != nil { + src.StoreSQLitePath = true + } } // ── Layer 3: environment variables ─────────────────────────────────── if err := applyEnv(&cfg); err != nil { - return Config{}, err + return Config{}, Sources{}, err + } + if os.Getenv("SQI_STORE_SQLITE_PATH") != "" { + src.StoreSQLitePath = true } // ── Layer 4: CLI flag overrides ─────────────────────────────────────── applyFlags(&cfg, flags) - return cfg, nil + return cfg, src, nil } // ── File layer ──────────────────────────────────────────────────────────────── @@ -122,6 +158,12 @@ type fileConfig struct { Addr *string `yaml:"addr"` DataDir *string `yaml:"data_dir"` MaxStoreMB *int `yaml:"max_store_mb"` + Auth *struct { + Enabled *bool `yaml:"enabled"` + JoinTokenTTL *string `yaml:"join_token_ttl"` + JoinTokenSingleUse *bool `yaml:"join_token_single_use"` + EnrollmentEndpointEnabled *bool `yaml:"enrollment_endpoint_enabled"` + } `yaml:"auth"` } `yaml:"nats"` Store *struct { @@ -222,27 +264,29 @@ type fileConfig struct { } `yaml:"auth"` } -func applyFile(cfg *Config, explicit string) error { +// loadFileConfig resolves and parses the config file, returning a nil +// *fileConfig (not an error) when none is found and filePath was not set +// explicitly. +func loadFileConfig(explicit string) (*fileConfig, error) { path, err := resolveFilePath(explicit) if err != nil { - return err + return nil, err } if path == "" { - return nil // no file found; not an error + return nil, nil } data, err := os.ReadFile(path) if err != nil { - return fmt.Errorf("read config file %q: %w", path, err) + return nil, fmt.Errorf("read config file %q: %w", path, err) } var fc fileConfig if err := yaml.Unmarshal(data, &fc); err != nil { - return fmt.Errorf("parse config file %q: %w", path, err) + return nil, fmt.Errorf("parse config file %q: %w", path, err) } - mergeFileConfig(cfg, fc) - return nil + return &fc, nil } // resolveFilePath returns the path to use for file loading. @@ -306,6 +350,36 @@ func mergeNATSFile(cfg *Config, fc fileConfig) { if fc.NATS.MaxStoreMB != nil { cfg.NATS.MaxStoreMB = *fc.NATS.MaxStoreMB } + mergeNATSAuthFile(cfg, fc.NATS.Auth) +} + +// mergeNATSAuthFile overlays the nats.auth sub-fields from fc onto cfg. Split +// out of [mergeNATSFile] to keep its cyclomatic complexity under the lint +// threshold. +func mergeNATSAuthFile(cfg *Config, a *struct { + Enabled *bool `yaml:"enabled"` + JoinTokenTTL *string `yaml:"join_token_ttl"` + JoinTokenSingleUse *bool `yaml:"join_token_single_use"` + EnrollmentEndpointEnabled *bool `yaml:"enrollment_endpoint_enabled"` +}, +) { + if a == nil { + return + } + if a.Enabled != nil { + cfg.NATS.Auth.Enabled = *a.Enabled + } + if a.JoinTokenTTL != nil { + if d, err := time.ParseDuration(*a.JoinTokenTTL); err == nil { + cfg.NATS.Auth.JoinTokenTTL = d + } + } + if a.JoinTokenSingleUse != nil { + cfg.NATS.Auth.JoinTokenSingleUse = *a.JoinTokenSingleUse + } + if a.EnrollmentEndpointEnabled != nil { + cfg.NATS.Auth.EnrollmentEndpointEnabled = *a.EnrollmentEndpointEnabled + } } func mergeStoreFile(cfg *Config, fc fileConfig) { @@ -641,6 +715,10 @@ func applyEnv(cfg *Config) error { setString(&cfg.NATS.Addr, "SQI_NATS_ADDR") setString(&cfg.NATS.DataDir, "SQI_NATS_DATA_DIR") collect(setInt(&cfg.NATS.MaxStoreMB, "SQI_NATS_MAX_STORE_MB")) + collect(setBool(&cfg.NATS.Auth.Enabled, "SQI_NATS_AUTH_ENABLED")) + collect(setDuration(&cfg.NATS.Auth.JoinTokenTTL, "SQI_NATS_AUTH_JOIN_TOKEN_TTL")) + collect(setBool(&cfg.NATS.Auth.JoinTokenSingleUse, "SQI_NATS_AUTH_JOIN_TOKEN_SINGLE_USE")) + collect(setBool(&cfg.NATS.Auth.EnrollmentEndpointEnabled, "SQI_NATS_AUTH_ENROLLMENT_ENDPOINT_ENABLED")) setString(&cfg.Store.SQLitePath, "SQI_STORE_SQLITE_PATH") collect(setDuration(&cfg.Store.CheckpointInterval, "SQI_STORE_CHECKPOINT_INTERVAL")) diff --git a/internal/config/validate.go b/internal/config/validate.go index ad07e29e..a938ef75 100644 --- a/internal/config/validate.go +++ b/internal/config/validate.go @@ -135,6 +135,22 @@ func validateNATS(cfg NATSConfig) []ValidationError { Message: fmt.Sprintf("must be > 0, got %d; set SQI_NATS_MAX_STORE_MB or nats.max_store_mb", cfg.MaxStoreMB), }) } + // A disabled auth block must never produce validation errors — the same + // rule validateAuth follows. Turning a feature off must not stop a + // server from starting because of a value nobody is reading. + if !cfg.Auth.Enabled { + return errs + } + if d := cfg.Auth.JoinTokenTTL; d < MinNATSAuthJoinTokenTTL || d > MaxNATSAuthJoinTokenTTL { + errs = append(errs, ValidationError{ + Field: "nats.auth.join_token_ttl", + Message: fmt.Sprintf( + "must be between %s and %s, got %s; set %s or %s", + MinNATSAuthJoinTokenTTL, MaxNATSAuthJoinTokenTTL, d, + "SQI_NATS_AUTH_JOIN_TOKEN_TTL", "nats.auth.join_token_ttl", + ), + }) + } return errs } diff --git a/internal/scheduler/assign.go b/internal/scheduler/assign.go index 1c6b8c4a..6b8fd223 100644 --- a/internal/scheduler/assign.go +++ b/internal/scheduler/assign.go @@ -18,7 +18,7 @@ package scheduler // parameters to concrete local paths (resolved mode). // // Lease semantics: -// Workers request work via the core-NATS work.lease. request/reply +// Workers request work via the core-NATS work.lease.. request/reply // protocol ([handleLeaseRequest]). The scheduler selects an eligible, fitting // batch of ready tasks, leases each atomically, and returns the marshaled // payloads built here in the reply. The server never pushes assignments to a diff --git a/internal/scheduler/assignment_test.go b/internal/scheduler/assignment_test.go index 3151ee33..f8fb49a6 100644 --- a/internal/scheduler/assignment_test.go +++ b/internal/scheduler/assignment_test.go @@ -51,7 +51,7 @@ func (*recordBus) SubscribeWorkerDiag(_ func(subject string, data []byte)) (*nat return nil, nil } -func (*recordBus) SubscribeLease(_ func(string, []byte) []byte) (*nats.Subscription, error) { +func (*recordBus) SubscribeLease(_ func(string, string, []byte) []byte) (*nats.Subscription, error) { return nil, nil } diff --git a/internal/scheduler/attemptcache.go b/internal/scheduler/attemptcache.go new file mode 100644 index 00000000..92e2a97c --- /dev/null +++ b/internal/scheduler/attemptcache.go @@ -0,0 +1,72 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package scheduler + +import "sync" + +// maxAttemptOwnerCacheEntries bounds [attemptOwnerCache] independently of +// terminal-status eviction. Normal operation never approaches it — every +// attempt is evicted when it reaches a terminal status (see +// handleTaskTerminal and handleTaskFailed) — so this only guards against an +// attempt whose terminal status this process never observes (e.g. one still +// in flight across a server restart), which would otherwise accumulate for +// the life of the process. +const maxAttemptOwnerCacheEntries = 8192 + +// attemptOwner is the (workerID, taskID) pair recorded for a task attempt at +// creation time. Both fields are immutable for the life of the attempt: no +// store implementation ever updates worker_id or task_id on an existing +// task_attempts row, so caching them is safe for as long as the entry exists. +type attemptOwner struct { + workerID string + taskID string +} + +// attemptOwnerCache is a bounded, concurrency-safe cache of task-attempt +// ownership, consulted by handleLogChunk before it falls back to +// store.GetTaskAttempt. Log chunks arrive roughly every 500ms per running +// task, so caching the two fields handleLogChunk actually checks turns +// hundreds of repeated reads of the same immutable row into one. +type attemptOwnerCache struct { + mu sync.Mutex + entries map[string]attemptOwner +} + +// newAttemptOwnerCache returns an empty attemptOwnerCache. +func newAttemptOwnerCache() *attemptOwnerCache { + return &attemptOwnerCache{entries: make(map[string]attemptOwner)} +} + +// get returns the cached owner for attemptID, if present. +func (c *attemptOwnerCache) get(attemptID string) (attemptOwner, bool) { + c.mu.Lock() + defer c.mu.Unlock() + o, ok := c.entries[attemptID] + return o, ok +} + +// put records attemptID's owner. If the cache is already at capacity, one +// arbitrary existing entry is dropped first — Go map iteration order is +// randomized, so this is not LRU, but it only ever runs when +// maxAttemptOwnerCacheEntries has already been reached, which normal +// terminal-status eviction is designed to prevent. +func (c *attemptOwnerCache) put(attemptID, workerID, taskID string) { + c.mu.Lock() + defer c.mu.Unlock() + if _, exists := c.entries[attemptID]; !exists && len(c.entries) >= maxAttemptOwnerCacheEntries { + for k := range c.entries { + delete(c.entries, k) + break + } + } + c.entries[attemptID] = attemptOwner{workerID: workerID, taskID: taskID} +} + +// evict removes attemptID's entry, if present. Called once an attempt +// reaches a terminal status, since it can never receive another log chunk +// after that. +func (c *attemptOwnerCache) evict(attemptID string) { + c.mu.Lock() + defer c.mu.Unlock() + delete(c.entries, attemptID) +} diff --git a/internal/scheduler/attemptcache_test.go b/internal/scheduler/attemptcache_test.go new file mode 100644 index 00000000..21b3335c --- /dev/null +++ b/internal/scheduler/attemptcache_test.go @@ -0,0 +1,97 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package scheduler + +import ( + "strconv" + "sync" + "testing" +) + +func TestAttemptOwnerCache_PutGetEvict(t *testing.T) { + c := newAttemptOwnerCache() + + if _, ok := c.get("missing"); ok { + t.Fatal("expected miss on empty cache") + } + + c.put("a1", "worker-1", "task-1") + got, ok := c.get("a1") + if !ok { + t.Fatal("expected hit after put") + } + if got.workerID != "worker-1" || got.taskID != "task-1" { + t.Errorf("got %+v, want {worker-1 task-1}", got) + } + + c.evict("a1") + if _, ok := c.get("a1"); ok { + t.Error("expected miss after evict") + } + + // Evicting a key that was never present must not panic. + c.evict("never-there") +} + +func TestAttemptOwnerCache_BoundedUnderContinuousGrowth(t *testing.T) { + c := newAttemptOwnerCache() + + // Simulate a server that never observes a terminal status: put far more + // entries than the cap without ever calling evict. + for i := range maxAttemptOwnerCacheEntries * 3 { + id := strconv.Itoa(i) + c.put(id, "worker-1", "task-"+id) + } + + c.mu.Lock() + n := len(c.entries) + c.mu.Unlock() + if n > maxAttemptOwnerCacheEntries { + t.Errorf("cache grew to %d entries, want <= %d", n, maxAttemptOwnerCacheEntries) + } +} + +// TestAttemptOwnerCache_ConcurrentAccess drives put/get/evict concurrently +// under -race, then checks a real final invariant rather than just surviving +// without a panic or a race report: every id gets its own unique key (no two +// goroutines contend for the same entry), so the final state is fully +// determined by each goroutine's own i%3==0 evict decision. That lets the +// assertion also catch a put/get keying mistake — e.g. get reading back a +// neighboring entry's taskID — which a mere "did it panic" check cannot. +func TestAttemptOwnerCache_ConcurrentAccess(t *testing.T) { + c := newAttemptOwnerCache() + const n = 50 + var wg sync.WaitGroup + for i := range n { + wg.Add(1) + go func(i int) { + defer wg.Done() + id := strconv.Itoa(i) + c.put(id, "worker-1", "task-"+id) + c.get(id) + if i%3 == 0 { + c.evict(id) + } + }(i) + } + wg.Wait() + + for i := range n { + id := strconv.Itoa(i) + got, ok := c.get(id) + if i%3 == 0 { + if ok { + t.Errorf("id %s: expected evict to have removed the entry, got %+v", id, got) + } + continue + } + if !ok { + t.Errorf("id %s: expected entry to remain present", id) + continue + } + wantTaskID := "task-" + id + if got.taskID != wantTaskID { + t.Errorf("id %s: taskID = %q, want %q (own key, not a neighbor's)", id, got.taskID, wantTaskID) + } + } +} diff --git a/internal/scheduler/cancellation_test.go b/internal/scheduler/cancellation_test.go index 6ceca71f..1eaef4f5 100644 --- a/internal/scheduler/cancellation_test.go +++ b/internal/scheduler/cancellation_test.go @@ -49,7 +49,7 @@ func (*stubBus) SubscribeWorkerDiag(_ func(subject string, data []byte)) (*nats. return nil, nil } -func (*stubBus) SubscribeLease(_ func(string, []byte) []byte) (*nats.Subscription, error) { +func (*stubBus) SubscribeLease(_ func(string, string, []byte) []byte) (*nats.Subscription, error) { return nil, nil } diff --git a/internal/scheduler/exprcaps_test.go b/internal/scheduler/exprcaps_test.go index 12d33e58..d607f097 100644 --- a/internal/scheduler/exprcaps_test.go +++ b/internal/scheduler/exprcaps_test.go @@ -555,7 +555,7 @@ func TestHandleWorkerRegister_PersistsAdvertisedExprCaps(t *testing.T) { AssignmentRetainedBytes: 4_444_444, } msg := &fakeJSMsg{ - subject: bus.SubjectWorkerRegister, + subject: bus.WorkerRegisterSubject("w-1"), data: workerMsgJSON(t, protocol.RegisterMsg{ Version: protocol.ProtocolVersion, Type: protocol.TypeRegister, WorkerID: "w-1", FarmID: "farm-1", Hostname: "node-1", OS: "linux", @@ -588,7 +588,7 @@ func TestHandleWorkerRegister_ExprCapWarningIsDeDuplicated(t *testing.T) { register := func(positions int64) { s.handleWorkerMessage(&fakeJSMsg{ - subject: bus.SubjectWorkerRegister, + subject: bus.WorkerRegisterSubject("w-1"), data: workerMsgJSON(t, protocol.RegisterMsg{ Version: protocol.ProtocolVersion, Type: protocol.TypeRegister, WorkerID: "w-1", FarmID: "farm-1", Hostname: "node-1", OS: "linux", @@ -749,7 +749,7 @@ func TestHandleWorkerRegister_EveryWireFieldReachesTheStore(t *testing.T) { }, } - msg := &fakeJSMsg{subject: bus.SubjectWorkerRegister, data: workerMsgJSON(t, sent)} + msg := &fakeJSMsg{subject: bus.WorkerRegisterSubject(sent.WorkerID), data: workerMsgJSON(t, sent)} s.handleWorkerMessage(msg) w, err := st.GetWorker(t.Context(), sent.WorkerID) diff --git a/internal/scheduler/failure.go b/internal/scheduler/failure.go index 22b5be44..0c4deb64 100644 --- a/internal/scheduler/failure.go +++ b/internal/scheduler/failure.go @@ -61,6 +61,10 @@ func (s *Scheduler) handleTaskFailed(ctx context.Context, attempt store.TaskAtte if err != nil { return err } + // RecordTaskFailure closes the attempt as failed whether or not this is + // the message that first closed it, so it can never receive another log + // chunk from here on — evict unconditionally, same as handleTaskTerminal. + s.attemptCache.evict(attempt.ID) if !firstClose { // The attempt was already terminal when this message arrived. That is diff --git a/internal/scheduler/failure_test.go b/internal/scheduler/failure_test.go index cc73853d..96b6a325 100644 --- a/internal/scheduler/failure_test.go +++ b/internal/scheduler/failure_test.go @@ -186,7 +186,7 @@ func (h *failureHarness) reportFailedWithMessage(taskID, message string) { Message: message, At: time.Now().UTC(), } - if err := h.s.processTaskStatus(h.t.Context(), msg); err != nil { + if err := h.s.processTaskStatus(h.t.Context(), attempt.WorkerID, msg); err != nil { h.t.Fatalf("processTaskStatus(failed): %v", err) } } @@ -339,6 +339,28 @@ func TestHandleTaskFailed_RetriesUntilCeiling(t *testing.T) { } } +// TestHandleTaskFailed_Retry_EvictsAttemptCache proves the RETRY branch of +// handleTaskFailed evicts the attempt-owner cache entry itself. That branch +// returns from retryTaskAfterFailure without ever reaching handleTaskTerminal, +// so the terminal path's own evict call cannot cover it — deleting +// handleTaskFailed's evict would leave every other test in this file green. +func TestHandleTaskFailed_Retry_EvictsAttemptCache(t *testing.T) { + h := newFailureHarness(t, RetryPolicy{MaxAttempts: 2, RetryDelay: 0, FailureLimit: 0}) + h.seedRunningTask("j1", "t1", "w1") + + attempt := h.current["t1"] + h.s.attemptCache.put(attempt.ID, attempt.WorkerID, attempt.TaskID) + + h.reportFailed("t1") // first failure: RETRY, not EXHAUSTED (ceiling is 2) + + if got := h.taskStatus("t1"); got != store.TaskStatusReady { + t.Fatalf("expected retry to requeue the task, got %s", got) + } + if _, ok := h.s.attemptCache.get(attempt.ID); ok { + t.Error("expected attempt-owner cache entry to be evicted on the retry branch") + } +} + // TestHandleTaskFailed_RedeliveryCountsOnce is the IMP-1 regression at the // scheduler layer. The task-status JetStream consumer is at-least-once, so the // same "failed" message can be delivered more than once (NAK, AckWait expiry, @@ -360,7 +382,7 @@ func TestHandleTaskFailed_RedeliveryCountsOnce(t *testing.T) { } // First delivery: one genuine failure → retry, task back to ready. - if err := h.s.processTaskStatus(t.Context(), msg); err != nil { + if err := h.s.processTaskStatus(t.Context(), h.current["t1"].WorkerID, msg); err != nil { t.Fatalf("first delivery: %v", err) } if got := h.taskFailedAttempts("t1"); got != 1 { @@ -377,7 +399,7 @@ func TestHandleTaskFailed_RedeliveryCountsOnce(t *testing.T) { // failed, so RecordTaskFailure returns the current counts without // re-incrementing, the retry decision is re-made identically, and the // requeue is re-applied harmlessly. - if err := h.s.processTaskStatus(t.Context(), msg); err != nil { + if err := h.s.processTaskStatus(t.Context(), h.current["t1"].WorkerID, msg); err != nil { t.Fatalf("redelivery: %v", err) } if got := h.taskFailedAttempts("t1"); got != 1 { @@ -616,7 +638,7 @@ func TestHandleTaskFailed_SupersededAttempt_LeavesReleasedTaskAlone(t *testing.T ExitCode: &exitCode, At: time.Now().UTC(), } - if err := h.s.processTaskStatus(ctx, msg); err != nil { + if err := h.s.processTaskStatus(ctx, stale.WorkerID, msg); err != nil { t.Fatalf("processTaskStatus(stale failed): %v", err) } diff --git a/internal/scheduler/jobdeps_test.go b/internal/scheduler/jobdeps_test.go index fc7d3c1c..2ce36841 100644 --- a/internal/scheduler/jobdeps_test.go +++ b/internal/scheduler/jobdeps_test.go @@ -15,11 +15,18 @@ import ( "github.com/google/uuid" + "github.com/uberware/sqi/internal/bus" "github.com/uberware/sqi/internal/store" fakestore "github.com/uberware/sqi/internal/store/fake" "github.com/uberware/sqi/internal/worker/protocol" ) +// jobDepsWorkerID is the worker seedRunnableJob's attempt is opened on, and +// the worker whose subject completeJob publishes the "succeeded" status on — +// they must match, since handleTaskStatusMessage now requires the subject's +// worker ID to match the attempt's recorded owner. +const jobDepsWorkerID = "worker-1" + // ── test harness ──────────────────────────────────────────────────────────── // jobDepsHarness bundles a Scheduler wired over a fresh fake store with a @@ -189,6 +196,7 @@ func (h *jobDepsHarness) seedRunnableJob(t *testing.T) store.Job { h.attempt, err = h.store.CreateTaskAttempt(ctx, store.TaskAttempt{ ID: uuid.NewString(), TaskID: task.ID, + WorkerID: jobDepsWorkerID, AttemptNumber: 1, Status: store.AttemptStatusRunning, StartedAt: time.Now(), @@ -212,6 +220,7 @@ func (h *jobDepsHarness) completeJob(t *testing.T, jobID string) { } exitCode := 0 msg := &fakeJSMsg{ + subject: bus.TaskStatusSubject(jobDepsWorkerID, h.task.JobID), data: taskStatusMsgJSON(t, protocol.TaskStatusMsg{ Version: protocol.ProtocolVersion, TaskID: h.task.ID, diff --git a/internal/scheduler/lease.go b/internal/scheduler/lease.go index 3f8bab70..968800c5 100644 --- a/internal/scheduler/lease.go +++ b/internal/scheduler/lease.go @@ -33,14 +33,40 @@ type leaseReply struct { // handleLeaseRequest decodes a lease request, leases a fitting batch, and on an // empty result parks the request in the waiter registry until new work appears // or leaseHoldTimeout elapses, then replies once more. -func (s *Scheduler) handleLeaseRequest(queueID string, data []byte) []byte { +// +// workerID is the identity carried by the request's subject; queueID is the +// queue the worker asked about. +func (s *Scheduler) handleLeaseRequest(workerID, queueID string, data []byte) []byte { ctx := s.ctx var req leaseRequest if err := json.Unmarshal(data, &req); err != nil || req.WorkerID == "" { + // Two cases, both refused the same way: a body this server cannot + // decode at all, and one that decodes but carries no identity to + // check against the subject. In neither is there anything to + // authorize, so the subject ID is all that is left to log. + // Debug, not warn: an unauthenticated broker lets anything publish here, + // so a warn would be a log-flood vector. + s.logger.DebugContext( + ctx, "scheduler: malformed lease request", + slog.String("subject_worker_id", workerID), + ) + return marshalLeaseReply(nil) + } + + // The subject is authoritative. A payload that names a different worker + // is either a stale client or an attempt to have tasks assigned to + // another worker while this connection receives the job code. No + // req.WorkerID == "" guard here: that case already returned above. + if req.WorkerID != workerID { + s.logger.WarnContext( + ctx, "scheduler: lease request whose payload identity differs from its subject — refusing", + slog.String("subject_worker_id", workerID), + slog.String("payload_worker_id", req.WorkerID), + ) return marshalLeaseReply(nil) } - worker, err := s.store.GetWorker(ctx, req.WorkerID) + worker, err := s.store.GetWorker(ctx, workerID) if err != nil { return marshalLeaseReply(nil) } @@ -49,7 +75,7 @@ func (s *Scheduler) handleLeaseRequest(queueID string, data []byte) []byte { if err != nil { s.logger.WarnContext( ctx, "scheduler: lease selection failed", - slog.String("worker_id", req.WorkerID), + slog.String("worker_id", workerID), slog.Any("error", err), ) return marshalLeaseReply(nil) @@ -62,7 +88,7 @@ func (s *Scheduler) handleLeaseRequest(queueID string, data []byte) []byte { // The park happens OUTSIDE the per-worker lock; only the selection below is // serialized, so a re-woken request reads the up-to-date committed cores. if s.waiters.wait(ctx, queueID, s.leaseHoldTimeout) { - if w2, err2 := s.store.GetWorker(ctx, req.WorkerID); err2 == nil { + if w2, err2 := s.store.GetWorker(ctx, workerID); err2 == nil { if batch2, err2 := s.selectLeaseBatchLocked(ctx, w2); err2 == nil { return marshalLeaseReply(batch2) } diff --git a/internal/scheduler/lease_test.go b/internal/scheduler/lease_test.go index e143fb2b..c7fb13d0 100644 --- a/internal/scheduler/lease_test.go +++ b/internal/scheduler/lease_test.go @@ -33,7 +33,7 @@ func TestHandleLeaseRequest_QueuelessWorkerWildcardToken(t *testing.T) { if err != nil { t.Fatalf("marshal request: %v", err) } - reply := s.handleLeaseRequest(bus.WildcardQueueToken, req) + reply := s.handleLeaseRequest(w.ID, bus.WildcardQueueToken, req) var got leaseReply if err := json.Unmarshal(reply, &got); err != nil { @@ -242,7 +242,7 @@ func TestHandleLeaseRequest_ReturnsBatch(t *testing.T) { if err != nil { t.Fatalf("marshal request: %v", err) } - reply := s.handleLeaseRequest("q1", req) + reply := s.handleLeaseRequest(w.ID, "q1", req) var got leaseReply if err := json.Unmarshal(reply, &got); err != nil { @@ -276,7 +276,7 @@ func TestHandleLeaseRequest_ConcurrentSameWorkerDoesNotOverLease(t *testing.T) { for range 2 { go func() { defer wg.Done() - _ = s.handleLeaseRequest("q1", req) + _ = s.handleLeaseRequest(w.ID, "q1", req) }() } wg.Wait() @@ -375,7 +375,7 @@ func TestHandleLeaseRequest_EmptyTimesOut(t *testing.T) { t.Fatalf("marshal request: %v", err) } start := time.Now() - reply := s.handleLeaseRequest("q1", req) + reply := s.handleLeaseRequest("w1", "q1", req) if elapsed := time.Since(start); elapsed < 30*time.Millisecond { t.Errorf("returned too fast (%v); expected to park until timeout", elapsed) } diff --git a/internal/scheduler/logingest.go b/internal/scheduler/logingest.go index 647a12a6..afac3bfc 100644 --- a/internal/scheduler/logingest.go +++ b/internal/scheduler/logingest.go @@ -5,7 +5,7 @@ package scheduler // Structured log ingestion that timestamps and persists log chunks // with monotonic sequence numbers per task attempt. // -// Workers publish [protocol.LogChunkMsg] values to task.logs. as their +// Workers publish [protocol.LogChunkMsg] values to task.logs.. as their // task produces stdout/stderr output. The server-side consumer here: // // 1. Decodes each [protocol.LogChunkMsg]. @@ -27,12 +27,14 @@ package scheduler import ( "context" "encoding/json" + "errors" "log/slog" "time" "github.com/google/uuid" "github.com/nats-io/nats.go/jetstream" + "github.com/uberware/sqi/internal/bus" "github.com/uberware/sqi/internal/store" "github.com/uberware/sqi/internal/worker/protocol" "github.com/uberware/sqi/internal/ws" @@ -46,8 +48,16 @@ func (s *Scheduler) startTaskLogsConsumer(ctx context.Context) error { return err } -// handleLogChunk is the JetStream message handler for task.logs. messages +// handleLogChunk is the JetStream message handler for task.logs.. messages // published by workers. +// +// [protocol.LogChunkMsg] carries no worker-identity field of its own, so the +// subject is the ONLY identity available: it is resolved to an attempt, and +// that attempt's recorded WorkerID and TaskID are both checked before a +// chunk is persisted or fanned out — the worker check alone would let a +// worker holding a live attempt of its own pair that attempt with a +// different task in the payload and inject log content there. See the +// auth-off note on [Scheduler.handleTaskStatusMessage]: it applies here too. func (s *Scheduler) handleLogChunk(msg jetstream.Msg) { ctx := s.ctx @@ -70,6 +80,77 @@ func (s *Scheduler) handleLogChunk(msg jetstream.Msg) { return } + // The subject is the only identity NATS itself can vouch for; a message + // on a subject that does not carry one concrete worker ID cannot be + // attributed to anyone and is discarded rather than acted on. + subjectWorkerID, _, ok := bus.ParseWorkerSubject(msg.Subject()) + if !ok { + s.discardUnexpectedSubject(ctx, msg, "task.logs") + return + } + + // Resolve the attempt this chunk claims to belong to so its recorded + // WorkerID and TaskID can be checked below. Both fields are immutable for + // the life of the attempt, so a cache hit is as good as a fresh store + // read; on a miss, fall back to the store exactly as before — including + // the transient-vs-permanent distinction below, since the store read + // also confirms the attempt still exists (a chunk for a vanished attempt + // must still be discarded, never assumed present). + owner, ok := s.attemptCache.get(m.AttemptID) + if !ok { + attempt, err := s.store.GetTaskAttempt(ctx, m.AttemptID) + if errors.Is(err, store.ErrNotFound) { + s.logger.WarnContext( + ctx, "scheduler: task.logs for unknown attempt — discarding", + slog.String("attempt_id", m.AttemptID), + slog.String("task_id", m.TaskID), + ) + s.ackMsg(ctx, msg) + return + } + if err != nil { + s.logger.WarnContext( + ctx, "scheduler: task.logs: lookup attempt failed — will redeliver", + slog.String("attempt_id", m.AttemptID), + slog.Any("error", err), + ) + s.nakMsg(ctx, msg) + return + } + owner = attemptOwner{workerID: attempt.WorkerID, taskID: attempt.TaskID} + s.attemptCache.put(m.AttemptID, owner.workerID, owner.taskID) + } + // The attempt is real, but is it for the task this chunk claims? Without + // this, a worker holding a live attempt of its own could pair that + // attempt's ID with a different task's ID in the payload — the worker-ID + // check below would pass (the attempt really is the subject worker's), + // but the log would land against a task that worker was never assigned. + if owner.taskID != m.TaskID { + s.logger.WarnContext( + ctx, "scheduler: task.logs attempt task_id mismatch — discarding", + slog.String("attempt_id", m.AttemptID), + slog.String("msg_task_id", m.TaskID), + slog.String("attempt_task_id", owner.taskID), + ) + s.ackMsg(ctx, msg) + return + } + // The subject's worker ID was enforced by NATS when broker auth is on; + // with LogChunkMsg carrying no worker field of its own, it is the only + // identity this handler has to check at all. Treat a mismatch as + // permanent — redelivery cannot make a forged message legal. + if subjectWorkerID != owner.workerID { + s.logger.WarnContext( + ctx, "scheduler: task.logs from a worker that does not hold this task — discarding", + slog.String("task_id", m.TaskID), + slog.String("attempt_id", m.AttemptID), + slog.String("subject_worker_id", subjectWorkerID), + slog.String("attempt_worker_id", owner.workerID), + ) + s.ackMsg(ctx, msg) + return + } + // Extract the NATS JetStream sequence number from the message metadata. // This is the stable cursor used by the log-tail REST endpoint. // NATS stream sequences are uint64; int64 is used here to match SQLite's @@ -101,6 +182,12 @@ func (s *Scheduler) handleLogChunk(msg jetstream.Msg) { } if _, err := s.store.CreateTaskLog(ctx, log); err != nil { + // The cached owner may be stale: the attempt row can be deleted along + // with its job while a chunk is in flight (e.g. DELETE /jobs/{id} on + // an active job). Dropping the entry sends the redelivery through the + // store path above, which discards the chunk instead of failing this + // write again. + s.attemptCache.evict(m.AttemptID) s.logger.WarnContext( ctx, "scheduler: persist log chunk failed — will redeliver", slog.String("task_id", m.TaskID), diff --git a/internal/scheduler/logingest_test.go b/internal/scheduler/logingest_test.go index a5f80756..abc7412d 100644 --- a/internal/scheduler/logingest_test.go +++ b/internal/scheduler/logingest_test.go @@ -13,6 +13,8 @@ package scheduler import ( "context" "encoding/json" + "errors" + "fmt" "log/slog" "testing" "time" @@ -20,12 +22,18 @@ import ( "github.com/google/uuid" "github.com/nats-io/nats.go/jetstream" + "github.com/uberware/sqi/internal/bus" "github.com/uberware/sqi/internal/store" "github.com/uberware/sqi/internal/store/fake" "github.com/uberware/sqi/internal/worker/protocol" "github.com/uberware/sqi/internal/ws" ) +// logTestWorkerID is the worker whose subject these tests publish log chunks +// on. handleLogChunk now checks it against the chunk's attempt, so any test +// exercising the persist path must open the attempt on this same worker. +const logTestWorkerID = "worker-1" + // ── fakeJSMsg: minimal jetstream.Msg for log ingest tests ──────────────────── // fakeJSMsg embeds jetstream.Msg (nil) and overrides only the methods called @@ -98,7 +106,15 @@ func TestHandleLogChunk_ValidStdout(t *testing.T) { taskID := uuid.NewString() now := time.Now().UTC() + if _, err := st.CreateTaskAttempt(t.Context(), store.TaskAttempt{ + ID: attemptID, TaskID: taskID, WorkerID: logTestWorkerID, + AttemptNumber: 1, Status: store.AttemptStatusRunning, StartedAt: now, + }); err != nil { + t.Fatalf("CreateTaskAttempt: %v", err) + } + msg := &fakeJSMsg{ + subject: bus.TaskLogsSubject(logTestWorkerID, taskID), natsSeq: 42, data: msgJSON(t, protocol.LogChunkMsg{ TaskID: taskID, @@ -141,13 +157,23 @@ func TestHandleLogChunk_ValidStderr(t *testing.T) { s.ctx = t.Context() attemptID := uuid.NewString() + taskID := uuid.NewString() + now := time.Now().UTC() + if _, err := st.CreateTaskAttempt(t.Context(), store.TaskAttempt{ + ID: attemptID, TaskID: taskID, WorkerID: logTestWorkerID, + AttemptNumber: 1, Status: store.AttemptStatusRunning, StartedAt: now, + }); err != nil { + t.Fatalf("CreateTaskAttempt: %v", err) + } + msg := &fakeJSMsg{ + subject: bus.TaskLogsSubject(logTestWorkerID, taskID), natsSeq: 1, data: msgJSON(t, protocol.LogChunkMsg{ - TaskID: uuid.NewString(), + TaskID: taskID, AttemptID: attemptID, SeqNum: 1, - At: time.Now().UTC(), + At: now, Stream: "stderr", Data: "error output", }), @@ -174,10 +200,19 @@ func TestHandleLogChunk_ZeroAtUsesServerTime(t *testing.T) { before := time.Now().UTC() attemptID := uuid.NewString() + taskID := uuid.NewString() + if _, err := st.CreateTaskAttempt(t.Context(), store.TaskAttempt{ + ID: attemptID, TaskID: taskID, WorkerID: logTestWorkerID, + AttemptNumber: 1, Status: store.AttemptStatusRunning, StartedAt: before, + }); err != nil { + t.Fatalf("CreateTaskAttempt: %v", err) + } + msg := &fakeJSMsg{ + subject: bus.TaskLogsSubject(logTestWorkerID, taskID), natsSeq: 1, data: msgJSON(t, protocol.LogChunkMsg{ - TaskID: uuid.NewString(), + TaskID: taskID, AttemptID: attemptID, SeqNum: 1, At: time.Time{}, // zero → server clock @@ -259,14 +294,23 @@ func TestHandleLogChunk_MissingAttemptID_Acked(t *testing.T) { func TestHandleLogChunk_StoreFailure_Nacked(t *testing.T) { inner := fake.New() + taskID := uuid.NewString() + attemptID := uuid.NewString() + if _, err := inner.CreateTaskAttempt(t.Context(), store.TaskAttempt{ + ID: attemptID, TaskID: taskID, WorkerID: logTestWorkerID, + AttemptNumber: 1, Status: store.AttemptStatusRunning, StartedAt: time.Now().UTC(), + }); err != nil { + t.Fatalf("CreateTaskAttempt: %v", err) + } est := &logIngestErrSt{Store: inner} s := newLogTestScheduler(est) s.ctx = t.Context() msg := &fakeJSMsg{ + subject: bus.TaskLogsSubject(logTestWorkerID, taskID), data: msgJSON(t, protocol.LogChunkMsg{ - TaskID: uuid.NewString(), - AttemptID: uuid.NewString(), + TaskID: taskID, + AttemptID: attemptID, SeqNum: 1, At: time.Now().UTC(), Stream: "stdout", @@ -291,10 +335,19 @@ func TestHandleLogChunk_MetadataError_NATSSeqZero(t *testing.T) { s.ctx = t.Context() attemptID := uuid.NewString() + taskID := uuid.NewString() + if _, err := st.CreateTaskAttempt(t.Context(), store.TaskAttempt{ + ID: attemptID, TaskID: taskID, WorkerID: logTestWorkerID, + AttemptNumber: 1, Status: store.AttemptStatusRunning, StartedAt: time.Now().UTC(), + }); err != nil { + t.Fatalf("CreateTaskAttempt: %v", err) + } + msg := &fakeJSMsg{ + subject: bus.TaskLogsSubject(logTestWorkerID, taskID), metaErr: context.DeadlineExceeded, // metadata unavailable data: msgJSON(t, protocol.LogChunkMsg{ - TaskID: uuid.NewString(), + TaskID: taskID, AttemptID: attemptID, SeqNum: 7, At: time.Now().UTC(), @@ -319,6 +372,113 @@ func TestHandleLogChunk_MetadataError_NATSSeqZero(t *testing.T) { } } +// ── attemptOwnerCache tests ─────────────────────────────────────────────────── + +// countingAttemptStore wraps a store.Store and counts GetTaskAttempt calls, +// so tests can prove the attempt-owner cache does (or does not) avoid a +// store read. +type countingAttemptStore struct { + store.Store + + getTaskAttemptCalls int +} + +func (s *countingAttemptStore) GetTaskAttempt(ctx context.Context, id string) (store.TaskAttempt, error) { + s.getTaskAttemptCalls++ + return s.Store.GetTaskAttempt(ctx, id) +} + +// TestHandleLogChunk_RepeatedChunk_CachedAfterFirstRead proves a second log +// chunk for the same attempt does not re-read the store: the first chunk +// misses the cache and reads through, the second hits. +func TestHandleLogChunk_RepeatedChunk_CachedAfterFirstRead(t *testing.T) { + cst := &countingAttemptStore{Store: fake.New()} + s := newLogTestScheduler(cst) + s.ctx = t.Context() + + attemptID := uuid.NewString() + taskID := uuid.NewString() + now := time.Now().UTC() + if _, err := cst.CreateTaskAttempt(t.Context(), store.TaskAttempt{ + ID: attemptID, TaskID: taskID, WorkerID: logTestWorkerID, + AttemptNumber: 1, Status: store.AttemptStatusRunning, StartedAt: now, + }); err != nil { + t.Fatalf("CreateTaskAttempt: %v", err) + } + + newChunk := func(seq int64) *fakeJSMsg { + return &fakeJSMsg{ + subject: bus.TaskLogsSubject(logTestWorkerID, taskID), + natsSeq: uint64(seq), // test data, small positive constant + data: msgJSON(t, protocol.LogChunkMsg{ + TaskID: taskID, + AttemptID: attemptID, + SeqNum: seq, + At: now, + Stream: "stdout", + Data: "line", + }), + } + } + + first := newChunk(1) + s.handleLogChunk(first) + if !first.acked { + t.Fatal("expected first chunk to be acked") + } + if cst.getTaskAttemptCalls != 1 { + t.Fatalf("getTaskAttemptCalls after first chunk = %d, want 1 (cache miss falls back to the store)", cst.getTaskAttemptCalls) + } + + second := newChunk(2) + s.handleLogChunk(second) + if !second.acked { + t.Fatal("expected second chunk to be acked") + } + if cst.getTaskAttemptCalls != 1 { + t.Errorf("getTaskAttemptCalls after second chunk = %d, want still 1 (cache hit must not re-read the store)", cst.getTaskAttemptCalls) + } + + logs, err := cst.ListTaskLogs(t.Context(), attemptID, 0, 100) + if err != nil { + t.Fatalf("ListTaskLogs: %v", err) + } + if len(logs) != 2 { + t.Fatalf("expected 2 log rows, got %d", len(logs)) + } +} + +// TestHandleLogChunk_CacheMiss_VanishedAttempt_StillDiscarded proves that on +// a cache miss for an attempt the store has never heard of, handleLogChunk +// still falls back to the store (rather than assuming ownership) and +// discards the chunk exactly as it would with no cache at all. +func TestHandleLogChunk_CacheMiss_VanishedAttempt_StillDiscarded(t *testing.T) { + cst := &countingAttemptStore{Store: fake.New()} + s := newLogTestScheduler(cst) + s.ctx = t.Context() + + msg := &fakeJSMsg{ + subject: bus.TaskLogsSubject(logTestWorkerID, uuid.NewString()), + data: msgJSON(t, protocol.LogChunkMsg{ + TaskID: uuid.NewString(), + AttemptID: uuid.NewString(), // never created — vanished/unknown + SeqNum: 1, + At: time.Now().UTC(), + Stream: "stdout", + Data: "line", + }), + } + + s.handleLogChunk(msg) + + if !msg.acked { + t.Error("chunk for an unknown attempt should be acked (discarded)") + } + if cst.getTaskAttemptCalls != 1 { + t.Errorf("getTaskAttemptCalls = %d, want 1 (a cache miss must still consult the store)", cst.getTaskAttemptCalls) + } +} + // ── logIngestErrSt: store that fails CreateTaskLog ──────────────────────────── type logIngestErrSt struct { @@ -330,3 +490,148 @@ func (*logIngestErrSt) CreateTaskLog(_ context.Context, _ store.TaskLog) (store. } var errInjectedLog = context.DeadlineExceeded + +// ── attemptOwnerCache: stale hit on a deleted attempt ───────────────────── + +// fkEnforcingLogStore wraps a store.Store and makes CreateTaskLog fail +// exactly as the SQLite backend does when a log row's AttemptID no longer +// names an existing task_attempts row: a FOREIGN KEY constraint failure, +// which mapErr turns into store.ErrConflict. The fake store enforces no such +// constraint (internal/store/fake/task_log.go appends unconditionally), so a +// test exercising this path needs this wrapper. +type fkEnforcingLogStore struct { + store.Store +} + +func (s *fkEnforcingLogStore) CreateTaskLog(ctx context.Context, log store.TaskLog) (store.TaskLog, error) { + if _, err := s.GetTaskAttempt(ctx, log.AttemptID); errors.Is(err, store.ErrNotFound) { + return store.TaskLog{}, fmt.Errorf("%w: FOREIGN KEY constraint failed", store.ErrConflict) + } else if err != nil { + return store.TaskLog{}, err + } + return s.Store.CreateTaskLog(ctx, log) +} + +// TestHandleLogChunk_CacheHit_DeletedAttempt_SelfHeals is the regression test +// for the hazard on the cache-HIT path: a cache hit against an attempt whose +// row has been deleted out from under it — e.g. DELETE /api/v1/jobs/{id} on +// an ACTIVE job cancels its tasks and then runs DeleteJob, which removes the +// task_attempts row, while a worker is still publishing chunks for the +// window before its process notices — must not turn into a NAK loop. +// +// The first chunk populates the cache. Deleting the job removes the attempt +// row underneath it, but nothing evicts the cache entry on that path (it is +// a bulk write, not a terminal task-status message). The next chunk hits the +// stale-but-identity-matching cache entry, so both ownership checks pass and +// the write is attempted — and must fail. That failure has to evict the +// entry so the redelivery takes the store path and discards the chunk +// cleanly, rather than hitting the same stale entry and failing forever. +func TestHandleLogChunk_CacheHit_DeletedAttempt_SelfHeals(t *testing.T) { + base := fake.New() + st := &fkEnforcingLogStore{Store: base} + s := newLogTestScheduler(st) + s.ctx = t.Context() + + ctx := t.Context() + now := time.Now().UTC() + + if _, err := base.CreateFarm(ctx, store.Farm{ID: "farm-1", Name: "farm-1"}); err != nil { + t.Fatalf("CreateFarm: %v", err) + } + if _, err := base.CreateQueue(ctx, store.Queue{ID: "queue-1", FarmID: "farm-1", Name: "queue-1"}); err != nil { + t.Fatalf("CreateQueue: %v", err) + } + job, err := base.CreateJob(ctx, store.Job{ + ID: uuid.NewString(), FarmID: "farm-1", QueueID: "queue-1", Name: "job", + Status: store.JobStatusRunning, TemplateFormat: store.TemplateFormatJSON, + CreatedAt: now, UpdatedAt: now, + }) + if err != nil { + t.Fatalf("CreateJob: %v", err) + } + step, err := base.CreateStep(ctx, store.Step{ + ID: uuid.NewString(), JobID: job.ID, Name: "s1", + Status: store.StepStatusRunning, CreatedAt: now, UpdatedAt: now, + }) + if err != nil { + t.Fatalf("CreateStep: %v", err) + } + task, err := base.CreateTask(ctx, store.Task{ + ID: uuid.NewString(), JobID: job.ID, StepID: step.ID, Name: "t1", + Status: store.TaskStatusRunning, AssignedWorkerID: logTestWorkerID, + CreatedAt: now, UpdatedAt: now, + }) + if err != nil { + t.Fatalf("CreateTask: %v", err) + } + attempt, err := base.CreateTaskAttempt(ctx, store.TaskAttempt{ + ID: uuid.NewString(), TaskID: task.ID, WorkerID: logTestWorkerID, + AttemptNumber: 1, Status: store.AttemptStatusRunning, StartedAt: now, CreatedAt: now, + }) + if err != nil { + t.Fatalf("CreateTaskAttempt: %v", err) + } + + newChunk := func(seq int64) *fakeJSMsg { + return &fakeJSMsg{ + subject: bus.TaskLogsSubject(logTestWorkerID, task.ID), + natsSeq: uint64(seq), // test data, small positive constant + data: msgJSON(t, protocol.LogChunkMsg{ + TaskID: task.ID, + AttemptID: attempt.ID, + SeqNum: seq, + At: now, + Stream: "stdout", + Data: "line", + }), + } + } + + // First chunk: cache miss, store read succeeds, cache is populated. + first := newChunk(1) + s.handleLogChunk(first) + if !first.acked { + t.Fatal("expected first chunk to be acked") + } + if _, ok := s.attemptCache.get(attempt.ID); !ok { + t.Fatal("expected first chunk to populate the attempt-owner cache") + } + + // The operator deletes the active job: tasks are canceled and DeleteJob + // removes the attempt row along with everything else. Nothing on this + // path evicts the cache entry, so it survives untouched. + if err := base.DeleteJob(ctx, job.ID); err != nil { + t.Fatalf("DeleteJob: %v", err) + } + if _, ok := s.attemptCache.get(attempt.ID); !ok { + t.Fatal("test setup invariant broken: cache entry should still be present after DeleteJob") + } + + // Second chunk: a cache HIT against the now-stale entry. Both identity + // checks pass (they check the cached, historical values), so the write + // is attempted and fails with the FK-style error. + second := newChunk(2) + s.handleLogChunk(second) + if !second.nacked { + t.Fatal("expected the write against a deleted attempt to nak for redelivery") + } + if second.acked { + t.Error("second chunk should not be acked when the write fails") + } + if _, ok := s.attemptCache.get(attempt.ID); ok { + t.Fatal("expected the failed write to evict the stale cache entry") + } + + // Redelivery: the cache is now empty, so this takes the store path, + // which correctly discards (acks) a chunk for a vanished attempt — + // self-healing, exactly like the pre-cache behavior, instead of nacking + // again and looping. + redelivery := newChunk(2) + s.handleLogChunk(redelivery) + if !redelivery.acked { + t.Error("expected the redelivery to be acked (discarded) once the cache entry is gone") + } + if redelivery.nacked { + t.Error("redelivery should not nak again — that would loop forever") + } +} diff --git a/internal/scheduler/protocolversion_test.go b/internal/scheduler/protocolversion_test.go index 3528fab0..35ab2e8d 100644 --- a/internal/scheduler/protocolversion_test.go +++ b/internal/scheduler/protocolversion_test.go @@ -4,7 +4,7 @@ package scheduler // Tests for the receiver-side wire-protocol version gate on the three worker → // server message types the server ACTS ON: worker.register, worker.heartbeat -// and task.status.. +// and task.status. // // WHY A GATE AT ALL. encoding/json silently drops every field the receiving // struct does not declare, so a message that merely decodes is not evidence @@ -40,7 +40,7 @@ import ( // sender). The empty case matters most — it is the one a reader assumes is // handled and the one an `if version != "" && version != current` gate lets // through. -var mismatchedVersions = []string{"1", "3", ""} +var mismatchedVersions = []string{"2", "4", ""} // ── worker.register ─────────────────────────────────────────────────────────── @@ -57,7 +57,7 @@ func TestHandleWorkerRegister_VersionMismatch_Discarded(t *testing.T) { s := newMetricsScheduler(st, &recordBus{}, "") msg := &fakeJSMsg{ - subject: bus.SubjectWorkerRegister, + subject: bus.WorkerRegisterSubject("w-1"), data: workerMsgJSON(t, protocol.RegisterMsg{ Version: v, Type: protocol.TypeRegister, @@ -110,7 +110,7 @@ func TestHandleWorkerHeartbeat_VersionMismatch_Discarded(t *testing.T) { } msg := &fakeJSMsg{ - subject: bus.SubjectWorkerHeartbeat, + subject: bus.WorkerHeartbeatSubject("w-1"), data: workerMsgJSON(t, protocol.HeartbeatMsg{ Version: v, Type: protocol.TypeHeartbeat, @@ -139,7 +139,7 @@ func TestHandleWorkerHeartbeat_VersionMismatch_Discarded(t *testing.T) { } } -// ── task.status. ───────────────────────────────────────────────────────── +// ── task.status.. ──────────────────────────────────────────────── // TestHandleTaskStatusMessage_VersionMismatch_Discarded is the one with a // cost: a terminal status discarded here means the server never learns the @@ -155,7 +155,7 @@ func TestHandleTaskStatusMessage_VersionMismatch_Discarded(t *testing.T) { job, _, task, attempt := seedStatusFixture(t, st, store.TaskStatusAssigned) msg := &fakeJSMsg{ - subject: "task.status." + job.ID, + subject: bus.TaskStatusSubject("w-1", job.ID), data: taskStatusMsgJSON(t, protocol.TaskStatusMsg{ Version: v, Type: protocol.TypeTaskStatus, diff --git a/internal/scheduler/provenance_test.go b/internal/scheduler/provenance_test.go new file mode 100644 index 00000000..eb0d3384 --- /dev/null +++ b/internal/scheduler/provenance_test.go @@ -0,0 +1,497 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package scheduler + +// Provenance tests: the subject a worker publishes on is the only identity +// NATS itself can vouch for, and these seven tests prove the server actually +// enforces it rather than trusting the payload's own WorkerID field. +// +// Each test has worker A publish on its OWN subject (the only one NATS would +// let it publish on) while the payload claims — or the store already holds — +// worker B's identity, and asserts on STORE STATE that the forgery was +// discarded, not merely that it was logged. + +import ( + "encoding/json" + "testing" + "time" + + "github.com/google/uuid" + + "github.com/uberware/sqi/internal/bus" + "github.com/uberware/sqi/internal/store" + "github.com/uberware/sqi/internal/store/fake" + "github.com/uberware/sqi/internal/worker/protocol" +) + +// mustMarshal marshals v to JSON, failing the test on error. +func mustMarshal(t *testing.T, v any) []byte { + t.Helper() + b, err := json.Marshal(v) + if err != nil { + t.Fatalf("marshal: %v", err) + } + return b +} + +// newFakeJetStreamMsg is an alias over the package's existing fakeJSMsg +// (logingest_test.go), which already implements jetstream.Msg with a +// settable Subject, Data, Ack and Nak. Given as a constructor matching the +// shape described for these tests so the subject — the whole point of a +// provenance test — is set at the call site. +func newFakeJetStreamMsg(_ *testing.T, subject string, data []byte) *fakeJSMsg { + return &fakeJSMsg{subject: subject, data: data} +} + +// seedRunnableTask creates a farm, queue, an online worker registered as +// workerID, a running job/step, and a task assigned to and running on that +// worker. Returns the job, step, and task. +// +// Farm/queue names are derived from workerID (fake.Store.CreateFarm rejects a +// duplicate Name) so a test that needs two independent workers — each with +// its own task — can call this twice without a spurious ErrConflict. +func seedRunnableTask(t *testing.T, st *fake.Store, workerID string) (store.Job, store.Step, store.Task) { + t.Helper() + ctx := t.Context() + now := time.Now().UTC() + + farm, err := st.CreateFarm(ctx, store.Farm{ID: uuid.NewString(), Name: "f-" + workerID}) + if err != nil { + t.Fatalf("CreateFarm: %v", err) + } + queue, err := st.CreateQueue(ctx, store.Queue{ID: uuid.NewString(), FarmID: farm.ID, Name: "q-" + workerID}) + if err != nil { + t.Fatalf("CreateQueue: %v", err) + } + if _, err := st.RegisterWorker(ctx, store.Worker{ + ID: workerID, FarmID: farm.ID, Hostname: workerID, + Status: store.WorkerStatusOnline, CPUCount: 4, LastHeartbeatAt: &now, + Tags: map[string]string{}, + }); err != nil { + t.Fatalf("RegisterWorker: %v", err) + } + job, err := st.CreateJob(ctx, store.Job{ + ID: uuid.NewString(), + FarmID: farm.ID, + QueueID: queue.ID, + Name: "job", + Status: store.JobStatusRunning, + TemplateFormat: store.TemplateFormatJSON, + CreatedAt: now, + UpdatedAt: now, + }) + if err != nil { + t.Fatalf("CreateJob: %v", err) + } + step, err := st.CreateStep(ctx, store.Step{ + ID: uuid.NewString(), JobID: job.ID, Name: "step", + Status: store.StepStatusRunning, CreatedAt: now, UpdatedAt: now, + }) + if err != nil { + t.Fatalf("CreateStep: %v", err) + } + task, err := st.CreateTask(ctx, store.Task{ + ID: uuid.NewString(), + JobID: job.ID, + StepID: step.ID, + Name: "task", + Status: store.TaskStatusRunning, + AssignedWorkerID: workerID, + CreatedAt: now, + UpdatedAt: now, + }) + if err != nil { + t.Fatalf("CreateTask: %v", err) + } + return job, step, task +} + +// ── task.status ──────────────────────────────────────────────────────────── + +// TestProvenance_StatusFromWrongWorker proves that a status message whose +// subject names worker A cannot complete a task held by worker B. +// +// Before this check, the handler read WorkerID only to log it. The only thing +// standing between an attacker and a forged completion was that attempt_id is +// an unguessable UUID — capability-by-obscurity, not authorization. +func TestProvenance_StatusFromWrongWorker(t *testing.T) { + st := fake.New() + ctx := t.Context() + + // Two workers, a task held by B, and a live attempt for it. + const workerA, workerB = "worker-a", "worker-b" + job, _, task := seedRunnableTask(t, st, workerB) // helper above + attempt, err := st.CreateTaskAttempt(ctx, store.TaskAttempt{ + ID: uuid.NewString(), + TaskID: task.ID, + WorkerID: workerB, + StartedAt: time.Now().UTC(), + }) + if err != nil { + t.Fatalf("CreateTaskAttempt: %v", err) + } + + s := newTestScheduler(st, &stubBus{}) // existing helper in this package + s.ctx = ctx + + // Worker A forges a completion on its OWN subject — which is the only + // subject NATS would let it publish on — for a task it does not hold. + forged := protocol.TaskStatusMsg{ + Version: protocol.ProtocolVersion, + Type: protocol.TypeTaskStatus, + TaskID: task.ID, + AttemptID: attempt.ID, + JobID: job.ID, + Status: "succeeded", + WorkerID: workerB, // the payload lies; the subject does not + At: time.Now().UTC(), + } + msg := newFakeJetStreamMsg( + t, + bus.TaskStatusSubject(workerA, job.ID), + mustMarshal(t, forged), + ) + + s.handleTaskStatusMessage(msg) + + // The task must be untouched. + got, err := st.GetTask(ctx, task.ID) + if err != nil { + t.Fatalf("GetTask: %v", err) + } + if got.Status == store.TaskStatusSucceeded { + t.Fatal("forged status was applied: worker A completed worker B's task") + } + if got.Status != task.Status { + t.Errorf("task status changed to %q, want unchanged %q", got.Status, task.Status) + } + // And the message must be acked, not redelivered forever. + if !msg.acked { + t.Error("forged message was not acked; it will redeliver in a loop") + } +} + +// ── task.logs ────────────────────────────────────────────────────────────── + +// TestProvenance_LogsFromWrongWorker proves that a log chunk whose subject +// names worker A is not persisted against a task held by worker B. +// +// LogChunkMsg carries no worker-identity field at all — the subject is the +// ONLY identity available for this channel — so this is the case where +// dropping the subject check would leave zero provenance signal whatsoever. +func TestProvenance_LogsFromWrongWorker(t *testing.T) { + st := fake.New() + ctx := t.Context() + + const workerA, workerB = "worker-a", "worker-b" + _, _, task := seedRunnableTask(t, st, workerB) + attempt, err := st.CreateTaskAttempt(ctx, store.TaskAttempt{ + ID: uuid.NewString(), + TaskID: task.ID, + WorkerID: workerB, + StartedAt: time.Now().UTC(), + }) + if err != nil { + t.Fatalf("CreateTaskAttempt: %v", err) + } + + s := newTestScheduler(st, &stubBus{}) + s.ctx = ctx + + forged := protocol.LogChunkMsg{ + TaskID: task.ID, + AttemptID: attempt.ID, + SeqNum: 1, + At: time.Now().UTC(), + Stream: "stdout", + Data: "injected by worker A", + } + msg := newFakeJetStreamMsg( + t, + bus.TaskLogsSubject(workerA, task.ID), + mustMarshal(t, forged), + ) + + s.handleLogChunk(msg) + + // afterNATSSeq=-1 (not 0): the fake message carries no NATS sequence + // metadata, so a persisted row would land at NATSSeq=0, and the store + // filters strictly greater than afterNATSSeq — 0 would silently exclude + // it and pass for the wrong reason regardless of whether the row exists. + logs, err := st.ListTaskLogs(ctx, attempt.ID, -1, 100) + if err != nil { + t.Fatalf("ListTaskLogs: %v", err) + } + if len(logs) != 0 { + t.Fatalf("forged log chunk was persisted: %d rows, want 0", len(logs)) + } + if !msg.acked { + t.Error("forged message was not acked; it will redeliver in a loop") + } +} + +// TestProvenance_LogsForAnotherWorkersTask proves that a worker holding a +// live attempt of its own cannot pair that attempt with a DIFFERENT task in +// the payload to inject log content there. This is distinct from +// TestProvenance_LogsFromWrongWorker: here the subject and the attempt's +// WorkerID genuinely agree (worker A really does hold attemptA), so the +// worker-ID check alone would let it through. The broker grant is +// task.logs..* — any trailing leaf is a valid subject for worker A +// to publish on — so the leaf cannot be relied on either; only comparing the +// attempt's own TaskID against the payload's TaskID catches this. +func TestProvenance_LogsForAnotherWorkersTask(t *testing.T) { + st := fake.New() + ctx := t.Context() + + const workerA, workerB = "worker-a", "worker-b" + _, _, taskA := seedRunnableTask(t, st, workerA) + attemptA, err := st.CreateTaskAttempt(ctx, store.TaskAttempt{ + ID: uuid.NewString(), + TaskID: taskA.ID, + WorkerID: workerA, + StartedAt: time.Now().UTC(), + }) + if err != nil { + t.Fatalf("CreateTaskAttempt(A): %v", err) + } + // Worker B's own task — uninvolved in the publish below except as the + // target the forged payload names. + _, _, taskB := seedRunnableTask(t, st, workerB) + + s := newTestScheduler(st, &stubBus{}) + s.ctx = ctx + + forged := protocol.LogChunkMsg{ + TaskID: taskB.ID, // names worker B's task… + AttemptID: attemptA.ID, // …paired with worker A's own, genuinely-held attempt + SeqNum: 1, + At: time.Now().UTC(), + Stream: "stdout", + Data: "injected into B's task via A's own attempt", + } + msg := newFakeJetStreamMsg( + t, + bus.TaskLogsSubject(workerA, "anything"), // any leaf is a valid subject for A + mustMarshal(t, forged), + ) + + s.handleLogChunk(msg) + + logs, err := st.ListTaskLogs(ctx, attemptA.ID, -1, 100) + if err != nil { + t.Fatalf("ListTaskLogs: %v", err) + } + if len(logs) != 0 { + t.Fatalf("forged log chunk was persisted against another worker's task: %d rows, want 0", len(logs)) + } + if !msg.acked { + t.Error("forged message was not acked; it will redeliver in a loop") + } +} + +// ── worker.register ──────────────────────────────────────────────────────── + +// TestProvenance_RegisterOfAnotherWorker proves that worker A cannot +// overwrite worker B's row — CPU count, tags, EXPR caps and all — by +// publishing a registration on its own subject with B's ID in the payload. +func TestProvenance_RegisterOfAnotherWorker(t *testing.T) { + st := fake.New() + ctx := t.Context() + + const workerA, workerB = "worker-a", "worker-b" + now := time.Now().UTC() + if _, err := st.RegisterWorker(ctx, store.Worker{ + ID: workerB, FarmID: "farm-1", Hostname: "real-host", CPUCount: 8, + Status: store.WorkerStatusOnline, LastHeartbeatAt: &now, + Tags: map[string]string{"gpu": "true"}, + }); err != nil { + t.Fatalf("RegisterWorker(B): %v", err) + } + + s := newMetricsScheduler(st, &recordBus{}, "") + + forged := protocol.RegisterMsg{ + Version: protocol.ProtocolVersion, + Type: protocol.TypeRegister, + WorkerID: workerB, // the payload lies; the subject does not + FarmID: "farm-1", + Hostname: "forged-host", + OS: "linux", + CPUCount: 1, + } + msg := newFakeJetStreamMsg( + t, + bus.WorkerRegisterSubject(workerA), + mustMarshal(t, forged), + ) + + s.handleWorkerMessage(msg) + + w, err := st.GetWorker(ctx, workerB) + if err != nil { + t.Fatalf("GetWorker(B): %v", err) + } + if w.Hostname != "real-host" || w.CPUCount != 8 { + t.Fatalf("worker B's row was overwritten by worker A's forged registration: hostname=%q cpu=%d", + w.Hostname, w.CPUCount) + } + if _, err := st.GetWorker(ctx, workerA); err == nil { + t.Error("worker A should not have been registered from a mismatched payload either") + } + if !msg.acked { + t.Error("forged message was not acked; it will redeliver in a loop") + } +} + +// ── worker.heartbeat ─────────────────────────────────────────────────────── + +// TestProvenance_HeartbeatOfAnotherWorker proves that worker A cannot keep +// worker B looking alive — and therefore un-reclaimed — by publishing a +// heartbeat on its own subject with B's ID in the payload. +func TestProvenance_HeartbeatOfAnotherWorker(t *testing.T) { + st := fake.New() + ctx := t.Context() + + const workerA, workerB = "worker-a", "worker-b" + stale := time.Now().UTC().Add(-time.Hour) + if _, err := st.RegisterWorker(ctx, store.Worker{ + ID: workerB, FarmID: "farm-1", Status: store.WorkerStatusOnline, LastHeartbeatAt: &stale, + }); err != nil { + t.Fatalf("RegisterWorker(B): %v", err) + } + + s := newMetricsScheduler(st, &recordBus{}, "") + + forged := protocol.HeartbeatMsg{ + Version: protocol.ProtocolVersion, + Type: protocol.TypeHeartbeat, + WorkerID: workerB, // the payload lies; the subject does not + At: time.Now().UTC(), + } + msg := newFakeJetStreamMsg( + t, + bus.WorkerHeartbeatSubject(workerA), + mustMarshal(t, forged), + ) + + s.handleWorkerMessage(msg) + + w, err := st.GetWorker(ctx, workerB) + if err != nil { + t.Fatalf("GetWorker(B): %v", err) + } + if w.LastHeartbeatAt == nil || !w.LastHeartbeatAt.Equal(stale) { + t.Fatalf("worker B's heartbeat was refreshed by worker A's forged heartbeat: LastHeartbeatAt = %v, want unchanged %v", + w.LastHeartbeatAt, stale) + } + if !msg.acked { + t.Error("forged message was not acked; it will redeliver in a loop") + } +} + +// ── worker.deregister ───────────────────────────────────────────────────── + +// TestProvenance_DeregisterOfAnotherWorker proves that worker A cannot mark +// worker B offline by publishing on its own subject with B's ID in the +// payload — a denial-of-service against the farm otherwise available to any +// worker that can reach the broker. +func TestProvenance_DeregisterOfAnotherWorker(t *testing.T) { + st := fake.New() + ctx := t.Context() + + const workerA, workerB = "worker-a", "worker-b" + now := time.Now().UTC() + if _, err := st.RegisterWorker(ctx, store.Worker{ + ID: workerA, FarmID: "farm-1", Status: store.WorkerStatusOnline, LastHeartbeatAt: &now, + }); err != nil { + t.Fatalf("RegisterWorker(A): %v", err) + } + if _, err := st.RegisterWorker(ctx, store.Worker{ + ID: workerB, FarmID: "farm-1", Status: store.WorkerStatusOnline, LastHeartbeatAt: &now, + }); err != nil { + t.Fatalf("RegisterWorker(B): %v", err) + } + + s := newMetricsScheduler(st, &recordBus{}, "") + + forged := struct { + WorkerID string `json:"worker_id"` + Reason string `json:"reason,omitempty"` + }{WorkerID: workerB, Reason: "forged shutdown"} + + msg := newFakeJetStreamMsg( + t, + bus.WorkerDeregisterSubject(workerA), + mustMarshal(t, forged), + ) + + s.handleWorkerMessage(msg) + + w, err := st.GetWorker(ctx, workerB) + if err != nil { + t.Fatalf("GetWorker(B): %v", err) + } + if w.Status != store.WorkerStatusOnline { + t.Fatalf("worker B was marked %q by worker A's forged deregister, want online", w.Status) + } + if !msg.acked { + t.Error("forged message was not acked; it will redeliver in a loop") + } +} + +// ── work.lease ───────────────────────────────────────────────────────────── + +// TestProvenance_LeaseAsAnotherWorker proves that a lease request published +// on worker A's subject cannot have its assignments credited to worker B: the +// reply must be empty and no task may end up assigned to B. +func TestProvenance_LeaseAsAnotherWorker(t *testing.T) { + st := fake.New() + one := 1 + // seedLeaseFixture (lease_test.go) registers a real, eligible worker "w1" + // with a ready task — workerB below is that real worker, the one being + // impersonated. + _, taskIDs := seedLeaseFixture(t, st, []*int{&one}) + + const workerA, workerB = "worker-a", "w1" + + // Worker A is ALSO real, eligible and in the same farm/queue, so an + // unforged request from it would legitimately be offered this same ready + // task. That is what pins the assertion below to the check: if the reply + // were empty merely because worker A didn't exist or wasn't eligible, + // removing the check wouldn't turn this test red. + now := time.Now().UTC() + if _, err := st.RegisterWorker(t.Context(), store.Worker{ + ID: workerA, FarmID: "f1", Hostname: workerA, + Status: store.WorkerStatusOnline, CPUCount: 4, LastHeartbeatAt: &now, + Tags: map[string]string{}, + }); err != nil { + t.Fatalf("RegisterWorker(A): %v", err) + } + + s := newMetricsScheduler(st, &recordBus{}, "f1") + + req, err := json.Marshal(leaseRequest{WorkerID: workerB}) + if err != nil { + t.Fatalf("marshal lease request: %v", err) + } + + reply := s.handleLeaseRequest(workerA, "q1", req) + + var got leaseReply + if err := json.Unmarshal(reply, &got); err != nil { + t.Fatalf("unmarshal reply: %v", err) + } + if len(got.Assignments) != 0 { + t.Fatalf("forged lease request returned %d assignments, want 0", len(got.Assignments)) + } + + for _, id := range taskIDs { + task, err := st.GetTask(t.Context(), id) + if err != nil { + t.Fatalf("GetTask(%s): %v", id, err) + } + if task.AssignedWorkerID != "" { + t.Fatalf("task %s was assigned to %q by a refused lease request, want unassigned", id, task.AssignedWorkerID) + } + } +} diff --git a/internal/scheduler/registry_test.go b/internal/scheduler/registry_test.go index 88f388ee..21b881b1 100644 --- a/internal/scheduler/registry_test.go +++ b/internal/scheduler/registry_test.go @@ -13,15 +13,18 @@ package scheduler import ( "context" "encoding/json" + "log/slog" "testing" "time" "github.com/google/uuid" "github.com/uberware/sqi/internal/bus" + "github.com/uberware/sqi/internal/metrics" "github.com/uberware/sqi/internal/store" "github.com/uberware/sqi/internal/store/fake" "github.com/uberware/sqi/internal/worker/protocol" + "github.com/uberware/sqi/internal/ws" ) // workerMsgJSON marshals any value to JSON bytes for a fakeJSMsg payload. @@ -41,7 +44,7 @@ func TestHandleWorkerRegister_Valid(t *testing.T) { s := newMetricsScheduler(st, &recordBus{}, "") msg := &fakeJSMsg{ - subject: bus.SubjectWorkerRegister, + subject: bus.WorkerRegisterSubject("w-1"), data: workerMsgJSON(t, protocol.RegisterMsg{ Version: protocol.ProtocolVersion, Type: protocol.TypeRegister, WorkerID: "w-1", FarmID: "farm-1", Name: "worker-2", Hostname: "node-1", OS: "linux", @@ -75,7 +78,7 @@ func TestHandleWorkerRegister_MalformedJSON_Acked(t *testing.T) { st := fake.New() s := newMetricsScheduler(st, &recordBus{}, "") - msg := &fakeJSMsg{subject: bus.SubjectWorkerRegister, data: []byte("{bad")} + msg := &fakeJSMsg{subject: bus.WorkerRegisterSubject("w-1"), data: []byte("{bad")} s.handleWorkerMessage(msg) if !msg.acked { @@ -83,18 +86,27 @@ func TestHandleWorkerRegister_MalformedJSON_Acked(t *testing.T) { } } -func TestHandleWorkerRegister_MissingWorkerID_Acked(t *testing.T) { +// TestHandleWorkerRegister_EmptyPayloadWorkerID_Acked drives a payload with no +// worker_id at all against a real subject. There is no separate "missing +// worker_id" code path any more: bus.ParseWorkerSubject guarantees the +// subject's worker token is never empty, so an empty m.WorkerID always fails +// the subject/payload mismatch check first and is discarded through that +// branch. +func TestHandleWorkerRegister_EmptyPayloadWorkerID_Acked(t *testing.T) { st := fake.New() s := newMetricsScheduler(st, &recordBus{}, "") msg := &fakeJSMsg{ - subject: bus.SubjectWorkerRegister, + subject: bus.WorkerRegisterSubject("w-1"), data: workerMsgJSON(t, protocol.RegisterMsg{Version: protocol.ProtocolVersion, WorkerID: "", FarmID: "farm-1"}), } s.handleWorkerMessage(msg) if !msg.acked { - t.Error("register missing worker_id should be acked") + t.Error("register with an empty payload worker_id should be acked (mismatch)") + } + if _, err := st.GetWorker(t.Context(), "w-1"); err == nil { + t.Error("no worker should have been registered from a mismatched payload") } } @@ -112,7 +124,7 @@ func TestHandleWorkerRegister_StoreError_Nacked(t *testing.T) { s := newMetricsScheduler(st, &recordBus{}, "") msg := &fakeJSMsg{ - subject: bus.SubjectWorkerRegister, + subject: bus.WorkerRegisterSubject("w-1"), data: workerMsgJSON(t, protocol.RegisterMsg{Version: protocol.ProtocolVersion, WorkerID: "w-1", FarmID: "farm-1"}), } s.handleWorkerMessage(msg) @@ -125,6 +137,112 @@ func TestHandleWorkerRegister_StoreError_Nacked(t *testing.T) { } } +// touchRecordingStore wraps a real store and records every +// TouchWorkerCredential call, so a test can prove registration does or does +// not reach it without depending on timing. +type touchRecordingStore struct { + store.Store + + touched []string +} + +func (s *touchRecordingStore) TouchWorkerCredential(ctx context.Context, workerID string, at time.Time) error { + s.touched = append(s.touched, workerID) + return s.Store.TouchWorkerCredential(ctx, workerID, at) +} + +// TestHandleWorkerRegister_TouchesActiveCredential_WhenAuthEnabled proves +// that registering a worker with an active broker credential sets +// LastSeenAt, and only when broker authentication is enabled. +func TestHandleWorkerRegister_TouchesActiveCredential_WhenAuthEnabled(t *testing.T) { + fk := fake.New() + if _, err := fk.CreateWorkerCredential(t.Context(), store.WorkerCredential{ + ID: uuid.NewString(), WorkerID: "w-1", PublicKey: "pub1", EnrolledAt: time.Now().UTC(), + }); err != nil { + t.Fatalf("seed CreateWorkerCredential: %v", err) + } + st := &touchRecordingStore{Store: fk} + + cfg := DefaultConfig() + cfg.NATSAuthEnabled = true + s := New(cfg, st, &recordBus{}, metrics.New(), slog.New(slog.DiscardHandler), ws.NoopNotifier{}, nil) + s.ctx = context.Background() + + msg := &fakeJSMsg{ + subject: bus.WorkerRegisterSubject("w-1"), + data: workerMsgJSON(t, protocol.RegisterMsg{ + Version: protocol.ProtocolVersion, Type: protocol.TypeRegister, + WorkerID: "w-1", FarmID: "farm-1", Hostname: "node-1", OS: "linux", + }), + } + s.handleWorkerMessage(msg) + + if !msg.acked { + t.Error("valid register should be acked") + } + if len(st.touched) != 1 || st.touched[0] != "w-1" { + t.Errorf("touched = %v, want exactly one call for w-1", st.touched) + } + cred, err := fk.GetActiveWorkerCredentialByWorkerID(t.Context(), "w-1") + if err != nil { + t.Fatalf("GetActiveWorkerCredentialByWorkerID: %v", err) + } + if cred.LastSeenAt == nil { + t.Error("expected LastSeenAt to be set after registration") + } +} + +// TestHandleWorkerRegister_NoTouchCall_WhenAuthDisabled asserts the +// auth-off default path does no extra store work: no credential rows exist +// on an auth-off farm, and the touch call must not even be attempted. +func TestHandleWorkerRegister_NoTouchCall_WhenAuthDisabled(t *testing.T) { + st := &touchRecordingStore{Store: fake.New()} + s := newMetricsScheduler(st, &recordBus{}, "") // DefaultConfig: NATSAuthEnabled false + + msg := &fakeJSMsg{ + subject: bus.WorkerRegisterSubject("w-1"), + data: workerMsgJSON(t, protocol.RegisterMsg{ + Version: protocol.ProtocolVersion, Type: protocol.TypeRegister, + WorkerID: "w-1", FarmID: "farm-1", Hostname: "node-1", OS: "linux", + }), + } + s.handleWorkerMessage(msg) + + if !msg.acked { + t.Error("valid register should be acked") + } + if len(st.touched) != 0 { + t.Errorf("touched = %v, want no calls with broker auth disabled", st.touched) + } +} + +// TestHandleWorkerRegister_NoActiveCredential_StillAcked proves a missing +// credential (store.ErrNotFound) never fails registration: the message is +// still acked and the worker is still registered. +func TestHandleWorkerRegister_NoActiveCredential_StillAcked(t *testing.T) { + st := fake.New() + cfg := DefaultConfig() + cfg.NATSAuthEnabled = true + s := New(cfg, st, &recordBus{}, metrics.New(), slog.New(slog.DiscardHandler), ws.NoopNotifier{}, nil) + s.ctx = context.Background() + + msg := &fakeJSMsg{ + subject: bus.WorkerRegisterSubject("w-1"), + data: workerMsgJSON(t, protocol.RegisterMsg{ + Version: protocol.ProtocolVersion, Type: protocol.TypeRegister, + WorkerID: "w-1", FarmID: "farm-1", Hostname: "node-1", OS: "linux", + }), + } + s.handleWorkerMessage(msg) + + if !msg.acked { + t.Error("register should be acked even with no active credential to touch") + } + if _, err := st.GetWorker(t.Context(), "w-1"); err != nil { + t.Errorf("worker should still be registered: %v", err) + } +} + // ── handleWorkerHeartbeat ───────────────────────────────────────────────────── func TestHandleWorkerHeartbeat_Valid(t *testing.T) { @@ -140,7 +258,7 @@ func TestHandleWorkerHeartbeat_Valid(t *testing.T) { hbAt := now.Add(5 * time.Second) msg := &fakeJSMsg{ - subject: bus.SubjectWorkerHeartbeat, + subject: bus.WorkerHeartbeatSubject("w-1"), data: workerMsgJSON(t, protocol.HeartbeatMsg{Version: protocol.ProtocolVersion, WorkerID: "w-1", At: hbAt}), } s.handleWorkerMessage(msg) @@ -170,7 +288,7 @@ func TestHandleWorkerHeartbeat_ZeroAt_UsesServerTime(t *testing.T) { } msg := &fakeJSMsg{ - subject: bus.SubjectWorkerHeartbeat, + subject: bus.WorkerHeartbeatSubject("w-1"), data: workerMsgJSON(t, protocol.HeartbeatMsg{Version: protocol.ProtocolVersion, WorkerID: "w-1"}), // zero At } s.handleWorkerMessage(msg) @@ -189,7 +307,7 @@ func TestHandleWorkerHeartbeat_UnknownWorker_Nacked(t *testing.T) { s := newMetricsScheduler(st, &recordBus{}, "") msg := &fakeJSMsg{ - subject: bus.SubjectWorkerHeartbeat, + subject: bus.WorkerHeartbeatSubject("ghost"), data: workerMsgJSON(t, protocol.HeartbeatMsg{Version: protocol.ProtocolVersion, WorkerID: "ghost", At: time.Now()}), } s.handleWorkerMessage(msg) @@ -199,13 +317,18 @@ func TestHandleWorkerHeartbeat_UnknownWorker_Nacked(t *testing.T) { } } -func TestHandleWorkerHeartbeat_MalformedAndMissingID_Acked(t *testing.T) { +// TestHandleWorkerHeartbeat_MalformedAndEmptyPayloadID_Acked covers a +// malformed body and a body with an empty worker_id. The latter has no +// separate "missing worker_id" code path any more: bus.ParseWorkerSubject +// guarantees the subject's worker token is never empty, so an empty +// m.WorkerID always fails the subject/payload mismatch check first. +func TestHandleWorkerHeartbeat_MalformedAndEmptyPayloadID_Acked(t *testing.T) { tests := []struct { name string data []byte }{ {"malformed", []byte("{bad")}, - {"missing id", nil}, // filled below + {"empty payload id (mismatch)", nil}, // filled below } tests[1].data = workerMsgJSON(t, protocol.HeartbeatMsg{Version: protocol.ProtocolVersion, WorkerID: ""}) @@ -213,7 +336,7 @@ func TestHandleWorkerHeartbeat_MalformedAndMissingID_Acked(t *testing.T) { t.Run(tt.name, func(t *testing.T) { st := fake.New() s := newMetricsScheduler(st, &recordBus{}, "") - msg := &fakeJSMsg{subject: bus.SubjectWorkerHeartbeat, data: tt.data} + msg := &fakeJSMsg{subject: bus.WorkerHeartbeatSubject("w-1"), data: tt.data} s.handleWorkerMessage(msg) if !msg.acked { t.Errorf("%s heartbeat should be acked (discarded)", tt.name) @@ -236,7 +359,7 @@ func TestHandleWorkerDeregister_Valid(t *testing.T) { } msg := &fakeJSMsg{ - subject: bus.SubjectWorkerDeregister, + subject: bus.WorkerDeregisterSubject("w-1"), data: workerMsgJSON(t, map[string]string{"worker_id": "w-1", "reason": "shutdown"}), } s.handleWorkerMessage(msg) @@ -302,7 +425,7 @@ func TestHandleWorkerDeregister_ReclaimsInFlightTasks(t *testing.T) { } msg := &fakeJSMsg{ - subject: bus.SubjectWorkerDeregister, + subject: bus.WorkerDeregisterSubject(workerID), data: workerMsgJSON(t, map[string]string{"worker_id": workerID, "reason": "shutdown"}), } s.handleWorkerMessage(msg) @@ -331,7 +454,7 @@ func TestHandleWorkerDeregister_UnknownWorker_Acked(t *testing.T) { s := newMetricsScheduler(st, &recordBus{}, "") msg := &fakeJSMsg{ - subject: bus.SubjectWorkerDeregister, + subject: bus.WorkerDeregisterSubject("ghost"), data: workerMsgJSON(t, map[string]string{"worker_id": "ghost"}), } s.handleWorkerMessage(msg) @@ -341,19 +464,24 @@ func TestHandleWorkerDeregister_UnknownWorker_Acked(t *testing.T) { } } -func TestHandleWorkerDeregister_MalformedAndMissingID_Acked(t *testing.T) { +// TestHandleWorkerDeregister_MalformedAndEmptyPayloadID_Acked covers a +// malformed body and a body with an empty worker_id. The latter has no +// separate "missing worker_id" code path any more: bus.ParseWorkerSubject +// guarantees the subject's worker token is never empty, so an empty +// m.WorkerID always fails the subject/payload mismatch check first. +func TestHandleWorkerDeregister_MalformedAndEmptyPayloadID_Acked(t *testing.T) { tests := []struct { name string data []byte }{ {"malformed", []byte("{bad")}, - {"missing id", workerMsgJSON(t, map[string]string{"worker_id": ""})}, + {"empty payload id (mismatch)", workerMsgJSON(t, map[string]string{"worker_id": ""})}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { st := fake.New() s := newMetricsScheduler(st, &recordBus{}, "") - msg := &fakeJSMsg{subject: bus.SubjectWorkerDeregister, data: tt.data} + msg := &fakeJSMsg{subject: bus.WorkerDeregisterSubject("w-1"), data: tt.data} s.handleWorkerMessage(msg) if !msg.acked { t.Errorf("%s deregister should be acked", tt.name) @@ -372,7 +500,7 @@ func TestRegistration_AutoRegistersComputeLocation(t *testing.T) { // Case 1: end-to-end — register a worker with a new location via // handleWorkerMessage; assert the entity is created in the store. msg := &fakeJSMsg{ - subject: bus.SubjectWorkerRegister, + subject: bus.WorkerRegisterSubject("w-loc-1"), data: workerMsgJSON(t, protocol.RegisterMsg{ Version: protocol.ProtocolVersion, Type: protocol.TypeRegister, WorkerID: "w-loc-1", FarmID: "farm-1", Hostname: "n1", OS: "linux", @@ -429,7 +557,7 @@ func TestRegistration_EnsureComputeLocation_StoreError(t *testing.T) { s := newMetricsScheduler(st, &recordBus{}, "") msg := &fakeJSMsg{ - subject: bus.SubjectWorkerRegister, + subject: bus.WorkerRegisterSubject("w-loc-err"), data: workerMsgJSON(t, protocol.RegisterMsg{ Version: protocol.ProtocolVersion, Type: protocol.TypeRegister, WorkerID: "w-loc-err", FarmID: "farm-1", Hostname: "n1", OS: "linux", @@ -482,7 +610,7 @@ func TestRegistration_EnsureComputeLocation_LookupError(t *testing.T) { s := newMetricsScheduler(st, &recordBus{}, "") msg := &fakeJSMsg{ - subject: bus.SubjectWorkerRegister, + subject: bus.WorkerRegisterSubject("w-loc-err2"), data: workerMsgJSON(t, protocol.RegisterMsg{ Version: protocol.ProtocolVersion, Type: protocol.TypeRegister, WorkerID: "w-loc-err2", FarmID: "farm-1", Hostname: "n1", OS: "linux", @@ -522,7 +650,7 @@ func TestHandleWorkerMessage_RegisterID(t *testing.T) { id := uuid.NewString() msg := &fakeJSMsg{ - subject: bus.SubjectWorkerRegister, + subject: bus.WorkerRegisterSubject(id), data: workerMsgJSON(t, protocol.RegisterMsg{Version: protocol.ProtocolVersion, WorkerID: id, FarmID: "farm-1"}), } s.handleWorkerMessage(msg) diff --git a/internal/scheduler/scheduler.go b/internal/scheduler/scheduler.go index 6cc0d4bd..6de77ab9 100644 --- a/internal/scheduler/scheduler.go +++ b/internal/scheduler/scheduler.go @@ -14,7 +14,7 @@ // payloads. When no work is available the request parks in the waiter // registry until new work appears or the hold elapses, then replies. // -// 2. Worker registry: a NATS push-consumer for worker.register messages that +// 2. Worker registry: a NATS push-consumer for worker.register. messages that // persists capability data via [store.WorkerStore.RegisterWorker] and keeps // the WorkersTotal Prometheus gauge current. // @@ -45,10 +45,10 @@ // // Status and log ingestion. A push-consumer on the SQI_TASK stream // ([handleTaskStatusMessage]) decodes [protocol.TaskStatusMsg] from -// task.status., updates the task/attempt, releases held usage pool slots, and +// task.status.., updates the task/attempt, releases held usage pool slots, and // drives step/job completion including [openjd.ResolveDependencies] for // multi-step jobs. A push-consumer on SQI_LOGS ([handleLogChunk]) persists each -// task.logs. chunk as a [store.TaskLog] row, recording both the +// task.logs.. chunk as a [store.TaskLog] row, recording both the // worker-assigned sequence number and the NATS stream sequence that serves as // the log-tail pagination cursor. // @@ -203,6 +203,14 @@ type Config struct { // // Zero fields normalize to the defaults in [New]. ExprLimits openjd.ExprLimits + + // NATSAuthEnabled mirrors config.NATSAuthConfig.Enabled (server.Config's + // NATSAuthEnabled). It gates one thing here: whether worker registration + // touches the worker's broker-credential LastSeenAt. With broker + // authentication off there are no credential rows at all, and the + // default no-config path must do no extra store work — see + // handleWorkerRegister. + NATSAuthEnabled bool } // busClient is the subset of [bus.Client] used by the Scheduler. Defined as @@ -213,7 +221,7 @@ type busClient interface { ConsumeTaskLogs(ctx context.Context, handler jetstream.MessageHandler) (jetstream.ConsumeContext, error) PublishTaskCancel(ctx context.Context, taskID string, data []byte) error SubscribeWorkerDiag(handler func(subject string, data []byte)) (*nats.Subscription, error) - SubscribeLease(handler func(queueID string, data []byte) []byte) (*nats.Subscription, error) + SubscribeLease(handler func(workerID, queueID string, data []byte) []byte) (*nats.Subscription, error) } // Scheduler owns the assignment loop, worker registry, and heartbeat sweep. @@ -241,6 +249,11 @@ type Scheduler struct { // waiters parks long-poll lease requests per queue; woken by wake triggers. waiters *waiterRegistry + // attemptCache holds recently-seen task-attempt ownership (workerID, + // taskID), consulted by handleLogChunk before it reads the store. See + // [attemptOwnerCache]. + attemptCache *attemptOwnerCache + // leaseLocks serializes lease selection per worker. Concurrent lease // requests for the SAME worker (one outstanding request per queue it // serves, plus retry overlap) must not both read the same committed-core @@ -329,6 +342,7 @@ func New(cfg Config, st store.Store, busClient busClient, m *metrics.Metrics, lo notifier: n, diagBuf: diagBuf, waiters: newWaiterRegistry(), + attemptCache: newAttemptOwnerCache(), leaseHoldTimeout: 30 * time.Second, retryWakeTimers: make(map[*time.Timer]struct{}), // ctx is overwritten with the derived cancellable context in Run. @@ -367,8 +381,8 @@ func (s *Scheduler) Run(ctx context.Context) error { ) // ── Worker NATS consumer ──────────────────────────────── - // A single JetStream push-consumer delivers both worker.register and - // worker.heartbeat messages. The handler dispatches by subject. + // A single JetStream push-consumer delivers the worker register, + // heartbeat and deregister messages. The handler dispatches by subject. _, err := s.bus.ConsumeWorker(ctx, s.handleWorkerMessage) if err != nil { return fmt.Errorf("scheduler: start worker consumer: %w", err) @@ -376,7 +390,7 @@ func (s *Scheduler) Run(ctx context.Context) error { s.logger.InfoContext(ctx, "scheduler: worker consumer started") // ── Task-status NATS consumer ──────────────────────────────── - // A JetStream push-consumer on SQI_TASK delivers task.status. + // A JetStream push-consumer on SQI_TASK delivers task.status.. // messages from workers. handleTaskStatusMessage updates the store, // closes attempt records, releases usage pool slots, and drives step/job // completion. @@ -386,7 +400,7 @@ func (s *Scheduler) Run(ctx context.Context) error { s.logger.InfoContext(ctx, "scheduler: task-status consumer started") // ── Task-logs NATS consumer ────────────────────────────────── - // A JetStream push-consumer on SQI_LOGS delivers task.logs. + // A JetStream push-consumer on SQI_LOGS delivers task.logs.. // messages from workers. handleLogChunk persists each chunk to the // task_logs table with NATS sequence as the pagination cursor. if err := s.startTaskLogsConsumer(ctx); err != nil { @@ -504,6 +518,10 @@ func (s *Scheduler) createAttemptAndClaimUsage( s.revertTaskToReady(ctx, task.ID, "attempt creation error") return store.TaskAttempt{}, fmt.Errorf("create task attempt for task %s: %w", task.ID, err) } + // The scheduler already knows both fields the log-ingest path needs, so + // populate the cache now rather than waiting for the first log chunk to + // pay for a store read. + s.attemptCache.put(attempt.ID, attempt.WorkerID, attempt.TaskID) // Re-check pool availability and create claim rows inside a single DB // transaction so no concurrent assignment can over-subscribe a pool. @@ -514,6 +532,10 @@ func (s *Scheduler) createAttemptAndClaimUsage( if err := s.store.TryClaimSlots(ctx, attempt.ID, claims, now); err != nil { s.revertTaskToReady(ctx, task.ID, "usage claim error") + // The attempt row survives this failure with no terminal status ever + // coming for it, so nothing else would evict its cache entry. Drop it + // now rather than let it sit as a stale, never-reused hit. + s.attemptCache.evict(attempt.ID) if errors.Is(err, store.ErrUsageAtCapacity) { s.logger.DebugContext( ctx, "scheduler: usage pool at capacity — deferring assignment", @@ -660,26 +682,39 @@ func (s *Scheduler) ReleaseTaskUsage(ctx context.Context, attemptID string) erro // ── Worker NATS consumer ───────────────────────────────────────────── -// handleWorkerMessage is the JetStream message handler for both -// worker.register and worker.heartbeat subjects (both flow through the +// handleWorkerMessage is the JetStream message handler for the worker +// register, heartbeat and deregister subjects (all three flow through the // SQI_WORKER stream and its single durable consumer). +// +// Each subject carries the publishing worker's ID as its last token, so the +// routing below recovers that identity before dispatching to the per-message +// handler. func (s *Scheduler) handleWorkerMessage(msg jetstream.Msg) { ctx := s.ctx subject := msg.Subject() - switch subject { - case bus.SubjectWorkerRegister: - s.handleWorkerRegister(ctx, msg) - case bus.SubjectWorkerHeartbeat: - s.handleWorkerHeartbeat(ctx, msg) - case bus.SubjectWorkerDeregister: - s.handleWorkerDeregister(ctx, msg) + workerID, _, ok := bus.ParseWorkerSubject(subject) + if !ok { + s.discardUnexpectedSubject(ctx, msg, "worker") + return + } + + switch { + case strings.HasPrefix(subject, bus.SubjectWorkerRegisterPrefix+"."): + s.handleWorkerRegister(ctx, msg, workerID) + case strings.HasPrefix(subject, bus.SubjectWorkerHeartbeatPrefix+"."): + s.handleWorkerHeartbeat(ctx, msg, workerID) + case strings.HasPrefix(subject, bus.SubjectWorkerDeregisterPrefix+"."): + s.handleWorkerDeregister(ctx, msg, workerID) default: - s.logger.WarnContext( - ctx, "scheduler: unexpected worker subject", - slog.String("subject", subject), - ) - s.ackMsg(ctx, msg) + // Defense in depth, not dead code: ParseWorkerSubject also accepts + // the four-token task.status/task.logs/work.lease shapes, so + // widening SQI_WORKER's subject filter to one of those would reach + // here rather than being mis-dispatched. A three-token worker.* + // subject with an unrecognized prefix cannot: ParseWorkerSubject + // whitelists the three it knows and rejects the rest, so that case + // is already handled by the !ok branch above. + s.discardUnexpectedSubject(ctx, msg, "worker") } } @@ -738,14 +773,57 @@ func (s *Scheduler) discardOnVersionMismatch( return true } +// discardUnexpectedSubject logs a " on unexpected subject — +// discarding" warning naming msg's subject and acks it away. noun identifies +// the calling consumer (e.g. "worker", "task.status", "task.logs") in the log +// line. Shared by every consumer that gives up on a message because its +// subject would not parse, or — for [Scheduler.handleWorkerMessage]'s +// unreachable-in-practice default case — parsed into a shape that consumer's +// dispatch does not recognize. A malformed or unrecognized subject cannot +// become valid on redelivery, so this always acks, never naks. +func (s *Scheduler) discardUnexpectedSubject(ctx context.Context, msg jetstream.Msg, noun string) { + s.logger.WarnContext( + ctx, "scheduler: "+noun+" on unexpected subject — discarding", + slog.String("subject", msg.Subject()), + ) + s.ackMsg(ctx, msg) +} + +// discardOnIdentityMismatch acks and discards msg if payloadWorkerID (from +// the decoded message body) disagrees with subjectWorkerID (from the NATS +// subject). The subject is the only identity NATS itself can enforce, so a +// payload claiming a different worker is treated as permanent — redelivery +// cannot make a forged or stale identity legal — the same reasoning as +// [Scheduler.discardOnVersionMismatch]. +// +// No separate "missing worker_id" check is needed: subjectWorkerID is always +// non-empty ([bus.ParseWorkerSubject] guarantees it), so an empty +// payloadWorkerID already fails this comparison. +func (s *Scheduler) discardOnIdentityMismatch(ctx context.Context, msg jetstream.Msg, subjectWorkerID, payloadWorkerID string) bool { + if payloadWorkerID == subjectWorkerID { + return false + } + s.logger.WarnContext( + ctx, "scheduler: worker message whose payload identity differs from its subject — discarding", + slog.String("subject_worker_id", subjectWorkerID), + slog.String("payload_worker_id", payloadWorkerID), + ) + s.ackMsg(ctx, msg) + return true +} + // handleWorkerRegister processes a worker.register message: // decodes the payload, upserts the worker in the store, and refreshes the // WorkersTotal Prometheus gauge. -func (s *Scheduler) handleWorkerRegister(ctx context.Context, msg jetstream.Msg) { +// +// subjectWorkerID is the worker the message's subject attributes it to. +func (s *Scheduler) handleWorkerRegister(ctx context.Context, msg jetstream.Msg, subjectWorkerID string) { var m protocol.RegisterMsg if err := json.Unmarshal(msg.Data(), &m); err != nil { + // The subject is the only identity left once the body will not decode. s.logger.WarnContext( ctx, "scheduler: malformed worker.register message", + slog.String("subject_worker_id", subjectWorkerID), slog.Any("error", err), ) s.ackMsg(ctx, msg) // ack to discard; re-delivery cannot fix a bad payload @@ -755,9 +833,7 @@ func (s *Scheduler) handleWorkerRegister(ctx context.Context, msg jetstream.Msg) "this worker is not registered and will be offered no work at all") { return } - if m.WorkerID == "" { - s.logger.WarnContext(ctx, "scheduler: worker.register missing worker_id") - s.ackMsg(ctx, msg) + if s.discardOnIdentityMismatch(ctx, msg, subjectWorkerID, m.WorkerID) { return } @@ -805,6 +881,15 @@ func (s *Scheduler) handleWorkerRegister(ctx context.Context, msg jetstream.Msg) s.ensureComputeLocation(ctx, m.ComputeLocation) + // Registration is "last seen" — connect and reconnect both go through + // here, and it is low-frequency, unlike heartbeat. Only touched when + // broker authentication is enabled: with it off there are no credential + // rows at all, and the default no-config path must do no extra store + // work. + if s.cfg.NATSAuthEnabled { + s.touchWorkerCredential(ctx, m.WorkerID, now) + } + s.warnOnExprCapShortfall(ctx, w) s.logger.InfoContext( @@ -894,11 +979,37 @@ func (s *Scheduler) ensureComputeLocation(ctx context.Context, name string) { } } +// touchWorkerCredential sets LastSeenAt on workerID's active broker +// credential to at. Best-effort, exactly like ensureComputeLocation: a +// credential bookkeeping write must never stop a worker coming online. +// [store.ErrNotFound] — no active credential for this worker — is not +// unusual enough to warrant more than a debug log: it is the normal shape +// for a worker that enrolled with broker auth off and was seen once auth +// was later turned on, or any other legitimate mismatch between "workers +// that exist" and "workers with a credential". Any other error is logged at +// warn, matching ensureComputeLocation's posture toward its own store +// writes. +func (s *Scheduler) touchWorkerCredential(ctx context.Context, workerID string, at time.Time) { + err := s.store.TouchWorkerCredential(ctx, workerID, at) + switch { + case err == nil: + return + case errors.Is(err, store.ErrNotFound): + s.logger.DebugContext(ctx, "scheduler: no active broker credential to touch on registration", + slog.String("worker_id", workerID)) + default: + s.logger.WarnContext(ctx, "scheduler: touch worker credential last-seen failed", + slog.String("worker_id", workerID), slog.Any("error", err)) + } +} + // handleWorkerDeregister processes a worker.deregister message published by a // worker on graceful shutdown. It marks the worker offline immediately so the // scheduler stops dispatching new assignments to it rather than waiting for // the heartbeat-timeout sweep. -func (s *Scheduler) handleWorkerDeregister(ctx context.Context, msg jetstream.Msg) { +// +// subjectWorkerID is the worker the message's subject attributes it to. +func (s *Scheduler) handleWorkerDeregister(ctx context.Context, msg jetstream.Msg, subjectWorkerID string) { // DeregisterMsg mirrors protocol.DeregisterMsg; we decode only the // fields the server needs without importing the worker protocol package. var m struct { @@ -906,16 +1017,16 @@ func (s *Scheduler) handleWorkerDeregister(ctx context.Context, msg jetstream.Ms Reason string `json:"reason,omitempty"` } if err := json.Unmarshal(msg.Data(), &m); err != nil { + // The subject is the only identity left once the body will not decode. s.logger.WarnContext( ctx, "scheduler: malformed worker.deregister message", + slog.String("subject_worker_id", subjectWorkerID), slog.Any("error", err), ) s.ackMsg(ctx, msg) return } - if m.WorkerID == "" { - s.logger.WarnContext(ctx, "scheduler: worker.deregister missing worker_id") - s.ackMsg(ctx, msg) + if s.discardOnIdentityMismatch(ctx, msg, subjectWorkerID, m.WorkerID) { return } @@ -975,11 +1086,15 @@ func (s *Scheduler) handleWorkerDeregister(ctx context.Context, msg jetstream.Ms // decoding into a narrower local struct is how the version gate below stops // meaning anything — every field outside the local set drops regardless of what // version says. -func (s *Scheduler) handleWorkerHeartbeat(ctx context.Context, msg jetstream.Msg) { +// +// subjectWorkerID is the worker the message's subject attributes it to. +func (s *Scheduler) handleWorkerHeartbeat(ctx context.Context, msg jetstream.Msg, subjectWorkerID string) { var m protocol.HeartbeatMsg if err := json.Unmarshal(msg.Data(), &m); err != nil { + // The subject is the only identity left once the body will not decode. s.logger.WarnContext( ctx, "scheduler: malformed worker.heartbeat message", + slog.String("subject_worker_id", subjectWorkerID), slog.Any("error", err), ) s.ackMsg(ctx, msg) @@ -991,8 +1106,7 @@ func (s *Scheduler) handleWorkerHeartbeat(ctx context.Context, msg jetstream.Msg "this worker's liveness signal is not recorded; the heartbeat sweep will retire it") { return } - if m.WorkerID == "" { - s.ackMsg(ctx, msg) + if s.discardOnIdentityMismatch(ctx, msg, subjectWorkerID, m.WorkerID) { return } diff --git a/internal/scheduler/taskstatus.go b/internal/scheduler/taskstatus.go index b0df55ba..099e4ad0 100644 --- a/internal/scheduler/taskstatus.go +++ b/internal/scheduler/taskstatus.go @@ -7,7 +7,7 @@ package scheduler // // This file implements the task-status consumer — the counterpart of the // worker-register and worker-heartbeat consumers already in scheduler.go. -// When a worker publishes a protocol.TaskStatusMsg to task.status., +// When a worker publishes a protocol.TaskStatusMsg to task.status.., // handleTaskStatusMessage updates the store, closes the attempt record, // releases usage pool slots, and drives step/job completion logic. // @@ -40,6 +40,7 @@ import ( "github.com/nats-io/nats.go/jetstream" + "github.com/uberware/sqi/internal/bus" "github.com/uberware/sqi/internal/openjd" "github.com/uberware/sqi/internal/store" "github.com/uberware/sqi/internal/worker/protocol" @@ -55,7 +56,18 @@ func (s *Scheduler) startTaskStatusConsumer(ctx context.Context) error { } // handleTaskStatusMessage is the JetStream message handler for -// task.status. messages published by workers. +// task.status.. messages published by workers. +// +// The subject's worker ID is authoritative: it is compared against the +// attempt's recorded owner in [Scheduler.processTaskStatus], and a mismatch +// discards the message rather than applying it — otherwise any worker could +// report completion or failure for a task another worker holds. +// +// NOTE ON AUTH-OFF. With broker authentication disabled — the default — the +// subject's worker ID is present but NOT enforced by NATS, so a hostile +// client can publish under any ID. These checks then catch honest bugs, not +// attackers. That is why sqi-server warns at startup when the broker is +// unauthenticated and non-loopback. func (s *Scheduler) handleTaskStatusMessage(msg jetstream.Msg) { ctx := s.ctx @@ -90,7 +102,16 @@ func (s *Scheduler) handleTaskStatusMessage(msg jetstream.Msg) { return } - if err := s.processTaskStatus(ctx, m); err != nil { + // The subject is the only identity NATS itself can vouch for; a message + // on a subject that does not carry one concrete worker ID cannot be + // attributed to anyone and is discarded rather than acted on. + subjectWorkerID, _, ok := bus.ParseWorkerSubject(msg.Subject()) + if !ok { + s.discardUnexpectedSubject(ctx, msg, "task.status") + return + } + + if err := s.processTaskStatus(ctx, subjectWorkerID, m); err != nil { // An illegal transition is permanent: the task has moved on (retried, // canceled, already terminal) and this message describes a past state. // Redelivering cannot make it legal, so discard it rather than Nak into @@ -121,7 +142,13 @@ func (s *Scheduler) handleTaskStatusMessage(msg jetstream.Msg) { } // processTaskStatus applies a single [protocol.TaskStatusMsg] to the store. -func (s *Scheduler) processTaskStatus(ctx context.Context, m protocol.TaskStatusMsg) error { +// +// subjectWorkerID is the worker the message's subject attributes it to — the +// only identity NATS itself can vouch for. It is compared against the +// attempt's recorded WorkerID below; m.WorkerID is not trusted for this +// decision, since it is asserted by whoever sent the message rather than +// enforced by the transport. +func (s *Scheduler) processTaskStatus(ctx context.Context, subjectWorkerID string, m protocol.TaskStatusMsg) error { // Verify the attempt still exists and is for the right task. attempt, err := s.store.GetTaskAttempt(ctx, m.AttemptID) if errors.Is(err, store.ErrNotFound) { @@ -147,6 +174,22 @@ func (s *Scheduler) processTaskStatus(ctx context.Context, m protocol.TaskStatus return nil } + // The subject's worker ID was enforced by NATS when broker auth is on; + // the payload's was asserted by whoever sent the message. Trust the + // subject, and treat a mismatch as permanent — redelivery cannot make a + // forged or stale message legal, so ack it away rather than Nak into a + // loop, the same reasoning as ErrInvalidTransition above. + if subjectWorkerID != attempt.WorkerID { + s.logger.WarnContext( + ctx, "scheduler: task.status from a worker that does not hold this task — discarding", + slog.String("task_id", m.TaskID), + slog.String("attempt_id", m.AttemptID), + slog.String("subject_worker_id", subjectWorkerID), + slog.String("attempt_worker_id", attempt.WorkerID), + ) + return nil + } + at := m.At if at.IsZero() { at = time.Now().UTC() @@ -283,6 +326,13 @@ func (s *Scheduler) handleTaskTerminal( return err } } + // The attempt is now terminal (closed above, or already closed by + // RecordTaskFailure for the failed path), so no further log chunks will + // be produced for it and its cached ownership entry is no longer needed. + // SQI_LOGS and SQI_STATUS are separate streams, so a chunk published just + // before this status can still be consumed after it — that is harmless, + // since a cache miss falls back to the store and re-reads correctly. + s.attemptCache.evict(attempt.ID) // ── Transition the task ─────────────────────────────────────────────── if err := s.store.UpdateTaskStatus(ctx, m.TaskID, taskStatus); err != nil { diff --git a/internal/scheduler/taskstatus_test.go b/internal/scheduler/taskstatus_test.go index 4bb20b95..ccdf2631 100644 --- a/internal/scheduler/taskstatus_test.go +++ b/internal/scheduler/taskstatus_test.go @@ -17,12 +17,24 @@ import ( "github.com/google/uuid" + "github.com/uberware/sqi/internal/bus" "github.com/uberware/sqi/internal/store" "github.com/uberware/sqi/internal/store/fake" "github.com/uberware/sqi/internal/worker/protocol" "github.com/uberware/sqi/internal/ws" ) +// statusTestWorkerID is the worker every fixture attempt in this file opens +// on; statusTestSubject is the matching task.status subject these tests +// publish on. handleTaskStatusMessage now requires the subject's worker ID +// to match the attempt's recorded owner, so the two must agree. The job leaf +// of the subject is not itself checked (only the worker token is), so a +// fixed placeholder is fine across every test regardless of which job it +// actually seeds. +const statusTestWorkerID = "worker-1" + +var statusTestSubject = bus.TaskStatusSubject(statusTestWorkerID, "job") + // ── helpers ─────────────────────────────────────────────────────────────────── func newStatusTestScheduler(st store.Store) *Scheduler { @@ -131,6 +143,7 @@ func seedStatusFixtureWithJobStatus( attempt, err = st.CreateTaskAttempt(ctx, store.TaskAttempt{ ID: uuid.NewString(), TaskID: task.ID, + WorkerID: statusTestWorkerID, AttemptNumber: 1, Status: store.AttemptStatusRunning, StartedAt: now, @@ -164,6 +177,7 @@ func TestHandleTaskStatusMessage_MissingTaskID(t *testing.T) { s.ctx = t.Context() msg := &fakeJSMsg{ + subject: statusTestSubject, data: taskStatusMsgJSON(t, protocol.TaskStatusMsg{ Version: protocol.ProtocolVersion, TaskID: "", @@ -184,6 +198,7 @@ func TestHandleTaskStatusMessage_UnknownAttemptID(t *testing.T) { s.ctx = t.Context() msg := &fakeJSMsg{ + subject: statusTestSubject, data: taskStatusMsgJSON(t, protocol.TaskStatusMsg{ Version: protocol.ProtocolVersion, TaskID: uuid.NewString(), @@ -209,6 +224,7 @@ func TestProcessTaskStatus_Running(t *testing.T) { sessionID := "openjd-session-abc" msg := &fakeJSMsg{ + subject: statusTestSubject, data: taskStatusMsgJSON(t, protocol.TaskStatusMsg{ Version: protocol.ProtocolVersion, TaskID: task.ID, @@ -252,6 +268,7 @@ func TestProcessTaskStatus_Running_PromotesPendingJob(t *testing.T) { ) msg := &fakeJSMsg{ + subject: statusTestSubject, data: taskStatusMsgJSON(t, protocol.TaskStatusMsg{ Version: protocol.ProtocolVersion, TaskID: task.ID, @@ -289,6 +306,7 @@ func TestProcessTaskStatus_Running_DoesNotUnpauseJob(t *testing.T) { ) msg := &fakeJSMsg{ + subject: statusTestSubject, data: taskStatusMsgJSON(t, protocol.TaskStatusMsg{ Version: protocol.ProtocolVersion, TaskID: task.ID, @@ -318,6 +336,7 @@ func TestProcessTaskStatus_Succeeded(t *testing.T) { exitCode := 0 msg := &fakeJSMsg{ + subject: statusTestSubject, data: taskStatusMsgJSON(t, protocol.TaskStatusMsg{ Version: protocol.ProtocolVersion, TaskID: task.ID, @@ -341,6 +360,37 @@ func TestProcessTaskStatus_Succeeded(t *testing.T) { } } +// TestProcessTaskStatus_Succeeded_EvictsAttemptCache proves handleTaskTerminal +// evicts the attempt-owner cache entry on a terminal status, not just the +// store row. Every other test in this file only asserts on the store, so a +// deleted evict call in handleTaskTerminal would leave them all green. +func TestProcessTaskStatus_Succeeded_EvictsAttemptCache(t *testing.T) { + st := fake.New() + s := newStatusTestScheduler(st) + s.ctx = t.Context() + + _, _, task, attempt := seedStatusFixture(t, st, store.TaskStatusRunning) + s.attemptCache.put(attempt.ID, attempt.WorkerID, task.ID) + exitCode := 0 + + msg := &fakeJSMsg{ + subject: statusTestSubject, + data: taskStatusMsgJSON(t, protocol.TaskStatusMsg{ + Version: protocol.ProtocolVersion, + TaskID: task.ID, + AttemptID: attempt.ID, + Status: "succeeded", + ExitCode: &exitCode, + At: time.Now().UTC(), + }), + } + s.handleTaskStatusMessage(msg) + + if _, ok := s.attemptCache.get(attempt.ID); ok { + t.Error("expected attempt-owner cache entry to be evicted on terminal status") + } +} + func TestProcessTaskStatus_Failed(t *testing.T) { st := fake.New() s := newStatusTestScheduler(st) @@ -350,6 +400,7 @@ func TestProcessTaskStatus_Failed(t *testing.T) { exitCode := 1 msg := &fakeJSMsg{ + subject: statusTestSubject, data: taskStatusMsgJSON(t, protocol.TaskStatusMsg{ Version: protocol.ProtocolVersion, TaskID: task.ID, @@ -378,6 +429,7 @@ func TestProcessTaskStatus_Canceled(t *testing.T) { _, _, task, attempt := seedStatusFixture(t, st, store.TaskStatusRunning) msg := &fakeJSMsg{ + subject: statusTestSubject, data: taskStatusMsgJSON(t, protocol.TaskStatusMsg{ Version: protocol.ProtocolVersion, TaskID: task.ID, @@ -409,6 +461,7 @@ func TestProcessTaskStatus_Canceled_PersistsMessageAndReason(t *testing.T) { _, _, task, attempt := seedStatusFixture(t, st, store.TaskStatusRunning) msg := &fakeJSMsg{ + subject: statusTestSubject, data: taskStatusMsgJSON(t, protocol.TaskStatusMsg{ Version: protocol.ProtocolVersion, TaskID: task.ID, @@ -464,6 +517,7 @@ func TestProcessTaskStatus_Canceled_EmptyWorkerEchoPreservesServerReason(t *test // The killed worker's terminal echo always carries an empty Message. msg := &fakeJSMsg{ + subject: statusTestSubject, data: taskStatusMsgJSON(t, protocol.TaskStatusMsg{ Version: protocol.ProtocolVersion, TaskID: task.ID, @@ -505,6 +559,7 @@ func TestProcessTaskStatus_AllTasksSucceeded_StepAndJobComplete(t *testing.T) { exitCode := 0 msg := &fakeJSMsg{ + subject: statusTestSubject, data: taskStatusMsgJSON(t, protocol.TaskStatusMsg{ Version: protocol.ProtocolVersion, TaskID: task.ID, @@ -547,6 +602,7 @@ func TestProcessTaskStatus_TaskFailed_JobFails(t *testing.T) { exitCode := 1 msg := &fakeJSMsg{ + subject: statusTestSubject, data: taskStatusMsgJSON(t, protocol.TaskStatusMsg{ Version: protocol.ProtocolVersion, TaskID: task.ID, @@ -629,6 +685,7 @@ func TestProcessTaskStatus_SucceededStep_UnblocksDependentStep(t *testing.T) { attempt1, err := st.CreateTaskAttempt(ctx, store.TaskAttempt{ ID: uuid.NewString(), TaskID: task1.ID, + WorkerID: statusTestWorkerID, AttemptNumber: 1, Status: store.AttemptStatusRunning, StartedAt: now, @@ -643,6 +700,7 @@ func TestProcessTaskStatus_SucceededStep_UnblocksDependentStep(t *testing.T) { exitCode := 0 msg := &fakeJSMsg{ + subject: statusTestSubject, data: taskStatusMsgJSON(t, protocol.TaskStatusMsg{ Version: protocol.ProtocolVersion, TaskID: task1.ID, @@ -707,6 +765,7 @@ func TestProcessTaskStatus_FailedStep_CascadeCancelsDependentAndCompletesJob(t * exitCode := 1 msg := &fakeJSMsg{ + subject: statusTestSubject, data: taskStatusMsgJSON(t, protocol.TaskStatusMsg{ Version: protocol.ProtocolVersion, TaskID: task1.ID, @@ -768,6 +827,7 @@ func TestProcessTaskStatus_CascadeCancel_NotifiesCanceledTasks(t *testing.T) { exitCode := 1 msg := &fakeJSMsg{ + subject: statusTestSubject, data: taskStatusMsgJSON(t, protocol.TaskStatusMsg{ Version: protocol.ProtocolVersion, TaskID: task1.ID, @@ -811,6 +871,7 @@ func TestProcessTaskStatus_CascadeCancel_StoreError_Nacked(t *testing.T) { exitCode := 1 msg := &fakeJSMsg{ + subject: statusTestSubject, data: taskStatusMsgJSON(t, protocol.TaskStatusMsg{ Version: protocol.ProtocolVersion, TaskID: task1.ID, @@ -848,6 +909,7 @@ func TestProcessTaskStatus_UpdateAttemptError_Nacked(t *testing.T) { exitCode := 0 msg := &fakeJSMsg{ + subject: statusTestSubject, data: taskStatusMsgJSON(t, protocol.TaskStatusMsg{ Version: protocol.ProtocolVersion, TaskID: task.ID, diff --git a/internal/scheduler/taskstatus_transition_test.go b/internal/scheduler/taskstatus_transition_test.go index c807a136..e892c1bf 100644 --- a/internal/scheduler/taskstatus_transition_test.go +++ b/internal/scheduler/taskstatus_transition_test.go @@ -36,6 +36,7 @@ func TestHandleTaskStatusMessage_InvalidTransitionIsAcked(t *testing.T) { } msg := &fakeJSMsg{ + subject: statusTestSubject, data: taskStatusMsgJSON(t, protocol.TaskStatusMsg{ Version: protocol.ProtocolVersion, TaskID: task.ID, @@ -74,6 +75,7 @@ func TestHandleTaskStatusMessage_DuplicateRunningIsAcked(t *testing.T) { for i := range 2 { msg := &fakeJSMsg{ + subject: statusTestSubject, data: taskStatusMsgJSON(t, protocol.TaskStatusMsg{ Version: protocol.ProtocolVersion, TaskID: task.ID, diff --git a/internal/server/exprlimits.go b/internal/server/exprlimits.go index 0ebc53b4..237fd67c 100644 --- a/internal/server/exprlimits.go +++ b/internal/server/exprlimits.go @@ -47,5 +47,6 @@ func ExprLimitsFromConfig(c config.OpenJDConfig) openjd.ExprLimits { func schedulerConfig(cfg Config) scheduler.Config { sched := cfg.Scheduler sched.ExprLimits = cfg.OpenJDExprLimits + sched.NATSAuthEnabled = cfg.NATSAuthEnabled return sched } diff --git a/internal/server/natsauthdeps_test.go b/internal/server/natsauthdeps_test.go new file mode 100644 index 00000000..280e3c54 --- /dev/null +++ b/internal/server/natsauthdeps_test.go @@ -0,0 +1,124 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package server + +import ( + "net/http" + "testing" + "time" + + "github.com/go-chi/chi/v5" + + "github.com/uberware/sqi/internal/api" + "github.com/uberware/sqi/internal/auth" + "github.com/uberware/sqi/internal/health" + "github.com/uberware/sqi/internal/metrics" + "github.com/uberware/sqi/internal/store/fake" +) + +// TestNATSAuthDeps_CopiesAllFourSettings covers the SECOND of the three hops +// that carry the nats.auth.* settings from the config file to the REST +// worker-enrollment surface: server.Config -> api.Deps. See +// TestServerConfig_CarriesTheBrokerAuthSettings in cmd/sqi-server for the +// first hop. +func TestNATSAuthDeps_CopiesAllFourSettings(t *testing.T) { + cfg := Config{ + NATSAuthEnabled: true, + NATSAuthEnrollmentEndpointEnabled: true, + NATSAuthJoinTokenTTL: 42 * time.Minute, + NATSAuthJoinTokenSingleUse: false, + } + var deps api.Deps + natsAuthDeps(cfg, &deps) + + if !deps.NATSAuthEnabled { + t.Error("deps.NATSAuthEnabled = false, want true") + } + if !deps.NATSAuthEnrollmentEndpointEnabled { + t.Error("deps.NATSAuthEnrollmentEndpointEnabled = false, want true") + } + if want := 42 * time.Minute; deps.JoinTokenTTL != want { + t.Errorf("deps.JoinTokenTTL = %s, want %s", deps.JoinTokenTTL, want) + } + if deps.JoinTokenSingleUse { + t.Error("deps.JoinTokenSingleUse = true, want false") + } +} + +// TestNATSAuthDeps_DoesNotDisturbFieldsItDoesNotOwn confirms natsAuthDeps +// only ever writes its own four fields, mirroring wireAuthDeps's contract: +// deps is built incrementally across several steps in start, so a mapping +// function that clobbers a sibling's field fails silently — the field it +// overwrote simply reverts to its zero value and the server boots anyway. +func TestNATSAuthDeps_DoesNotDisturbFieldsItDoesNotOwn(t *testing.T) { + deps := api.Deps{CookieName: "sqi_session", SessionTTL: time.Hour} + natsAuthDeps(Config{}, &deps) + + if deps.CookieName != "sqi_session" { + t.Errorf("CookieName clobbered: %q", deps.CookieName) + } + if deps.SessionTTL != time.Hour { + t.Errorf("SessionTTL clobbered: %s", deps.SessionTTL) + } +} + +// routeMounted reports whether method+pattern is registered on r. +func routeMounted(t *testing.T, r chi.Router, method, pattern string) bool { + t.Helper() + found := false + err := chi.Walk(r, func(m, route string, _ http.Handler, _ ...func(http.Handler) http.Handler) error { + if m == method && route == pattern { + found = true + } + return nil + }) + if err != nil { + t.Fatalf("chi.Walk: %v", err) + } + return found +} + +// TestNATSAuthDeps_EndToEnd_EnrollRouteMountedWhenConfigured is the THIRD +// hop: build the exact api.Deps a real boot would produce from a Config with +// broker auth and the enrollment endpoint both on, hand it to the real +// api.NewRouter, and confirm POST /api/v1/workers/enroll is actually live. +// +// Nothing else fails when this hop breaks: the server still boots, `config +// print` still echoes the operator's values, and the route simply never +// mounts. +func TestNATSAuthDeps_EndToEnd_EnrollRouteMountedWhenConfigured(t *testing.T) { + cfg := Config{ + NATSAuthEnabled: true, + NATSAuthEnrollmentEndpointEnabled: true, + } + deps := api.Deps{Store: fake.New(), Auth: auth.Anonymous()} + natsAuthDeps(cfg, &deps) + + r := api.NewRouter(api.Config{DisableRateLimit: true}, deps, testLogger(), metrics.New(), health.NewRegistry()) + if !routeMounted(t, r, http.MethodPost, "/api/v1/workers/enroll") { + t.Error("POST /api/v1/workers/enroll is not mounted with NATSAuthEnabled and " + + "NATSAuthEnrollmentEndpointEnabled both true — the config-to-Deps wiring is broken") + } +} + +// TestNATSAuthDeps_EndToEnd_EnrollRouteAbsentAtDefaults is the companion +// negative case: a server built from DefaultConfig() (broker auth off, the +// v0.3.0 and pre-H1 behavior) must never expose the enrollment endpoint. +func TestNATSAuthDeps_EndToEnd_EnrollRouteAbsentAtDefaults(t *testing.T) { + deps := api.Deps{Store: fake.New(), Auth: auth.Anonymous()} + natsAuthDeps(DefaultConfig(), &deps) + + r := api.NewRouter(api.Config{DisableRateLimit: true}, deps, testLogger(), metrics.New(), health.NewRegistry()) + if routeMounted(t, r, http.MethodPost, "/api/v1/workers/enroll") { + t.Error("POST /api/v1/workers/enroll is mounted at default configuration " + + "(broker auth off); it must not exist until nats.auth.enabled is turned on") + } + + // Confirm it collides with /workers/{id} the way router.go's own comment + // documents (workerenroll_test.go, internal/api, worked around this by + // walking the route table instead of asserting a status code): a GET or + // DELETE for that pattern IS registered, just never POST. + if !routeMounted(t, r, http.MethodGet, "/api/v1/workers/{id}") { + t.Fatal("test assumption broken: GET /api/v1/workers/{id} is not registered") + } +} diff --git a/internal/server/revokeworker_test.go b/internal/server/revokeworker_test.go new file mode 100644 index 00000000..61081a61 --- /dev/null +++ b/internal/server/revokeworker_test.go @@ -0,0 +1,366 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package server + +import ( + "context" + "errors" + "net" + "sync" + "testing" + "time" + + nats "github.com/nats-io/nats.go" + "github.com/nats-io/nkeys" + + "github.com/uberware/sqi/internal/brokerauth" + "github.com/uberware/sqi/internal/bus" + "github.com/uberware/sqi/internal/store" + "github.com/uberware/sqi/internal/store/fake" +) + +// Unit tests for [Server.RevokeWorker] — the method that turns DELETE +// /api/v1/workers/{id}/credential into a synchronous disconnect, as opposed +// to "sqi-server worker revoke", which only ever writes the store from a +// separate process holding no broker handle. +// +// These exercise RevokeWorker directly against a *Server built by hand +// (store + a real embedded broker, no HTTP layer, no scheduler) so the +// store-write-then-reload ordering and its failure mode are pinned at the +// unit level. The full HTTP-to-disconnect-to-reclaim path is covered by +// test/integration's broker-auth suite, which also proves the existing +// heartbeat-sweep/reclaim path — not anything reimplemented here — is what +// returns a revoked worker's task to ready. + +// freeLoopbackAddr asks the OS for an available loopback TCP port and +// returns "127.0.0.1:", releasing it immediately so the broker under +// test can bind it. +func freeLoopbackAddr(t *testing.T) string { + t.Helper() + lc := &net.ListenConfig{} + ln, err := lc.Listen(context.Background(), "tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("freeLoopbackAddr: listen: %v", err) + } + addr := ln.Addr().String() + if err := ln.Close(); err != nil { + t.Fatalf("freeLoopbackAddr: close: %v", err) + } + return addr +} + +// startTestBroker boots a real embedded broker with authentication enabled +// and enrolled with creds, and registers cleanup. +func startTestBroker(t *testing.T, creds []bus.WorkerCredentialRef) *bus.Broker { + t.Helper() + b := bus.New(bus.BrokerConfig{ + Addr: freeLoopbackAddr(t), + DataDir: t.TempDir() + "/nats", + MaxStoreMB: 64, + Auth: bus.BrokerAuthConfig{Enabled: true, Credentials: creds}, + }, testLogger()) + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + if err := b.Start(ctx); err != nil { + t.Fatalf("startTestBroker: Start: %v", err) + } + t.Cleanup(b.Shutdown) + return b +} + +// enrolledCredential generates a fresh nkey, seeds a matching +// [store.WorkerCredential] row in st, and returns the [bus.WorkerCredentialRef] +// and raw seed needed to connect as that worker. +func enrolledCredential(t *testing.T, st store.Store, workerID string) (bus.WorkerCredentialRef, []byte) { + t.Helper() + seed, pub, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("enrolledCredential: GenerateSeed: %v", err) + } + if _, err := st.CreateWorkerCredential(context.Background(), store.WorkerCredential{ + ID: workerID + "-cred", + WorkerID: workerID, + PublicKey: pub, + EnrolledAt: time.Now().UTC(), + }); err != nil { + t.Fatalf("enrolledCredential: CreateWorkerCredential: %v", err) + } + return bus.WorkerCredentialRef{WorkerID: workerID, PublicKey: pub}, seed +} + +// nkeyOption builds a nats.Option that authenticates as the nkey pair +// identified by pub, signing server challenges with seed. +func nkeyOption(t *testing.T, seed []byte, pub string) nats.Option { + t.Helper() + return nats.Nkey(pub, func(nonce []byte) ([]byte, error) { + kp, err := nkeys.FromSeed(seed) + if err != nil { + return nil, err + } + return kp.Sign(nonce) + }) +} + +// connectAsWorker dials b as the given nkey credential, with NoReconnect and +// a ClosedHandler feeding the returned channel — the same pattern +// internal/bus's own revocation tests use, so the disconnect assertion is +// not coupled to nats.go's reconnect/backoff timing. +func connectAsWorker(t *testing.T, b *bus.Broker, seed []byte, pub string) (*nats.Conn, <-chan struct{}) { + t.Helper() + closedCh := make(chan struct{}) + nc, err := nats.Connect( + b.ClientURL(), + nkeyOption(t, seed, pub), + nats.NoReconnect(), + nats.ClosedHandler(func(*nats.Conn) { close(closedCh) }), + ) + if err != nil { + t.Fatalf("connectAsWorker: Connect: %v", err) + } + t.Cleanup(func() { + if !nc.IsClosed() { + nc.Close() + } + }) + return nc, closedCh +} + +// TestRevokeWorker_NATSAuthDisabled_StoreWriteOnly covers the farm that runs +// without broker authentication at all: RevokeWorker must still perform the +// store write (so an operator can pre-provision revocations, or clean up +// after auth was turned off) but must never touch the broker — there is no +// authorized-key set to reload, and s.broker is left nil here specifically +// to prove that: touching it would panic. +func TestRevokeWorker_NATSAuthDisabled_StoreWriteOnly(t *testing.T) { + st := fake.New() + ref, _ := enrolledCredential(t, st, "worker-a") + + s := &Server{cfg: Config{NATSAuthEnabled: false}, store: st, logger: testLogger()} + + if err := s.RevokeWorker(context.Background(), ref.WorkerID); err != nil { + t.Fatalf("RevokeWorker: %v", err) + } + + if _, err := st.GetActiveWorkerCredentialByWorkerID(context.Background(), ref.WorkerID); !errors.Is(err, store.ErrNotFound) { + t.Fatalf("GetActiveWorkerCredentialByWorkerID after revoke: %v, want store.ErrNotFound", err) + } +} + +// TestRevokeWorker_UnknownWorker_ReturnsErrNotFoundWithoutTouchingBroker +// proves the store write happens FIRST: with no credential for "ghost" in +// the store, RevokeWorker must fail and return before ever reaching s.broker +// — which is nil here, so touching it would panic — even though +// NATSAuthEnabled is true. +func TestRevokeWorker_UnknownWorker_ReturnsErrNotFoundWithoutTouchingBroker(t *testing.T) { + st := fake.New() + s := &Server{cfg: Config{NATSAuthEnabled: true}, store: st, logger: testLogger()} + + err := s.RevokeWorker(context.Background(), "ghost") + if !errors.Is(err, store.ErrNotFound) { + t.Fatalf("RevokeWorker(\"ghost\") = %v, want store.ErrNotFound", err) + } +} + +// TestRevokeWorker_ReloadDisconnectsRevokedWorkerOnly is the core positive +// case: revoking worker A through Server.RevokeWorker disconnects A's live +// broker connection inside the call (nats-server's reloadAuthorization +// re-authorizes every connected client synchronously) and leaves B +// completely unaffected. +func TestRevokeWorker_ReloadDisconnectsRevokedWorkerOnly(t *testing.T) { + st := fake.New() + refA, seedA := enrolledCredential(t, st, "worker-a") + refB, seedB := enrolledCredential(t, st, "worker-b") + + broker := startTestBroker(t, []bus.WorkerCredentialRef{refA, refB}) + s := &Server{cfg: Config{NATSAuthEnabled: true}, store: st, broker: broker, logger: testLogger()} + + _, closedA := connectAsWorker(t, broker, seedA, refA.PublicKey) + ncB, closedB := connectAsWorker(t, broker, seedB, refB.PublicKey) + + if err := s.RevokeWorker(context.Background(), refA.WorkerID); err != nil { + t.Fatalf("RevokeWorker: %v", err) + } + + select { + case <-closedA: + case <-time.After(2 * time.Second): + t.Fatal("worker A's connection was not closed after revocation") + } + + select { + case <-closedB: + t.Fatal("worker B's connection was closed by an unrelated revocation") + case <-time.After(200 * time.Millisecond): + } + if err := ncB.Flush(); err != nil { + t.Fatalf("worker B's connection unusable after A's revocation: %v", err) + } + + if _, err := st.GetActiveWorkerCredentialByWorkerID(context.Background(), refA.WorkerID); !errors.Is(err, store.ErrNotFound) { + t.Errorf("A's credential still active after revoke: %v, want store.ErrNotFound", err) + } + if _, err := st.GetActiveWorkerCredentialByWorkerID(context.Background(), refB.WorkerID); err != nil { + t.Errorf("B's credential was disturbed by A's revocation: %v", err) + } +} + +// TestRevokeWorker_ReloadFailure_StoreStaysRevoked pins the defensible +// outcome of a reload failure after a successful store write: the credential +// row stays revoked (the store is never rolled back) and the failure is +// still surfaced to the caller, since the synchronous disconnect this method +// exists to provide did not actually happen. A shut-down broker stands in +// for "the reload call itself failed" — ReloadCredentials returns "broker +// not started" once Shutdown has run, which is the same shape of failure as +// any other ReloadOptions error from the caller's point of view. +func TestRevokeWorker_ReloadFailure_StoreStaysRevoked(t *testing.T) { + st := fake.New() + ref, _ := enrolledCredential(t, st, "worker-a") + + broker := startTestBroker(t, []bus.WorkerCredentialRef{ref}) + broker.Shutdown() // torn down before RevokeWorker runs + + s := &Server{cfg: Config{NATSAuthEnabled: true}, store: st, broker: broker, logger: testLogger()} + + err := s.RevokeWorker(context.Background(), ref.WorkerID) + if err == nil { + t.Fatal("RevokeWorker: want an error when the broker reload fails, got nil") + } + + // The store write must stand regardless: rolling it back on a reload + // failure would leave a credential the operator explicitly revoked + // silently trusted again, which is worse than the reload simply not + // having taken effect yet. + if _, err := st.GetActiveWorkerCredentialByWorkerID(context.Background(), ref.WorkerID); !errors.Is(err, store.ErrNotFound) { + t.Errorf("credential still active after a reload failure: %v, want store.ErrNotFound", err) + } +} + +// ── An unserialized read-then-reload span loses an update ────────────────── +// +// An unlocked "read the active credential set, then call ReloadCredentials" +// sequence is vulnerable to a lost update. Two concurrent revocations of +// DIFFERENT workers can interleave like this: +// +// 1. revoke(A) commits its store write. +// 2. revoke(A) reads the active set — B is still active, so the set +// contains B. +// 3. revoke(B) commits its store write. +// 4. revoke(B) reads the active set — correctly excludes A and B — and +// reloads. +// 5. revoke(A)'s reload, built from the STALE step-2 read, applies LAST +// and reintroduces B into the broker's trusted key set. +// +// raceMarkerKey and blockingListStore below reproduce this deterministically +// instead of hoping many iterations happen to hit the right interleaving: +// A's read is captured, then paused (holding s.brokerReloadMu) until the +// test explicitly releases it — by which point B's own revoke has fully +// run. This pins the exact scenario s.brokerReloadMu closes, not just "some +// race somewhere." + +// raceMarkerKey tags a context so blockingListStore knows which caller's +// ListActiveWorkerCredentials call to pause. +type raceMarkerKey struct{} + +// blockingListStore wraps a store.Store and, only for the call whose +// context carries blockFor, pauses AFTER reading the real result (so the +// caller holds a real, but soon-to-be-stale, snapshot) until release is +// closed. entered is closed the moment the pause begins, so the test knows +// the blocked call has already captured its snapshot before letting the +// other revoke proceed. +type blockingListStore struct { + store.Store + + blockFor string + entered chan struct{} + release chan struct{} +} + +func (s *blockingListStore) ListActiveWorkerCredentials(ctx context.Context) ([]store.WorkerCredential, error) { + creds, err := s.Store.ListActiveWorkerCredentials(ctx) + if v, ok := ctx.Value(raceMarkerKey{}).(string); ok && v == s.blockFor { + close(s.entered) + <-s.release + } + return creds, err +} + +// TestRevokeWorker_ConcurrentRevocationsOfDifferentWorkers_BothStayRevoked +// deterministically forces the interleaving described above and asserts the +// broker ends up trusting NEITHER worker — not just that the store shows +// both revoked (the store was never the vulnerable part; the broker's +// authorized-key set was). +func TestRevokeWorker_ConcurrentRevocationsOfDifferentWorkers_BothStayRevoked(t *testing.T) { + st := fake.New() + refA, seedA := enrolledCredential(t, st, "worker-a") + refB, seedB := enrolledCredential(t, st, "worker-b") + + broker := startTestBroker(t, []bus.WorkerCredentialRef{refA, refB}) + + wrapped := &blockingListStore{ + Store: st, + blockFor: "A", + entered: make(chan struct{}), + release: make(chan struct{}), + } + s := &Server{cfg: Config{NATSAuthEnabled: true}, store: wrapped, broker: broker, logger: testLogger()} + + ctxA := context.WithValue(context.Background(), raceMarkerKey{}, "A") + ctxB := context.WithValue(context.Background(), raceMarkerKey{}, "B") // never matches blockFor; B is not paused + + var wg sync.WaitGroup + var errA, errB error + wg.Go(func() { + errA = s.RevokeWorker(ctxA, refA.WorkerID) + }) + + // Wait until A has committed its store write, read the (still-stale, + // B-included) active set, and is now paused holding that snapshot. + <-wrapped.entered + + // Run B's revoke to completion while A is paused. RevokeWorker's own + // call blocks trying to acquire s.brokerReloadMu (A is still inside the + // critical section) until A finishes — so this goroutine only returns + // after A's whole RevokeWorker call has completed. An unserialized + // implementation would instead let B run immediately and finish well + // before A resumes. + wg.Go(func() { + errB = s.RevokeWorker(ctxB, refB.WorkerID) + }) + + // Give B's goroutine a moment to either finish (unserialized) or block + // on the mutex (serialized) before releasing A — the property under + // test does not depend on this sleep's exact duration, only that B has + // had the chance to run to the point it would reach if it were going to. + time.Sleep(50 * time.Millisecond) + close(wrapped.release) + + wg.Wait() + + if errA != nil { + t.Errorf("RevokeWorker(A): %v", errA) + } + if errB != nil { + t.Errorf("RevokeWorker(B): %v", errB) + } + + // Both must be gone from the store... + if _, err := st.GetActiveWorkerCredentialByWorkerID(context.Background(), refA.WorkerID); !errors.Is(err, store.ErrNotFound) { + t.Errorf("worker A still active in the store: %v, want store.ErrNotFound", err) + } + if _, err := st.GetActiveWorkerCredentialByWorkerID(context.Background(), refB.WorkerID); !errors.Is(err, store.ErrNotFound) { + t.Errorf("worker B still active in the store: %v, want store.ErrNotFound", err) + } + + // ...and, the property this test exists to pin, neither can connect to + // the broker: a fresh connection attempt with either's key must be + // refused. An unserialized read-then-reload would let B's stale reload + // reintroduce it here even though the store already shows it revoked. + if nc, err := nats.Connect(broker.ClientURL(), nkeyOption(t, seedA, refA.PublicKey), nats.NoReconnect()); err == nil { + nc.Close() + t.Error("worker A connected to the broker after concurrent revocation — its credential was reintroduced") + } + if nc, err := nats.Connect(broker.ClientURL(), nkeyOption(t, seedB, refB.PublicKey), nats.NoReconnect()); err == nil { + nc.Close() + t.Error("worker B connected to the broker after concurrent revocation — its credential was reintroduced by a stale reload") + } +} diff --git a/internal/server/routerconfig.go b/internal/server/routerconfig.go index 6fa6dcce..2592eb12 100644 --- a/internal/server/routerconfig.go +++ b/internal/server/routerconfig.go @@ -35,3 +35,26 @@ func routerConfig(cfg Config, workerOfflineThreshold time.Duration) api.Config { ExprLimits: cfg.OpenJDExprLimits, } } + +// natsAuthDeps copies this server's broker-auth-derived settings onto deps: +// whether POST /api/v1/workers/enroll is mounted at all, and the join-token +// defaults its handlers consult. It mutates deps rather than returning an +// api.Deps, unlike routerConfig, because Deps is built incrementally across +// several steps in start (wireAuthDeps, the preset library, diagnostics) and +// a returned value here would either overwrite those or need its own merge. +// +// Split out for the same reason routerConfig is (see its doc comment): a +// struct-field assignment inline in start can only be reached by booting a +// whole server, so nothing would catch a dropped or transposed line. +// +// Nothing else fails when this hop breaks: the server still boots, `config +// print` still echoes the operator's values, and the route simply never +// mounts — a deployed server would run with all four at their zero value +// regardless of nats.auth.* configuration, the enrollment endpoint never +// mounted and a minted token carrying a zero TTL. +func natsAuthDeps(cfg Config, deps *api.Deps) { + deps.NATSAuthEnabled = cfg.NATSAuthEnabled + deps.NATSAuthEnrollmentEndpointEnabled = cfg.NATSAuthEnrollmentEndpointEnabled + deps.JoinTokenTTL = cfg.NATSAuthJoinTokenTTL + deps.JoinTokenSingleUse = cfg.NATSAuthJoinTokenSingleUse +} diff --git a/internal/server/server.go b/internal/server/server.go index ae31e50e..712ca3e2 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -20,6 +20,7 @@ import ( "net" "net/http" "slices" + "sync" "time" "github.com/uberware/sqi/internal/api" @@ -72,11 +73,37 @@ type Config struct { // NATSAddr is the TCP address the embedded NATS server listens on. // It defaults to all interfaces so that workers discovering the server - // via mDNS can connect to NATS at the advertised LAN host. (Broker - // authentication does not exist: any host that can reach this port can - // register as a worker and receive assignments. Deferred to Phase 4.) + // via mDNS can connect to NATS at the advertised LAN host. + // Broker authentication is opt-in (nats.auth.enabled) and off by + // default; when off, any host that can reach this port can register as + // a worker. warnIfBrokerUnauthenticated logs this at startup. NATSAddr string // default "0.0.0.0:4222" + // NATSAuthEnabled reports whether the broker requires a per-worker + // credential. Used for the startup warning and to configure the broker + // itself, including which workers it authorizes. Also gates, together + // with NATSAuthEnrollmentEndpointEnabled, whether POST + // /api/v1/workers/enroll is mounted at all. + NATSAuthEnabled bool + + // NATSAuthEnrollmentEndpointEnabled mirrors + // config.NATSAuthConfig.EnrollmentEndpointEnabled: whether POST + // /api/v1/workers/enroll is mounted. Meaningful only when NATSAuthEnabled + // is true; a site that provisions every credential by hand (`sqi-server + // worker enroll`) can turn this off to remove the self-service surface + // entirely. + NATSAuthEnrollmentEndpointEnabled bool + + // NATSAuthJoinTokenTTL mirrors config.NATSAuthConfig.JoinTokenTTL: how + // long a join token minted by POST /api/v1/workers/join-tokens remains + // valid. Already bounds-checked at config load. + NATSAuthJoinTokenTTL time.Duration + + // NATSAuthJoinTokenSingleUse mirrors + // config.NATSAuthConfig.JoinTokenSingleUse: whether a join token is + // rejected on a second enrollment attempt after its first successful use. + NATSAuthJoinTokenSingleUse bool + // NATSDataDir is the directory used by JetStream for file-backed stream // storage. It is created at startup if it does not exist. NATSDataDir string // default "data/nats" @@ -240,6 +267,18 @@ type Server struct { wsHub *ws.Hub // WebSocket fan-out hub discovery *discovery.Responder // mDNS advertisement + // brokerReloadMu serializes every "read the active worker-credential set + // from the store, then reload it into the broker" span — see + // reloadBrokerCredentials. Both a revocation and an enrollment trigger + // that span, from different goroutines (different HTTP requests), and + // the span is not safe to interleave: a reload built from a read that + // started before another writer's store commit can still finish AFTER + // that writer's own reload, silently reintroducing whatever the other + // writer just removed (or omitting whatever it just added). Locking only + // the call into the broker, not the read that precedes it, would not + // close this — the stale READ is what makes the reload wrong. + brokerReloadMu sync.Mutex + // diagBuf is the in-memory diagnostic-log ring buffer. It is created in the // serve command before the logger (so the server's own logs are captured // from the first line) and threaded here. Nil when diagnostics are disabled. @@ -268,6 +307,177 @@ func (s *Server) Metrics() *metrics.Metrics { return s.metrics } +// warnIfBrokerUnauthenticated emits a WARN when the NATS broker is reachable +// from outside this machine and has no credential requirement. +// +// This is the highest-value line in the whole broker-auth component: broker +// auth is opt-in and off by default, so this is the only thing that tells an +// operator who turned on auth.enabled that their worker transport is still +// wide open. It stays silent on loopback, where the exposure does not exist, +// and silent on an unparseable address, where config validation has already +// produced a better message. +func warnIfBrokerUnauthenticated(ctx context.Context, natsAddr string, brokerAuthEnabled bool, logger *slog.Logger) { + if brokerAuthEnabled { + return + } + host, _, err := net.SplitHostPort(natsAddr) + if err != nil { + return + } + if isLoopbackHost(host) { + return + } + logger.WarnContext( + ctx, "the NATS broker is unauthenticated and reachable beyond this host", + slog.String("nats_addr", natsAddr), + slog.String("impact", "any host that can reach this port can register as a worker and execute submitted job code"), + slog.String("remediation", "set nats.auth.enabled: true and enroll your workers, or bind nats.addr to 127.0.0.1:4222 for single-machine use"), + ) +} + +// isLoopbackHost reports whether host names only this machine. An empty host +// (from ":4222") means all interfaces, which is not loopback. +func isLoopbackHost(host string) bool { + if host == "" { + return false + } + if host == "localhost" { + return true + } + ip := net.ParseIP(host) + return ip != nil && ip.IsLoopback() +} + +// startBroker constructs and starts the embedded NATS broker, warning about +// an unauthenticated broker exposed beyond this host and loading the +// enrolled worker credential set when broker authentication is enabled. +func (s *Server) startBroker(ctx context.Context) (*bus.Broker, error) { + warnIfBrokerUnauthenticated(ctx, s.cfg.NATSAddr, s.cfg.NATSAuthEnabled, s.logger) + + brokerAuth, err := loadBrokerAuthConfig(ctx, s.store, s.cfg.NATSAuthEnabled) + if err != nil { + return nil, fmt.Errorf("load worker credentials: %w", err) + } + + broker := bus.New(bus.BrokerConfig{ + Addr: s.cfg.NATSAddr, + DataDir: s.cfg.NATSDataDir, + MaxStoreMB: s.cfg.NATSMaxStoreMB, + Auth: brokerAuth, + }, s.logger) + if err := broker.Start(ctx); err != nil { + return nil, err + } + return broker, nil +} + +// loadBrokerAuthConfig builds the broker's authorization configuration. +// +// It queries the store for the enrolled worker credential set only when +// enabled is true, so that a server running with broker authentication off +// never touches the credential table — the default startup path stays +// byte-for-byte what it was before broker authentication existed. +func loadBrokerAuthConfig(ctx context.Context, st store.WorkerCredentialStore, enabled bool) (bus.BrokerAuthConfig, error) { + if !enabled { + return bus.BrokerAuthConfig{}, nil + } + creds, err := st.ListActiveWorkerCredentials(ctx) + if err != nil { + return bus.BrokerAuthConfig{}, err + } + refs := make([]bus.WorkerCredentialRef, 0, len(creds)) + for _, c := range creds { + refs = append(refs, bus.WorkerCredentialRef{ + WorkerID: c.WorkerID, + PublicKey: c.PublicKey, + }) + } + return bus.BrokerAuthConfig{Enabled: true, Credentials: refs}, nil +} + +// RevokeWorker revokes workerID's active broker credential and, when broker +// authentication is enabled, disconnects it and reclaims its in-flight work. +// It implements [api.WorkerRevoker] and is what makes DELETE +// /api/v1/workers/{id}/credential the synchronous revocation path: the +// offline "sqi-server worker revoke" CLI command writes the same store row +// from a separate process holding no broker handle, so it can only apply at +// the broker's next start; this method runs inside the server process where +// the broker handle lives, so it can act on a running broker immediately. +// +// The store is written FIRST and the broker reloaded SECOND, on purpose: a +// reload failure after a successful store write leaves the worst-case state +// "revoked in the store, still trusted by the running broker" — recoverable +// at the broker's next start or next successful reload, and never worse than +// what the offline CLI path already promises. The reverse ordering could +// instead leave a credential trusted by neither the store nor a +// findable-again authorized-key set, which nothing could repair. The reload +// failure is still returned to the caller — the synchronous guarantee this +// method exists to provide was not met — but it never triggers a rollback of +// the store write. +func (s *Server) RevokeWorker(ctx context.Context, workerID string) error { + if err := s.store.RevokeWorkerCredential(ctx, workerID, time.Now().UTC()); err != nil { + return err + } + return s.reloadBrokerCredentials(ctx) +} + +// ReloadBrokerCredentials implements [api.BrokerCredentialReloader]. It is +// called after a new worker credential is created (self-service REST +// enrollment) so that worker can connect to THIS running broker without an +// operator restarting it — loadBrokerAuthConfig otherwise only ever runs +// once, at Start, so without this call a freshly-enrolled worker's +// connection is refused by a broker whose Options.Nkeys was built before +// that credential existed, and the worker exits fatally naming that +// rejection. +// +// Unlike a revoke's reload failure — which the caller must hear about, +// because it means the broker still trusts a credential the store says is +// gone — a failure here means the opposite direction: the broker is +// (temporarily) too STRICT, not too permissive. The credential the caller +// asked to create is genuinely created and durable; the worker simply cannot +// connect until the next successful reload or restart, and can retry then. +// Telling the enrolling caller "enrollment failed" would be false, so the +// REST handler logs this failure and still reports success — see +// workerenroll.go's enroll. +func (s *Server) ReloadBrokerCredentials(ctx context.Context) error { + return s.reloadBrokerCredentials(ctx) +} + +// reloadBrokerCredentials re-reads the active worker-credential set from the +// store and reloads it into the broker's authorized-key set. It backs both +// RevokeWorker and ReloadBrokerCredentials, which is deliberate: the two +// are the same operation ("make the broker's authorized-key set match the +// store's active rows right now"), triggered by opposite events. +// +// The read-then-reload span is serialized by s.brokerReloadMu — see that +// field's doc comment for why an unserialized version loses updates. +// Locking only the call into the broker (not the read that precedes it) +// would not close that race: the read is what goes stale. +// +// Returns nil immediately, without touching the store or the broker, when +// broker authentication is disabled: there is no authorized-key set to +// reload, and this keeps the auth-off path free of the extra store read. +func (s *Server) reloadBrokerCredentials(ctx context.Context) error { + if !s.cfg.NATSAuthEnabled { + return nil + } + if s.broker == nil { + return errors.New("reload broker credentials: broker not started") + } + + s.brokerReloadMu.Lock() + defer s.brokerReloadMu.Unlock() + + brokerAuth, err := loadBrokerAuthConfig(ctx, s.store, true) + if err != nil { + return fmt.Errorf("reload broker credentials: read active credential set: %w", err) + } + if err := s.broker.ReloadCredentials(brokerAuth.Credentials); err != nil { + return fmt.Errorf("reload broker credentials: %w", err) + } + return nil +} + // Health returns the [*health.Registry] owned by this server. Components // (store, bus) call Health().Register(...) during startup to participate in // readiness gating via GET /readyz. @@ -338,12 +548,8 @@ func (s *Server) start(ctx context.Context) error { // ── Message bus (NATS JetStream) ─────────────────────────────────────── // Embed NATS server, enable JetStream, provision streams. // Typed client wrapper, consumers, reconnect, drain. - broker := bus.New(bus.BrokerConfig{ - Addr: s.cfg.NATSAddr, - DataDir: s.cfg.NATSDataDir, - MaxStoreMB: s.cfg.NATSMaxStoreMB, - }, s.logger) - if err := broker.Start(ctx); err != nil { + broker, err := s.startBroker(ctx) + if err != nil { return fmt.Errorf("start bus: %w", err) } s.broker = broker @@ -405,10 +611,12 @@ func (s *Server) start(ctx context.Context) error { EnforceLimits: s.cfg.EnforceOpenJDLimits, ExprLimits: s.cfg.OpenJDExprLimits, }), - Products: product.NewCatalog(s.store), - Scheduler: s.sched, - Hub: s.wsHub, - Version: version.Get(), + Products: product.NewCatalog(s.store), + Scheduler: s.sched, + Hub: s.wsHub, + Version: version.Get(), + WorkerRevoker: s, + BrokerCredentialReloader: s, } // Only expose the diagnostics reader when diagnostics are enabled. Leaving // DiagReader as a nil interface (rather than a typed-nil *diag.Buffer) makes @@ -428,6 +636,7 @@ func (s *Server) start(ctx context.Context) error { deps.SessionTTL = s.cfg.AuthSessionTTL deps.CookieName = s.cfg.AuthCookieName deps.CookieSecure = s.cfg.AuthCookieSecure + natsAuthDeps(s.cfg, &deps) router := api.NewRouter( routerConfig(s.cfg, s.sched.WorkerTimeout()), deps, diff --git a/internal/server/server_test.go b/internal/server/server_test.go index 52012674..0c308112 100644 --- a/internal/server/server_test.go +++ b/internal/server/server_test.go @@ -2,7 +2,13 @@ package server -import "testing" +import ( + "bytes" + "context" + "log/slog" + "strings" + "testing" +) func TestBrowseURL(t *testing.T) { tests := []struct { @@ -25,3 +31,37 @@ func TestBrowseURL(t *testing.T) { }) } } + +func TestWarnIfBrokerUnauthenticated(t *testing.T) { + tests := []struct { + name string + addr string + authEnabled bool + wantWarn bool + }{ + {"non-loopback, auth off", "0.0.0.0:4222", false, true}, + // An empty host means "all interfaces", exactly as 0.0.0.0 does — + // isLoopbackHost("") returning false is what makes both warn. + {"empty host, auth off", ":4222", false, true}, + {"specific LAN ip, auth off", "192.168.1.10:4222", false, true}, + {"ipv6 any, auth off", "[::]:4222", false, true}, + {"loopback v4, auth off", "127.0.0.1:4222", false, false}, + {"loopback v6, auth off", "[::1]:4222", false, false}, + {"localhost name, auth off", "localhost:4222", false, false}, + {"non-loopback, auth on", "0.0.0.0:4222", true, false}, + {"unparseable addr, auth off", "not-an-addr", false, false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var buf bytes.Buffer + logger := slog.New(slog.NewTextHandler(&buf, &slog.HandlerOptions{Level: slog.LevelWarn})) + + warnIfBrokerUnauthenticated(context.Background(), tt.addr, tt.authEnabled, logger) + + got := strings.Contains(buf.String(), "broker is unauthenticated") + if got != tt.wantWarn { + t.Errorf("warn emitted = %v, want %v; log was %q", got, tt.wantWarn, buf.String()) + } + }) + } +} diff --git a/internal/store/fake/fake.go b/internal/store/fake/fake.go index 818b96df..48e9735b 100644 --- a/internal/store/fake/fake.go +++ b/internal/store/fake/fake.go @@ -13,25 +13,27 @@ import ( // Store is an in-memory implementation of [store.Store] for unit tests that // must avoid touching the filesystem. type Store struct { - mu sync.Mutex - farms map[string]store.Farm - queues map[string]store.Queue - storageLocations map[string]store.StorageLocation - computeLocations map[string]store.ComputeLocation - products map[string]store.Product - usagePools map[string]store.UsagePool - usageClaims map[string]store.UsageClaim - workers map[string]store.Worker - jobs map[string]store.Job - jobDependencies map[string][]string // jobID -> upstream IDs (insertion order) - steps map[string]store.Step - tasks map[string]store.Task - taskAttempts map[string]store.TaskAttempt - taskLogs []store.TaskLog - auditEntries []store.AuditEntry - users map[string]store.User - sessions map[string]store.Session - apiKeys map[string]store.APIKey + mu sync.Mutex + farms map[string]store.Farm + queues map[string]store.Queue + storageLocations map[string]store.StorageLocation + computeLocations map[string]store.ComputeLocation + products map[string]store.Product + usagePools map[string]store.UsagePool + usageClaims map[string]store.UsageClaim + workers map[string]store.Worker + jobs map[string]store.Job + jobDependencies map[string][]string // jobID -> upstream IDs (insertion order) + steps map[string]store.Step + tasks map[string]store.Task + taskAttempts map[string]store.TaskAttempt + taskLogs []store.TaskLog + auditEntries []store.AuditEntry + users map[string]store.User + sessions map[string]store.Session + apiKeys map[string]store.APIKey + workerCredentials map[string]store.WorkerCredential + workerJoinTokens map[string]store.WorkerJoinToken } var _ store.Store = (*Store)(nil) @@ -39,24 +41,26 @@ var _ store.Store = (*Store)(nil) // New returns a ready-to-use in-memory store. func New() *Store { return &Store{ - farms: make(map[string]store.Farm), - queues: make(map[string]store.Queue), - storageLocations: make(map[string]store.StorageLocation), - computeLocations: make(map[string]store.ComputeLocation), - products: make(map[string]store.Product), - usagePools: make(map[string]store.UsagePool), - usageClaims: make(map[string]store.UsageClaim), - workers: make(map[string]store.Worker), - jobs: make(map[string]store.Job), - jobDependencies: make(map[string][]string), - steps: make(map[string]store.Step), - tasks: make(map[string]store.Task), - taskAttempts: make(map[string]store.TaskAttempt), - taskLogs: make([]store.TaskLog, 0), - auditEntries: make([]store.AuditEntry, 0), - users: make(map[string]store.User), - sessions: make(map[string]store.Session), - apiKeys: make(map[string]store.APIKey), + farms: make(map[string]store.Farm), + queues: make(map[string]store.Queue), + storageLocations: make(map[string]store.StorageLocation), + computeLocations: make(map[string]store.ComputeLocation), + products: make(map[string]store.Product), + usagePools: make(map[string]store.UsagePool), + usageClaims: make(map[string]store.UsageClaim), + workers: make(map[string]store.Worker), + jobs: make(map[string]store.Job), + jobDependencies: make(map[string][]string), + steps: make(map[string]store.Step), + tasks: make(map[string]store.Task), + taskAttempts: make(map[string]store.TaskAttempt), + taskLogs: make([]store.TaskLog, 0), + auditEntries: make([]store.AuditEntry, 0), + users: make(map[string]store.User), + sessions: make(map[string]store.Session), + apiKeys: make(map[string]store.APIKey), + workerCredentials: make(map[string]store.WorkerCredential), + workerJoinTokens: make(map[string]store.WorkerJoinToken), } } @@ -89,4 +93,6 @@ func (s *Store) Reset() { s.users = make(map[string]store.User) s.sessions = make(map[string]store.Session) s.apiKeys = make(map[string]store.APIKey) + s.workerCredentials = make(map[string]store.WorkerCredential) + s.workerJoinTokens = make(map[string]store.WorkerJoinToken) } diff --git a/internal/store/fake/task_log.go b/internal/store/fake/task_log.go index c91d9721..eee6ad05 100644 --- a/internal/store/fake/task_log.go +++ b/internal/store/fake/task_log.go @@ -9,7 +9,11 @@ import ( "github.com/uberware/sqi/internal/store" ) -// CreateTaskLog implements [store.TaskLogStore]. +// CreateTaskLog implements [store.TaskLogStore]. It does not enforce the +// attempt_id foreign key that the SQLite backend does — a log row is +// appended even when AttemptID names no existing task_attempts row. Tests +// that exercise a log-ingest error path tied to a missing/deleted attempt +// need the real SQLite store, or a wrapper around this one, to observe it. func (s *Store) CreateTaskLog(_ context.Context, log store.TaskLog) (store.TaskLog, error) { s.mu.Lock() defer s.mu.Unlock() diff --git a/internal/store/fake/workercredential.go b/internal/store/fake/workercredential.go new file mode 100644 index 00000000..4fcd25a8 --- /dev/null +++ b/internal/store/fake/workercredential.go @@ -0,0 +1,199 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package fake + +import ( + "context" + "sort" + "time" + + "github.com/uberware/sqi/internal/store" +) + +// CreateWorkerCredential implements [store.WorkerCredentialStore]. +func (s *Store) CreateWorkerCredential(_ context.Context, c store.WorkerCredential) (store.WorkerCredential, error) { + s.mu.Lock() + defer s.mu.Unlock() + if _, ok := s.workerCredentials[c.ID]; ok { + return store.WorkerCredential{}, store.ErrConflict + } + for _, ex := range s.workerCredentials { + // Mirror SQLite's constraints: public_key is UNIQUE across every row, + // active or revoked, but worker_id is only unique among ACTIVE rows + // (worker_credentials_active, a partial index) — a revoked row must + // never block the same worker ID from enrolling again with a new key. + if ex.PublicKey == c.PublicKey { + return store.WorkerCredential{}, store.ErrConflict + } + if ex.WorkerID == c.WorkerID && ex.RevokedAt == nil { + return store.WorkerCredential{}, store.ErrConflict + } + } + s.workerCredentials[c.ID] = c + return c, nil +} + +// GetActiveWorkerCredentialByWorkerID implements [store.WorkerCredentialStore]. +func (s *Store) GetActiveWorkerCredentialByWorkerID(_ context.Context, workerID string) (store.WorkerCredential, error) { + s.mu.Lock() + defer s.mu.Unlock() + for _, c := range s.workerCredentials { + if c.WorkerID == workerID && c.RevokedAt == nil { + return c, nil + } + } + return store.WorkerCredential{}, store.ErrNotFound +} + +// ListActiveWorkerCredentials implements [store.WorkerCredentialStore]. +func (s *Store) ListActiveWorkerCredentials(_ context.Context) ([]store.WorkerCredential, error) { + s.mu.Lock() + defer s.mu.Unlock() + var out []store.WorkerCredential + for _, c := range s.workerCredentials { + if c.RevokedAt == nil { + out = append(out, c) + } + } + sort.Slice(out, func(i, j int) bool { return out[i].EnrolledAt.Before(out[j].EnrolledAt) }) + return out, nil +} + +// RevokeWorkerCredential implements [store.WorkerCredentialStore]. +func (s *Store) RevokeWorkerCredential(_ context.Context, workerID string, at time.Time) error { + s.mu.Lock() + defer s.mu.Unlock() + // Scan the WHOLE map before giving up. Go's map iteration order is + // randomized, and after a key rotation a worker can legitimately have + // both a revoked row and an active one for the same worker_id — mirror + // SQLite's "UPDATE ... WHERE worker_id = ? AND revoked_at IS NULL", + // which matches by predicate regardless of row count, rather than + // stopping at the first row this map happens to yield. Returning + // ErrNotFound on hitting a revoked row before the active one would make + // a legitimate revoke fail non-deterministically. + for id, c := range s.workerCredentials { + if c.WorkerID != workerID || c.RevokedAt != nil { + continue + } + c.RevokedAt = &at + s.workerCredentials[id] = c + return nil + } + return store.ErrNotFound +} + +// TouchWorkerCredential implements [store.WorkerCredentialStore]. +func (s *Store) TouchWorkerCredential(_ context.Context, workerID string, at time.Time) error { + s.mu.Lock() + defer s.mu.Unlock() + // Scan the whole map rather than stopping at the first match, for the + // same reason RevokeWorkerCredential does: a worker can have both a + // revoked row and an active one after a key rotation, and map iteration + // order is randomized. + for id, c := range s.workerCredentials { + if c.WorkerID != workerID || c.RevokedAt != nil { + continue + } + c.LastSeenAt = &at + s.workerCredentials[id] = c + return nil + } + return store.ErrNotFound +} + +// CreateWorkerJoinToken implements [store.WorkerCredentialStore]. +func (s *Store) CreateWorkerJoinToken(_ context.Context, t store.WorkerJoinToken) (store.WorkerJoinToken, error) { + s.mu.Lock() + defer s.mu.Unlock() + if _, ok := s.workerJoinTokens[t.ID]; ok { + return store.WorkerJoinToken{}, store.ErrConflict + } + for _, ex := range s.workerJoinTokens { + if ex.TokenHash == t.TokenHash { + return store.WorkerJoinToken{}, store.ErrConflict + } + } + s.workerJoinTokens[t.ID] = t + return t, nil +} + +// GetWorkerJoinTokenByHash implements [store.WorkerCredentialStore]. +func (s *Store) GetWorkerJoinTokenByHash(_ context.Context, hash string) (store.WorkerJoinToken, error) { + s.mu.Lock() + defer s.mu.Unlock() + for _, t := range s.workerJoinTokens { + if t.TokenHash == hash { + return t, nil + } + } + return store.WorkerJoinToken{}, store.ErrNotFound +} + +// MarkWorkerJoinTokenUsed implements [store.WorkerCredentialStore]. +func (s *Store) MarkWorkerJoinTokenUsed(_ context.Context, id string, at time.Time) error { + s.mu.Lock() + defer s.mu.Unlock() + t, ok := s.workerJoinTokens[id] + if !ok { + return store.ErrNotFound + } + t.UsedAt = &at + s.workerJoinTokens[id] = t + return nil +} + +// RedeemWorkerJoinToken implements [store.WorkerCredentialStore]. +// +// Mirrors SQLite's transaction: the token claim and the credential creation +// happen under ONE hold of the mutex, so a caller never observes the token +// consumed without the credential existing, or the reverse. An unknown, +// expired or already-claimed token is store.ErrNotFound with cred left +// uncreated; a cred that collides with an existing worker_id or public_key +// is store.ErrConflict with the token left unclaimed — mirroring SQLite's +// rollback, since nothing here is written until both checks pass. +func (s *Store) RedeemWorkerJoinToken(_ context.Context, hash string, now time.Time, cred store.WorkerCredential) (store.WorkerCredential, error) { + s.mu.Lock() + defer s.mu.Unlock() + + var tokID string + found := false + for id, t := range s.workerJoinTokens { + if t.TokenHash != hash || t.UsedAt != nil { + continue + } + // Strictly after, matching SQL's "expires_at > ?": a token whose + // expiry is exactly now is expired. + if !t.ExpiresAt.After(now) { + continue + } + tokID = id + found = true + break + } + if !found { + return store.WorkerCredential{}, store.ErrNotFound + } + + // Same conflict rules as CreateWorkerCredential, inlined rather than + // called: that method takes s.mu itself, and this method already holds + // it for the whole claim-and-create span. + if _, ok := s.workerCredentials[cred.ID]; ok { + return store.WorkerCredential{}, store.ErrConflict + } + for _, ex := range s.workerCredentials { + if ex.PublicKey == cred.PublicKey { + return store.WorkerCredential{}, store.ErrConflict + } + if ex.WorkerID == cred.WorkerID && ex.RevokedAt == nil { + return store.WorkerCredential{}, store.ErrConflict + } + } + + tok := s.workerJoinTokens[tokID] + at := now + tok.UsedAt = &at + s.workerJoinTokens[tokID] = tok + + s.workerCredentials[cred.ID] = cred + return cred, nil +} diff --git a/internal/store/fake/workercredential_test.go b/internal/store/fake/workercredential_test.go new file mode 100644 index 00000000..571fe2f5 --- /dev/null +++ b/internal/store/fake/workercredential_test.go @@ -0,0 +1,528 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package fake_test + +import ( + "context" + "errors" + "fmt" + "sync" + "testing" + "time" + + "github.com/uberware/sqi/internal/store" + "github.com/uberware/sqi/internal/store/fake" +) + +// TestWorkerCredential_RevokeActiveAfterRotation exercises +// RevokeWorkerCredential against a worker that has one revoked row and one +// active row for the same worker_id — the state a key rotation (revoke, then +// re-enroll with a new key) leaves behind, and a state that only became +// legal once worker_id uniqueness was scoped to active rows. +// +// Go's map iteration order is randomized. A version of RevokeWorkerCredential +// that stops at the FIRST row matching the worker ID, rather than scanning +// for one that is also active, fails non-deterministically here: whenever +// iteration reaches the already-revoked row before the active one, it +// returns ErrNotFound for what is otherwise a completely legitimate revoke. +// +// The test seeds a pile of unrelated "noise" workers so the map has enough +// entries for order to matter, and repeats the whole scenario across many +// subtests (fresh store, fresh map, fresh iteration order each time) so a +// fix that only happens to work by accident of one run's order cannot pass +// by luck. +func TestWorkerCredential_RevokeActiveAfterRotation(t *testing.T) { + for run := range 50 { + t.Run(fmt.Sprintf("run-%d", run), func(t *testing.T) { + s := fake.New() + defer s.Close() + ctx := context.Background() + now := time.Now().UTC() + + // Noise: several other fully-revoked workers, planted before and + // after the worker under test, so a worker_id match with no + // revoked_at filter has other rows it could land on first. + for i := range 20 { + wid := fmt.Sprintf("noise-%d", i) + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "noise-wc-" + wid, WorkerID: wid, PublicKey: "noise-pub-" + wid, EnrolledAt: now, + }); err != nil { + t.Fatalf("CreateWorkerCredential (noise %d): %v", i, err) + } + } + + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: now, + }); err != nil { + t.Fatalf("CreateWorkerCredential (first): %v", err) + } + if err := s.RevokeWorkerCredential(ctx, "w1", now); err != nil { + t.Fatalf("RevokeWorkerCredential (first): %v", err) + } + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc2", WorkerID: "w1", PublicKey: "pub2", EnrolledAt: now, + }); err != nil { + t.Fatalf("CreateWorkerCredential (rotated): %v", err) + } + + // w1 now has one revoked row (wc1) and one active row (wc2). + // Revoking the active credential must succeed regardless of + // which row the map iteration reaches first. + if err := s.RevokeWorkerCredential(ctx, "w1", now.Add(time.Minute)); err != nil { + t.Fatalf("RevokeWorkerCredential (active credential, post-rotation): %v", err) + } + + if _, err := s.GetActiveWorkerCredentialByWorkerID(ctx, "w1"); !errors.Is(err, store.ErrNotFound) { + t.Errorf("expected ErrNotFound after revoking the only active credential, got %v", err) + } + }) + } +} + +// TestWorkerCredential_GetActive_ReturnsActiveRowAfterRotation verifies that +// GetActiveWorkerCredentialByWorkerID resolves the ambiguity a rotation +// creates: it must return the active (rotated) row, never the revoked one. +func TestWorkerCredential_GetActive_ReturnsActiveRowAfterRotation(t *testing.T) { + s := fake.New() + defer s.Close() + ctx := context.Background() + now := time.Now().UTC() + + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: now, + }); err != nil { + t.Fatalf("CreateWorkerCredential (first): %v", err) + } + if err := s.RevokeWorkerCredential(ctx, "w1", now); err != nil { + t.Fatalf("RevokeWorkerCredential: %v", err) + } + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc2", WorkerID: "w1", PublicKey: "pub2", EnrolledAt: now, + }); err != nil { + t.Fatalf("CreateWorkerCredential (rotated): %v", err) + } + + got, err := s.GetActiveWorkerCredentialByWorkerID(ctx, "w1") + if err != nil { + t.Fatalf("GetActiveWorkerCredentialByWorkerID: %v", err) + } + if got.PublicKey != "pub2" { + t.Errorf("PublicKey = %q, want %q (the active, rotated key)", got.PublicKey, "pub2") + } + if got.RevokedAt != nil { + t.Errorf("RevokedAt = %v, want nil", got.RevokedAt) + } +} + +// TestWorkerCredential_GetActive_RevokedOnlyReturnsNotFound verifies that a +// worker whose only credential has been revoked (no rotation) reports no +// active credential, mirroring the SQLite backend's contract exactly. +func TestWorkerCredential_GetActive_RevokedOnlyReturnsNotFound(t *testing.T) { + s := fake.New() + defer s.Close() + ctx := context.Background() + now := time.Now().UTC() + + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: now, + }); err != nil { + t.Fatalf("CreateWorkerCredential: %v", err) + } + if err := s.RevokeWorkerCredential(ctx, "w1", now); err != nil { + t.Fatalf("RevokeWorkerCredential: %v", err) + } + + _, err := s.GetActiveWorkerCredentialByWorkerID(ctx, "w1") + if !errors.Is(err, store.ErrNotFound) { + t.Errorf("expected ErrNotFound for a worker with only a revoked credential, got %v", err) + } +} + +// TestWorkerCredential_Touch mirrors sqlite_test.TestWorkerCredential_Touch: +// touching a worker's active credential sets LastSeenAt. +func TestWorkerCredential_Touch(t *testing.T) { + s := fake.New() + defer s.Close() + ctx := context.Background() + enrolledAt := time.Now().UTC() + + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: enrolledAt, + }); err != nil { + t.Fatalf("CreateWorkerCredential: %v", err) + } + + seenAt := enrolledAt.Add(time.Hour) + if err := s.TouchWorkerCredential(ctx, "w1", seenAt); err != nil { + t.Fatalf("TouchWorkerCredential: %v", err) + } + + got, err := s.GetActiveWorkerCredentialByWorkerID(ctx, "w1") + if err != nil { + t.Fatalf("GetActiveWorkerCredentialByWorkerID: %v", err) + } + if got.LastSeenAt == nil || !got.LastSeenAt.Equal(seenAt) { + t.Errorf("LastSeenAt = %v, want %v", got.LastSeenAt, seenAt) + } +} + +// TestWorkerCredential_TouchNotFound mirrors +// sqlite_test.TestWorkerCredential_TouchNotFound. +func TestWorkerCredential_TouchNotFound(t *testing.T) { + s := fake.New() + defer s.Close() + err := s.TouchWorkerCredential(context.Background(), "nope", time.Now().UTC()) + if !errors.Is(err, store.ErrNotFound) { + t.Errorf("expected ErrNotFound, got %v", err) + } +} + +// TestWorkerCredential_TouchOnlyMatchesActiveRow mirrors +// sqlite_test.TestWorkerCredential_TouchOnlyMatchesActiveRow: a worker with +// only a revoked credential must not be touchable through that stale row. +func TestWorkerCredential_TouchOnlyMatchesActiveRow(t *testing.T) { + s := fake.New() + defer s.Close() + ctx := context.Background() + now := time.Now().UTC() + + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: now, + }); err != nil { + t.Fatalf("CreateWorkerCredential: %v", err) + } + if err := s.RevokeWorkerCredential(ctx, "w1", now); err != nil { + t.Fatalf("RevokeWorkerCredential: %v", err) + } + + err := s.TouchWorkerCredential(ctx, "w1", now.Add(time.Hour)) + if !errors.Is(err, store.ErrNotFound) { + t.Errorf("expected ErrNotFound for a worker with only a revoked credential, got %v", err) + } +} + +// TestWorkerCredential_DuplicateWorkerID mirrors +// sqlite_test.TestWorkerCredential_DuplicateWorkerID: a worker ID with an +// existing ACTIVE credential cannot be double-enrolled. It exercises +// CreateWorkerCredential's own conflict logic directly, rather than only +// indirectly through the rotation tests above. +func TestWorkerCredential_DuplicateWorkerID(t *testing.T) { + s := fake.New() + defer s.Close() + ctx := context.Background() + now := time.Now().UTC() + + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: now, + }); err != nil { + t.Fatalf("CreateWorkerCredential: %v", err) + } + _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc2", WorkerID: "w1", PublicKey: "pub2", EnrolledAt: now, + }) + if !errors.Is(err, store.ErrConflict) { + t.Errorf("expected ErrConflict for duplicate active worker_id, got %v", err) + } +} + +// TestWorkerCredential_DuplicatePublicKey mirrors +// sqlite_test.TestWorkerCredential_DuplicatePublicKey: two different +// workers can never share a public key, active-vs-active. +func TestWorkerCredential_DuplicatePublicKey(t *testing.T) { + s := fake.New() + defer s.Close() + ctx := context.Background() + now := time.Now().UTC() + + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: now, + }); err != nil { + t.Fatalf("CreateWorkerCredential: %v", err) + } + _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc2", WorkerID: "w2", PublicKey: "pub1", EnrolledAt: now, + }) + if !errors.Is(err, store.ErrConflict) { + t.Errorf("expected ErrConflict for duplicate public_key, got %v", err) + } +} + +// TestWorkerCredential_RevokedPublicKeyStillUnique mirrors +// sqlite_test.TestWorkerCredential_RevokedPublicKeyStillUnique. The whole +// point of leaving public_key globally UNIQUE (not scoped to active rows the +// way worker_id now is) is that a rotated-away key must never become usable +// again — this pins that in the fake as well as SQLite, so public_key +// uniqueness can never silently regress to active-rows-only scoping without +// a test failing in both backends. If it did regress, a revoked credential's +// key would become available for reuse, defeating the reason rotation is +// safe in the first place. +func TestWorkerCredential_RevokedPublicKeyStillUnique(t *testing.T) { + s := fake.New() + defer s.Close() + ctx := context.Background() + now := time.Now().UTC() + + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: now, + }); err != nil { + t.Fatalf("CreateWorkerCredential: %v", err) + } + if err := s.RevokeWorkerCredential(ctx, "w1", now); err != nil { + t.Fatalf("RevokeWorkerCredential: %v", err) + } + + _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc2", WorkerID: "w2", PublicKey: "pub1", EnrolledAt: now, + }) + if !errors.Is(err, store.ErrConflict) { + t.Errorf("expected ErrConflict reusing a revoked credential's public_key, got %v", err) + } +} + +// TestWorkerCredential_ListActiveAfterRotationHasExactlyOneRow mirrors +// sqlite_test.TestWorkerCredential_ListActiveAfterRotationHasExactlyOneRow: +// after enroll, revoke, re-enroll, ListActiveWorkerCredentials must return +// exactly one row for the rotated worker — never zero (the rotation +// silently failing) and never two. Two simultaneously-active rows for one +// worker_id would mean the broker's authorized-key set (which this method +// feeds) accepts both the old and the new key for that worker at once, +// which is not a rotation at all. +func TestWorkerCredential_ListActiveAfterRotationHasExactlyOneRow(t *testing.T) { + s := fake.New() + defer s.Close() + ctx := context.Background() + now := time.Now().UTC() + + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: now, + }); err != nil { + t.Fatalf("CreateWorkerCredential (first): %v", err) + } + if err := s.RevokeWorkerCredential(ctx, "w1", now); err != nil { + t.Fatalf("RevokeWorkerCredential: %v", err) + } + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc2", WorkerID: "w1", PublicKey: "pub2", EnrolledAt: now, + }); err != nil { + t.Fatalf("CreateWorkerCredential (rotated): %v", err) + } + + active, err := s.ListActiveWorkerCredentials(ctx) + if err != nil { + t.Fatalf("ListActiveWorkerCredentials: %v", err) + } + var forW1 []store.WorkerCredential + for _, c := range active { + if c.WorkerID == "w1" { + forW1 = append(forW1, c) + } + } + if len(forW1) != 1 { + t.Fatalf("want exactly 1 active credential for w1 after rotation, got %d", len(forW1)) + } + if forW1[0].PublicKey != "pub2" { + t.Errorf("active credential PublicKey = %q, want %q (the rotated key)", forW1[0].PublicKey, "pub2") + } +} + +// ── RedeemWorkerJoinToken ──────────────────────────────────────────────────── +// +// Mirrors sqlite_test's RedeemWorkerJoinToken cases. The fake is what every +// internal/api test runs against, so a fake whose claim is not atomic with +// the credential creation — or whose expiry boundary differs from SQLite's +// "expires_at > ?" — would let the handler's concurrency and rollback tests +// pass over a store that does not behave like the real one. + +func TestWorkerJoinToken_Redeem(t *testing.T) { + s := fake.New() + defer s.Close() + ctx := context.Background() + now := time.Now().UTC() + + if _, err := s.CreateWorkerJoinToken(ctx, store.WorkerJoinToken{ + ID: "jt1", TokenHash: "hash1", Prefix: "sqiwjt_abcd", + ExpiresAt: now.Add(time.Hour), CreatedAt: now, + }); err != nil { + t.Fatalf("CreateWorkerJoinToken: %v", err) + } + + claimedAt := now.Add(time.Minute) + got, err := s.RedeemWorkerJoinToken(ctx, "hash1", claimedAt, store.WorkerCredential{ + ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: claimedAt, + }) + if err != nil { + t.Fatalf("RedeemWorkerJoinToken: %v", err) + } + if got.ID != "wc1" || got.WorkerID != "w1" { + t.Errorf("got %+v, want the created credential for w1", got) + } + + stored, err := s.GetWorkerJoinTokenByHash(ctx, "hash1") + if err != nil { + t.Fatalf("GetWorkerJoinTokenByHash: %v", err) + } + if stored.UsedAt == nil || !stored.UsedAt.Equal(claimedAt) { + t.Errorf("UsedAt: got %v, want %v", stored.UsedAt, claimedAt) + } + + storedCred, err := s.GetActiveWorkerCredentialByWorkerID(ctx, "w1") + if err != nil { + t.Fatalf("GetActiveWorkerCredentialByWorkerID: %v", err) + } + if storedCred.PublicKey != "pub1" { + t.Errorf("PublicKey: got %q, want %q", storedCred.PublicKey, "pub1") + } + + if _, err := s.RedeemWorkerJoinToken(ctx, "hash1", claimedAt.Add(time.Second), store.WorkerCredential{ + ID: "wc2", WorkerID: "w2", PublicKey: "pub2", EnrolledAt: claimedAt, + }); !errors.Is(err, store.ErrNotFound) { + t.Errorf("second RedeemWorkerJoinToken: got %v, want store.ErrNotFound", err) + } +} + +func TestWorkerJoinToken_RedeemExpired(t *testing.T) { + s := fake.New() + defer s.Close() + ctx := context.Background() + now := time.Now().UTC() + + if _, err := s.CreateWorkerJoinToken(ctx, store.WorkerJoinToken{ + ID: "jt1", TokenHash: "hash1", Prefix: "sqiwjt_abcd", + ExpiresAt: now, CreatedAt: now.Add(-time.Hour), + }); err != nil { + t.Fatalf("CreateWorkerJoinToken: %v", err) + } + cred := store.WorkerCredential{ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: now} + + // Exactly at expiry counts as expired, matching SQLite's expires_at > ?. + if _, err := s.RedeemWorkerJoinToken(ctx, "hash1", now, cred); !errors.Is(err, store.ErrNotFound) { + t.Errorf("RedeemWorkerJoinToken at the expiry instant: got %v, want store.ErrNotFound", err) + } + if _, err := s.RedeemWorkerJoinToken(ctx, "hash1", now.Add(time.Hour), cred); !errors.Is(err, store.ErrNotFound) { + t.Errorf("RedeemWorkerJoinToken after expiry: got %v, want store.ErrNotFound", err) + } + + stored, err := s.GetWorkerJoinTokenByHash(ctx, "hash1") + if err != nil { + t.Fatalf("GetWorkerJoinTokenByHash: %v", err) + } + if stored.UsedAt != nil { + t.Error("a refused claim marked the token used") + } + if _, err := s.GetActiveWorkerCredentialByWorkerID(ctx, "w1"); !errors.Is(err, store.ErrNotFound) { + t.Errorf("expected no credential to have been created, got %v", err) + } +} + +func TestWorkerJoinToken_RedeemUnknown(t *testing.T) { + s := fake.New() + defer s.Close() + cred := store.WorkerCredential{ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: time.Now().UTC()} + if _, err := s.RedeemWorkerJoinToken(context.Background(), "nope", time.Now().UTC(), cred); !errors.Is(err, store.ErrNotFound) { + t.Errorf("expected ErrNotFound, got %v", err) + } +} + +// TestWorkerJoinToken_RedeemConflictRollsBackClaim mirrors +// sqlite_test.TestWorkerJoinToken_RedeemConflictRollsBackClaim: a credential +// conflict must roll back the token claim too. +func TestWorkerJoinToken_RedeemConflictRollsBackClaim(t *testing.T) { + s := fake.New() + defer s.Close() + ctx := context.Background() + now := time.Now().UTC() + + if _, err := s.CreateWorkerJoinToken(ctx, store.WorkerJoinToken{ + ID: "jt1", TokenHash: "hash1", Prefix: "sqiwjt_abcd", + ExpiresAt: now.Add(time.Hour), CreatedAt: now, + }); err != nil { + t.Fatalf("CreateWorkerJoinToken: %v", err) + } + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc-existing", WorkerID: "w1", PublicKey: "pub-existing", EnrolledAt: now, + }); err != nil { + t.Fatalf("seed CreateWorkerCredential: %v", err) + } + + _, err := s.RedeemWorkerJoinToken(ctx, "hash1", now.Add(time.Minute), store.WorkerCredential{ + ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: now, + }) + if !errors.Is(err, store.ErrConflict) { + t.Fatalf("RedeemWorkerJoinToken: got %v, want store.ErrConflict", err) + } + + stored, err := s.GetWorkerJoinTokenByHash(ctx, "hash1") + if err != nil { + t.Fatalf("GetWorkerJoinTokenByHash: %v", err) + } + if stored.UsedAt != nil { + t.Error("a conflicting redemption spent the token; it must roll back") + } + + got, err := s.RedeemWorkerJoinToken(ctx, "hash1", now.Add(time.Minute), store.WorkerCredential{ + ID: "wc2", WorkerID: "w3", PublicKey: "pub3", EnrolledAt: now, + }) + if err != nil { + t.Fatalf("re-redemption after conflict: %v", err) + } + if got.WorkerID != "w3" { + t.Errorf("WorkerID = %q, want w3", got.WorkerID) + } +} + +// TestWorkerJoinToken_RedeemIsAtomic is the fake's half of the invariant +// SQLite gets from one transaction: concurrent redemptions of one +// single-use token, each for a distinct (non-conflicting) worker ID and +// public key, must yield exactly one winner. Repeated rounds, because a +// single burst can serialize by chance. +func TestWorkerJoinToken_RedeemIsAtomic(t *testing.T) { + const rounds = 50 + const attempts = 8 + + for round := range rounds { + s := fake.New() + ctx := context.Background() + now := time.Now().UTC() + if _, err := s.CreateWorkerJoinToken(ctx, store.WorkerJoinToken{ + ID: "jt1", TokenHash: "hash1", Prefix: "sqiwjt_abcd", + ExpiresAt: now.Add(time.Hour), CreatedAt: now, + }); err != nil { + t.Fatalf("CreateWorkerJoinToken: %v", err) + } + + errs := make([]error, attempts) + var ready, done sync.WaitGroup + ready.Add(attempts) + done.Add(attempts) + start := make(chan struct{}) + for i := range attempts { + go func() { + defer done.Done() + ready.Done() + <-start + _, errs[i] = s.RedeemWorkerJoinToken(ctx, "hash1", now.Add(time.Minute), store.WorkerCredential{ + ID: fmt.Sprintf("wc%d", i), WorkerID: fmt.Sprintf("w%d", i), + PublicKey: fmt.Sprintf("pub%d", i), EnrolledAt: now, + }) + }() + } + ready.Wait() + close(start) + done.Wait() + + won := 0 + for i, err := range errs { + switch { + case err == nil: + won++ + case errors.Is(err, store.ErrNotFound): + default: + t.Fatalf("round %d attempt %d: unexpected error %v", round, i, err) + } + } + if won != 1 { + t.Fatalf("round %d: %d of %d concurrent redemptions succeeded, want exactly 1", round, won, attempts) + } + s.Close() + } +} diff --git a/internal/store/migrations/00030_broker_auth.sql b/internal/store/migrations/00030_broker_auth.sql new file mode 100644 index 00000000..033637c5 --- /dev/null +++ b/internal/store/migrations/00030_broker_auth.sql @@ -0,0 +1,46 @@ +-- SPDX-License-Identifier: AGPL-3.0-or-later + +-- +goose Up + +-- worker_credentials binds a worker's self-chosen ID to an Ed25519 nkey +-- public key. worker_id is deliberately NOT a foreign key to workers(id): a +-- credential is issued before the worker has ever registered, and on +-- auth-off farms worker rows exist with no credential at all. +-- +-- worker_id is uniqued by a PARTIAL index over active rows only (below), not +-- a column-level UNIQUE: a revoked row must not block the same worker ID +-- from enrolling again with a new key, or revocation would be a one-way +-- door and a worker with a lost or compromised seed could never recover its +-- identity. public_key stays globally UNIQUE — a key must never be reused, +-- even after the credential that held it is revoked. +CREATE TABLE worker_credentials ( + id TEXT PRIMARY KEY, + worker_id TEXT NOT NULL, + public_key TEXT NOT NULL UNIQUE, + name TEXT NOT NULL DEFAULT '', + enrolled_at TEXT NOT NULL, + last_seen_at TEXT, + revoked_at TEXT +); +CREATE UNIQUE INDEX worker_credentials_active ON worker_credentials (worker_id) WHERE revoked_at IS NULL; + +-- worker_join_tokens mirrors api_keys (00022): only the hash is stored, and +-- the prefix exists so an operator can identify a token in a list without +-- the server ever holding the secret. +CREATE TABLE worker_join_tokens ( + id TEXT PRIMARY KEY, + token_hash TEXT NOT NULL UNIQUE, + prefix TEXT NOT NULL, + name TEXT NOT NULL DEFAULT '', + expires_at TEXT NOT NULL, + used_at TEXT, + created_by TEXT NOT NULL DEFAULT '', + created_at TEXT NOT NULL +); +CREATE INDEX worker_join_tokens_expires ON worker_join_tokens (expires_at); + +-- +goose Down +DROP INDEX worker_join_tokens_expires; +DROP TABLE worker_join_tokens; +DROP INDEX worker_credentials_active; +DROP TABLE worker_credentials; diff --git a/internal/store/sqlite/store.go b/internal/store/sqlite/store.go index fbfb5dd1..fcc7b70f 100644 --- a/internal/store/sqlite/store.go +++ b/internal/store/sqlite/store.go @@ -247,6 +247,16 @@ type Store struct { stmtListAPIKeysForUser *sql.Stmt stmtRevokeAPIKey *sql.Stmt stmtTouchAPIKeyLastUsed *sql.Stmt + + // ── worker credentials & join tokens ──────────────────────────────────── + stmtInsertWorkerCredential *sql.Stmt + stmtGetActiveWorkerCredentialByWorkerID *sql.Stmt + stmtListActiveWorkerCredentials *sql.Stmt + stmtRevokeWorkerCredential *sql.Stmt + stmtTouchWorkerCredential *sql.Stmt + stmtInsertWorkerJoinToken *sql.Stmt + stmtGetWorkerJoinTokenByHash *sql.Stmt + stmtMarkWorkerJoinTokenUsed *sql.Stmt } // Open opens (or creates) the SQLite database at path, applies connection @@ -813,5 +823,34 @@ func (s *Store) prepareAll(ctx context.Context) error { return err } + // ── worker credentials & join tokens ───────────────────────────────────── + if s.stmtInsertWorkerCredential, err = s.prepare(ctx, sqlInsertWorkerCredential); err != nil { + return err + } + if s.stmtGetActiveWorkerCredentialByWorkerID, err = s.prepare(ctx, sqlGetActiveWorkerCredentialByWorkerID); err != nil { + return err + } + if s.stmtListActiveWorkerCredentials, err = s.prepare(ctx, sqlListActiveWorkerCredentials); err != nil { + return err + } + if s.stmtRevokeWorkerCredential, err = s.prepare(ctx, sqlRevokeWorkerCredential); err != nil { + return err + } + if s.stmtTouchWorkerCredential, err = s.prepare(ctx, sqlTouchWorkerCredential); err != nil { + return err + } + if s.stmtInsertWorkerJoinToken, err = s.prepare(ctx, sqlInsertWorkerJoinToken); err != nil { + return err + } + if s.stmtGetWorkerJoinTokenByHash, err = s.prepare(ctx, sqlGetWorkerJoinTokenByHash); err != nil { + return err + } + if s.stmtMarkWorkerJoinTokenUsed, err = s.prepare(ctx, sqlMarkWorkerJoinTokenUsed); err != nil { + return err + } + // sqlConsumeWorkerJoinToken has no prepared statement of its own: it is + // used only as raw SQL text inside RedeemWorkerJoinToken's transaction + // (workercredential.go), matching CreateJobSubmission's idiom. + return nil } diff --git a/internal/store/sqlite/workercredential.go b/internal/store/sqlite/workercredential.go new file mode 100644 index 00000000..2804a32d --- /dev/null +++ b/internal/store/sqlite/workercredential.go @@ -0,0 +1,197 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package sqlite + +import ( + "context" + "database/sql" + "time" + + "github.com/uberware/sqi/internal/store" +) + +const ( + sqlInsertWorkerCredential = `INSERT INTO worker_credentials (id, worker_id, public_key, name, enrolled_at, last_seen_at, revoked_at) VALUES (?, ?, ?, ?, ?, ?, ?) RETURNING id, worker_id, public_key, name, enrolled_at, last_seen_at, revoked_at` + + // gosec G101 flags these consts on the false-positive belief that an + // identifier containing "Credential" is a hardcoded credential; the + // value is SQL text, not a secret. The nolint must be a trailing comment + // on the declaration's own line (where gosec attributes the finding), so + // these stay single-line rather than the multi-line style used + // elsewhere. + sqlGetActiveWorkerCredentialByWorkerID = `SELECT id, worker_id, public_key, name, enrolled_at, last_seen_at, revoked_at FROM worker_credentials WHERE worker_id = ? AND revoked_at IS NULL` //nolint:gosec // G101: SQL text, not a credential + + sqlListActiveWorkerCredentials = `SELECT id, worker_id, public_key, name, enrolled_at, last_seen_at, revoked_at FROM worker_credentials WHERE revoked_at IS NULL ORDER BY enrolled_at` //nolint:gosec // G101: SQL text, not a credential + + sqlRevokeWorkerCredential = `UPDATE worker_credentials SET revoked_at = ? WHERE worker_id = ? AND revoked_at IS NULL` //nolint:gosec // G101: SQL text, not a credential + + sqlTouchWorkerCredential = `UPDATE worker_credentials SET last_seen_at = ? WHERE worker_id = ? AND revoked_at IS NULL` //nolint:gosec // G101: SQL text, not a credential + + sqlInsertWorkerJoinToken = `INSERT INTO worker_join_tokens (id, token_hash, prefix, name, expires_at, used_at, created_by, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?) RETURNING id, token_hash, prefix, name, expires_at, used_at, created_by, created_at` + + sqlGetWorkerJoinTokenByHash = `SELECT id, token_hash, prefix, name, expires_at, used_at, created_by, created_at FROM worker_join_tokens WHERE token_hash = ?` //nolint:gosec // G101: SQL text, not a credential + + sqlMarkWorkerJoinTokenUsed = `UPDATE worker_join_tokens SET used_at = ? WHERE id = ?` + + // sqlConsumeWorkerJoinToken claims a single-use token in ONE statement: + // the used_at IS NULL and expires_at > ? predicates are the validity + // check, and the SET is the claim. Two concurrent enrollments presenting + // the same token therefore cannot both succeed — the second matches no + // row and gets store.ErrNotFound, indistinguishable from an unknown + // token, which is what the unauthenticated enroll endpoint needs. Used + // only inside RedeemWorkerJoinToken's transaction (raw text via tx, not + // a prepared statement — see that method's doc comment), never on its + // own: a token claimed here has to be followed by successfully creating + // the credential it authorizes, in the same transaction, or the claim + // itself must roll back. + sqlConsumeWorkerJoinToken = `UPDATE worker_join_tokens SET used_at = ? WHERE token_hash = ? AND used_at IS NULL AND expires_at > ? RETURNING id, token_hash, prefix, name, expires_at, used_at, created_by, created_at` +) + +func scanWorkerCredential(row scanner) (store.WorkerCredential, error) { + var c store.WorkerCredential + var lastSeenAt, revokedAt sql.NullString + var enrolledAt string + if err := row.Scan(&c.ID, &c.WorkerID, &c.PublicKey, &c.Name, + &enrolledAt, &lastSeenAt, &revokedAt); err != nil { + return store.WorkerCredential{}, err + } + c.EnrolledAt = mustTime(enrolledAt) + c.LastSeenAt = nullTextToTime(lastSeenAt) + c.RevokedAt = nullTextToTime(revokedAt) + return c, nil +} + +func scanWorkerJoinToken(row scanner) (store.WorkerJoinToken, error) { + var t store.WorkerJoinToken + var usedAt sql.NullString + var expiresAt, createdAt string + if err := row.Scan(&t.ID, &t.TokenHash, &t.Prefix, &t.Name, + &expiresAt, &usedAt, &t.CreatedBy, &createdAt); err != nil { + return store.WorkerJoinToken{}, err + } + t.ExpiresAt = mustTime(expiresAt) + t.UsedAt = nullTextToTime(usedAt) + t.CreatedAt = mustTime(createdAt) + return t, nil +} + +// CreateWorkerCredential implements [store.WorkerCredentialStore]. +func (s *Store) CreateWorkerCredential(ctx context.Context, c store.WorkerCredential) (store.WorkerCredential, error) { + row := s.stmtInsertWorkerCredential.QueryRowContext(ctx, c.ID, c.WorkerID, c.PublicKey, c.Name, + timeToText(c.EnrolledAt), nullTimeToText(c.LastSeenAt), nullTimeToText(c.RevokedAt)) + out, err := scanWorkerCredential(row) + return out, mapErr(err) +} + +// GetActiveWorkerCredentialByWorkerID implements [store.WorkerCredentialStore]. +func (s *Store) GetActiveWorkerCredentialByWorkerID(ctx context.Context, workerID string) (store.WorkerCredential, error) { + row := s.stmtGetActiveWorkerCredentialByWorkerID.QueryRowContext(ctx, workerID) + out, err := scanWorkerCredential(row) + return out, mapErr(err) +} + +// ListActiveWorkerCredentials implements [store.WorkerCredentialStore]. +func (s *Store) ListActiveWorkerCredentials(ctx context.Context) ([]store.WorkerCredential, error) { + rows, err := s.stmtListActiveWorkerCredentials.QueryContext(ctx) + if err != nil { + return nil, mapErr(err) + } + defer rows.Close() + var out []store.WorkerCredential + for rows.Next() { + c, err := scanWorkerCredential(rows) + if err != nil { + return nil, mapErr(err) + } + out = append(out, c) + } + return out, mapErr(rows.Err()) +} + +// RevokeWorkerCredential implements [store.WorkerCredentialStore]. +func (s *Store) RevokeWorkerCredential(ctx context.Context, workerID string, at time.Time) error { + res, err := s.stmtRevokeWorkerCredential.ExecContext(ctx, timeToText(at), workerID) + if err != nil { + return mapErr(err) + } + return checkRowsAffected(res) +} + +// TouchWorkerCredential implements [store.WorkerCredentialStore]. +func (s *Store) TouchWorkerCredential(ctx context.Context, workerID string, at time.Time) error { + res, err := s.stmtTouchWorkerCredential.ExecContext(ctx, timeToText(at), workerID) + if err != nil { + return mapErr(err) + } + return checkRowsAffected(res) +} + +// CreateWorkerJoinToken implements [store.WorkerCredentialStore]. +func (s *Store) CreateWorkerJoinToken(ctx context.Context, t store.WorkerJoinToken) (store.WorkerJoinToken, error) { + row := s.stmtInsertWorkerJoinToken.QueryRowContext(ctx, t.ID, t.TokenHash, t.Prefix, t.Name, + timeToText(t.ExpiresAt), nullTimeToText(t.UsedAt), t.CreatedBy, timeToText(t.CreatedAt)) + out, err := scanWorkerJoinToken(row) + return out, mapErr(err) +} + +// GetWorkerJoinTokenByHash implements [store.WorkerCredentialStore]. +func (s *Store) GetWorkerJoinTokenByHash(ctx context.Context, hash string) (store.WorkerJoinToken, error) { + row := s.stmtGetWorkerJoinTokenByHash.QueryRowContext(ctx, hash) + out, err := scanWorkerJoinToken(row) + return out, mapErr(err) +} + +// MarkWorkerJoinTokenUsed implements [store.WorkerCredentialStore]. +func (s *Store) MarkWorkerJoinTokenUsed(ctx context.Context, id string, at time.Time) error { + res, err := s.stmtMarkWorkerJoinTokenUsed.ExecContext(ctx, timeToText(at), id) + if err != nil { + return mapErr(err) + } + return checkRowsAffected(res) +} + +// RedeemWorkerJoinToken implements [store.WorkerCredentialStore]. +// +// Follows the same idiom as CreateJobSubmission (job.go): the whole +// operation runs inside one transaction, using raw SQL via tx rather than +// the prepared statements the standalone per-row methods use, and a +// deferred Rollback that is a no-op after a successful Commit. +// +// The token claim reuses sqlConsumeWorkerJoinToken's text — a single +// UPDATE ... RETURNING that carries both the validity check and the claim in +// one statement, so an unknown token, an expired one, and one another +// request claimed a moment earlier are all the same "matched no row" +// outcome. If that claim fails, the transaction is rolled back (via the +// deferred Rollback) and [store.ErrNotFound] is returned with the token +// untouched. +// +// The credential insert happens inside the SAME transaction, so a conflict +// there — a worker ID already bound to an active credential, or a public +// key already enrolled anywhere — rolls back the token claim too: the token +// is NOT consumed, and remains redeemable by a later, non-conflicting +// request. +func (s *Store) RedeemWorkerJoinToken(ctx context.Context, hash string, now time.Time, cred store.WorkerCredential) (store.WorkerCredential, error) { + tx, err := s.db.BeginTx(ctx, nil) + if err != nil { + return store.WorkerCredential{}, mapErr(err) + } + defer func() { _ = tx.Rollback() }() //nolint:errcheck // rollback after commit is a no-op + + nowText := timeToText(now) + tokRow := tx.QueryRowContext(ctx, sqlConsumeWorkerJoinToken, nowText, hash, nowText) + if _, err := scanWorkerJoinToken(tokRow); err != nil { + return store.WorkerCredential{}, mapErr(err) + } + + credRow := tx.QueryRowContext(ctx, sqlInsertWorkerCredential, cred.ID, cred.WorkerID, cred.PublicKey, cred.Name, + timeToText(cred.EnrolledAt), nullTimeToText(cred.LastSeenAt), nullTimeToText(cred.RevokedAt)) + out, err := scanWorkerCredential(credRow) + if err != nil { + return store.WorkerCredential{}, mapErr(err) + } + + if err := tx.Commit(); err != nil { + return store.WorkerCredential{}, mapErr(err) + } + return out, nil +} diff --git a/internal/store/sqlite/workercredential_test.go b/internal/store/sqlite/workercredential_test.go new file mode 100644 index 00000000..a19b9a9c --- /dev/null +++ b/internal/store/sqlite/workercredential_test.go @@ -0,0 +1,611 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package sqlite_test + +import ( + "context" + "errors" + "fmt" + "sync" + "testing" + "time" + + "github.com/uberware/sqi/internal/store" +) + +func TestWorkerCredential_CreateAndGet(t *testing.T) { + s := openTestStore(t) + ctx := context.Background() + enrolledAt := time.Now().UTC().Truncate(time.Second) + + c := store.WorkerCredential{ + ID: "wc1", + WorkerID: "w1", + PublicKey: "pub1", + Name: "worker-one", + EnrolledAt: enrolledAt, + } + created, err := s.CreateWorkerCredential(ctx, c) + if err != nil { + t.Fatalf("CreateWorkerCredential: %v", err) + } + if created.ID != c.ID { + t.Errorf("ID: got %q, want %q", created.ID, c.ID) + } + if !created.EnrolledAt.Equal(enrolledAt) { + t.Errorf("EnrolledAt: got %v, want %v", created.EnrolledAt, enrolledAt) + } + if created.LastSeenAt != nil { + t.Errorf("LastSeenAt: got %v, want nil", created.LastSeenAt) + } + if created.RevokedAt != nil { + t.Errorf("RevokedAt: got %v, want nil", created.RevokedAt) + } + + got, err := s.GetActiveWorkerCredentialByWorkerID(ctx, "w1") + if err != nil { + t.Fatalf("GetActiveWorkerCredentialByWorkerID: %v", err) + } + if got.ID != c.ID || got.PublicKey != c.PublicKey || got.Name != c.Name { + t.Errorf("got %+v, want fields matching %+v", got, c) + } +} + +func TestWorkerCredential_GetNotFound(t *testing.T) { + s := openTestStore(t) + _, err := s.GetActiveWorkerCredentialByWorkerID(context.Background(), "nope") + if !errors.Is(err, store.ErrNotFound) { + t.Errorf("expected ErrNotFound, got %v", err) + } +} + +func TestWorkerCredential_DuplicateWorkerID(t *testing.T) { + s := openTestStore(t) + ctx := context.Background() + now := time.Now().UTC().Truncate(time.Second) + + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: now, + }); err != nil { + t.Fatalf("CreateWorkerCredential: %v", err) + } + _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc2", WorkerID: "w1", PublicKey: "pub2", EnrolledAt: now, + }) + if !errors.Is(err, store.ErrConflict) { + t.Errorf("expected ErrConflict for duplicate worker_id, got %v", err) + } +} + +func TestWorkerCredential_DuplicatePublicKey(t *testing.T) { + s := openTestStore(t) + ctx := context.Background() + now := time.Now().UTC().Truncate(time.Second) + + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: now, + }); err != nil { + t.Fatalf("CreateWorkerCredential: %v", err) + } + _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc2", WorkerID: "w2", PublicKey: "pub1", EnrolledAt: now, + }) + if !errors.Is(err, store.ErrConflict) { + t.Errorf("expected ErrConflict for duplicate public_key, got %v", err) + } +} + +// TestWorkerCredential_RevokedPublicKeyStillUnique verifies that public_key +// uniqueness is untouched by the worker_id fix: a key that belonged to a +// now-revoked credential can never be reused, by any worker. +func TestWorkerCredential_RevokedPublicKeyStillUnique(t *testing.T) { + s := openTestStore(t) + ctx := context.Background() + now := time.Now().UTC().Truncate(time.Second) + + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: now, + }); err != nil { + t.Fatalf("CreateWorkerCredential: %v", err) + } + if err := s.RevokeWorkerCredential(ctx, "w1", now); err != nil { + t.Fatalf("RevokeWorkerCredential: %v", err) + } + + _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc2", WorkerID: "w2", PublicKey: "pub1", EnrolledAt: now, + }) + if !errors.Is(err, store.ErrConflict) { + t.Errorf("expected ErrConflict reusing a revoked credential's public_key, got %v", err) + } +} + +func TestWorkerCredential_ListActiveOmitsRevoked(t *testing.T) { + s := openTestStore(t) + ctx := context.Background() + now := time.Now().UTC().Truncate(time.Second) + + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: now, + }); err != nil { + t.Fatalf("CreateWorkerCredential w1: %v", err) + } + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc2", WorkerID: "w2", PublicKey: "pub2", EnrolledAt: now, + }); err != nil { + t.Fatalf("CreateWorkerCredential w2: %v", err) + } + if err := s.RevokeWorkerCredential(ctx, "w2", now); err != nil { + t.Fatalf("RevokeWorkerCredential: %v", err) + } + + active, err := s.ListActiveWorkerCredentials(ctx) + if err != nil { + t.Fatalf("ListActiveWorkerCredentials: %v", err) + } + if len(active) != 1 { + t.Fatalf("want 1 active credential, got %d", len(active)) + } + if active[0].WorkerID != "w1" { + t.Errorf("active[0].WorkerID: got %q, want %q", active[0].WorkerID, "w1") + } +} + +// TestWorkerCredential_ListActiveAfterRotationHasExactlyOneRow verifies that +// after a worker is enrolled, revoked, and re-enrolled with a new key, +// exactly one row is active — never zero (the rotation silently failing) and +// never two (both keys accepted by the broker at once). +func TestWorkerCredential_ListActiveAfterRotationHasExactlyOneRow(t *testing.T) { + s := openTestStore(t) + ctx := context.Background() + now := time.Now().UTC().Truncate(time.Second) + + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: now, + }); err != nil { + t.Fatalf("CreateWorkerCredential (first): %v", err) + } + if err := s.RevokeWorkerCredential(ctx, "w1", now); err != nil { + t.Fatalf("RevokeWorkerCredential: %v", err) + } + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc2", WorkerID: "w1", PublicKey: "pub2", EnrolledAt: now, + }); err != nil { + t.Fatalf("CreateWorkerCredential (rotated): %v", err) + } + + active, err := s.ListActiveWorkerCredentials(ctx) + if err != nil { + t.Fatalf("ListActiveWorkerCredentials: %v", err) + } + var forW1 []store.WorkerCredential + for _, c := range active { + if c.WorkerID == "w1" { + forW1 = append(forW1, c) + } + } + if len(forW1) != 1 { + t.Fatalf("want exactly 1 active credential for w1 after rotation, got %d", len(forW1)) + } + if forW1[0].PublicKey != "pub2" { + t.Errorf("active credential PublicKey = %q, want %q (the rotated key)", forW1[0].PublicKey, "pub2") + } +} + +// TestWorkerCredential_RotateAfterRevoke verifies that revocation is not a +// one-way door: once a worker's credential is revoked, that same worker ID +// can enroll again with a new public key, and both rows persist — the +// revoked one keeping its revoked_at, the new one active. +func TestWorkerCredential_RotateAfterRevoke(t *testing.T) { + s := openTestStore(t) + ctx := context.Background() + now := time.Now().UTC().Truncate(time.Second) + + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: now, + }); err != nil { + t.Fatalf("CreateWorkerCredential (first): %v", err) + } + if err := s.RevokeWorkerCredential(ctx, "w1", now); err != nil { + t.Fatalf("RevokeWorkerCredential: %v", err) + } + + created, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc2", WorkerID: "w1", PublicKey: "pub2", EnrolledAt: now, + }) + if err != nil { + t.Fatalf("CreateWorkerCredential (rotated): %v", err) + } + if created.RevokedAt != nil { + t.Errorf("rotated credential RevokedAt = %v, want nil", created.RevokedAt) + } + + // GetActiveWorkerCredentialByWorkerID's whole contract is to resolve this + // ambiguity: after a rotation there are two rows for w1, and it must + // return the active one (pub2), never the revoked one (pub1). + got, err := s.GetActiveWorkerCredentialByWorkerID(ctx, "w1") + if err != nil { + t.Fatalf("GetActiveWorkerCredentialByWorkerID: %v", err) + } + if got.WorkerID != "w1" { + t.Errorf("GetActiveWorkerCredentialByWorkerID WorkerID = %q, want %q", got.WorkerID, "w1") + } + if got.PublicKey != "pub2" { + t.Errorf("GetActiveWorkerCredentialByWorkerID PublicKey = %q, want %q (the active, rotated key)", got.PublicKey, "pub2") + } + if got.RevokedAt != nil { + t.Errorf("GetActiveWorkerCredentialByWorkerID RevokedAt = %v, want nil", got.RevokedAt) + } +} + +// TestWorkerCredential_GetActiveOnly_RevokedOnlyReturnsNotFound verifies that +// a worker whose only credential has been revoked (no rotation yet) is +// reported as having no active credential — GetActiveWorkerCredentialByWorkerID +// must never hand back a revoked row. +func TestWorkerCredential_GetActiveOnly_RevokedOnlyReturnsNotFound(t *testing.T) { + s := openTestStore(t) + ctx := context.Background() + now := time.Now().UTC().Truncate(time.Second) + + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: now, + }); err != nil { + t.Fatalf("CreateWorkerCredential: %v", err) + } + if err := s.RevokeWorkerCredential(ctx, "w1", now); err != nil { + t.Fatalf("RevokeWorkerCredential: %v", err) + } + + _, err := s.GetActiveWorkerCredentialByWorkerID(ctx, "w1") + if !errors.Is(err, store.ErrNotFound) { + t.Errorf("expected ErrNotFound for a worker with only a revoked credential, got %v", err) + } +} + +func TestWorkerCredential_RevokeNotFound(t *testing.T) { + s := openTestStore(t) + err := s.RevokeWorkerCredential(context.Background(), "nope", time.Now().UTC()) + if !errors.Is(err, store.ErrNotFound) { + t.Errorf("expected ErrNotFound, got %v", err) + } +} + +func TestWorkerCredential_Touch(t *testing.T) { + s := openTestStore(t) + ctx := context.Background() + enrolledAt := time.Now().UTC().Truncate(time.Second) + + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: enrolledAt, + }); err != nil { + t.Fatalf("CreateWorkerCredential: %v", err) + } + + seenAt := enrolledAt.Add(time.Hour) + if err := s.TouchWorkerCredential(ctx, "w1", seenAt); err != nil { + t.Fatalf("TouchWorkerCredential: %v", err) + } + + got, err := s.GetActiveWorkerCredentialByWorkerID(ctx, "w1") + if err != nil { + t.Fatalf("GetActiveWorkerCredentialByWorkerID: %v", err) + } + if got.LastSeenAt == nil || !got.LastSeenAt.Equal(seenAt) { + t.Errorf("LastSeenAt = %v, want %v", got.LastSeenAt, seenAt) + } +} + +func TestWorkerCredential_TouchNotFound(t *testing.T) { + s := openTestStore(t) + err := s.TouchWorkerCredential(context.Background(), "nope", time.Now().UTC()) + if !errors.Is(err, store.ErrNotFound) { + t.Errorf("expected ErrNotFound, got %v", err) + } +} + +func TestWorkerCredential_TouchOnlyMatchesActiveRow(t *testing.T) { + s := openTestStore(t) + ctx := context.Background() + now := time.Now().UTC().Truncate(time.Second) + + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: now, + }); err != nil { + t.Fatalf("CreateWorkerCredential: %v", err) + } + if err := s.RevokeWorkerCredential(ctx, "w1", now); err != nil { + t.Fatalf("RevokeWorkerCredential: %v", err) + } + + // Only a revoked row exists for w1: touching it must not resurrect it as + // a match, so it reports ErrNotFound exactly like the revoke and get + // paths do. + err := s.TouchWorkerCredential(ctx, "w1", now.Add(time.Hour)) + if !errors.Is(err, store.ErrNotFound) { + t.Errorf("expected ErrNotFound for a worker with only a revoked credential, got %v", err) + } +} + +func TestWorkerJoinToken_CreateAndGetByHash(t *testing.T) { + s := openTestStore(t) + ctx := context.Background() + expiresAt := time.Now().UTC().Add(24 * time.Hour).Truncate(time.Second) + createdAt := time.Now().UTC().Truncate(time.Second) + + tok := store.WorkerJoinToken{ + ID: "jt1", + TokenHash: "hash1", + Prefix: "sqiwjt_abcd", + Name: "farm-join", + ExpiresAt: expiresAt, + CreatedBy: "u1", + CreatedAt: createdAt, + } + created, err := s.CreateWorkerJoinToken(ctx, tok) + if err != nil { + t.Fatalf("CreateWorkerJoinToken: %v", err) + } + if created.UsedAt != nil { + t.Errorf("UsedAt: got %v, want nil", created.UsedAt) + } + + got, err := s.GetWorkerJoinTokenByHash(ctx, "hash1") + if err != nil { + t.Fatalf("GetWorkerJoinTokenByHash: %v", err) + } + if got.ID != tok.ID || got.Prefix != tok.Prefix || got.Name != tok.Name || got.CreatedBy != tok.CreatedBy { + t.Errorf("got %+v, want fields matching %+v", got, tok) + } + if !got.ExpiresAt.Equal(expiresAt) { + t.Errorf("ExpiresAt: got %v, want %v", got.ExpiresAt, expiresAt) + } +} + +func TestWorkerJoinToken_MarkUsed(t *testing.T) { + s := openTestStore(t) + ctx := context.Background() + now := time.Now().UTC().Truncate(time.Second) + + if _, err := s.CreateWorkerJoinToken(ctx, store.WorkerJoinToken{ + ID: "jt1", + TokenHash: "hash1", + Prefix: "sqiwjt_abcd", + ExpiresAt: now.Add(time.Hour), + CreatedAt: now, + }); err != nil { + t.Fatalf("CreateWorkerJoinToken: %v", err) + } + + usedAt := now.Add(time.Minute) + if err := s.MarkWorkerJoinTokenUsed(ctx, "jt1", usedAt); err != nil { + t.Fatalf("MarkWorkerJoinTokenUsed: %v", err) + } + + got, err := s.GetWorkerJoinTokenByHash(ctx, "hash1") + if err != nil { + t.Fatalf("GetWorkerJoinTokenByHash: %v", err) + } + if got.UsedAt == nil { + t.Fatal("UsedAt: got nil, want set") + } + if !got.UsedAt.Equal(usedAt) { + t.Errorf("UsedAt: got %v, want %v", *got.UsedAt, usedAt) + } +} + +func TestWorkerJoinToken_MarkUsedNotFound(t *testing.T) { + s := openTestStore(t) + err := s.MarkWorkerJoinTokenUsed(context.Background(), "nope", time.Now().UTC()) + if !errors.Is(err, store.ErrNotFound) { + t.Errorf("expected ErrNotFound, got %v", err) + } +} + +func TestWorkerJoinToken_GetByHashNotFound(t *testing.T) { + s := openTestStore(t) + _, err := s.GetWorkerJoinTokenByHash(context.Background(), "nope") + if !errors.Is(err, store.ErrNotFound) { + t.Errorf("expected ErrNotFound, got %v", err) + } +} + +// ── RedeemWorkerJoinToken ──────────────────────────────────────────────────── +// +// The atomic single-use claim, combined with the credential it authorizes in +// one transaction. Its semantics are mirrored by +// fake.Store.RedeemWorkerJoinToken; the same cases are asserted there. + +func TestWorkerJoinToken_Redeem(t *testing.T) { + s := openTestStore(t) + ctx := context.Background() + now := time.Now().UTC().Truncate(time.Second) + + if _, err := s.CreateWorkerJoinToken(ctx, store.WorkerJoinToken{ + ID: "jt1", TokenHash: "hash1", Prefix: "sqiwjt_abcd", + ExpiresAt: now.Add(time.Hour), CreatedAt: now, + }); err != nil { + t.Fatalf("CreateWorkerJoinToken: %v", err) + } + + claimedAt := now.Add(time.Minute) + got, err := s.RedeemWorkerJoinToken(ctx, "hash1", claimedAt, store.WorkerCredential{ + ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: claimedAt, + }) + if err != nil { + t.Fatalf("RedeemWorkerJoinToken: %v", err) + } + if got.ID != "wc1" || got.WorkerID != "w1" { + t.Errorf("got %+v, want the created credential for w1", got) + } + + storedTok, err := s.GetWorkerJoinTokenByHash(ctx, "hash1") + if err != nil { + t.Fatalf("GetWorkerJoinTokenByHash: %v", err) + } + if storedTok.UsedAt == nil { + t.Error("the claim was not persisted") + } else if !storedTok.UsedAt.Equal(claimedAt) { + t.Errorf("UsedAt: got %v, want %v", *storedTok.UsedAt, claimedAt) + } + + storedCred, err := s.GetActiveWorkerCredentialByWorkerID(ctx, "w1") + if err != nil { + t.Fatalf("GetActiveWorkerCredentialByWorkerID: %v", err) + } + if storedCred.PublicKey != "pub1" { + t.Errorf("PublicKey: got %q, want %q", storedCred.PublicKey, "pub1") + } + + // A second redemption of the same token matches nothing. + if _, err := s.RedeemWorkerJoinToken(ctx, "hash1", claimedAt.Add(time.Second), store.WorkerCredential{ + ID: "wc2", WorkerID: "w2", PublicKey: "pub2", EnrolledAt: claimedAt, + }); !errors.Is(err, store.ErrNotFound) { + t.Errorf("second RedeemWorkerJoinToken: got %v, want store.ErrNotFound", err) + } +} + +func TestWorkerJoinToken_RedeemExpired(t *testing.T) { + s := openTestStore(t) + ctx := context.Background() + now := time.Now().UTC().Truncate(time.Second) + + if _, err := s.CreateWorkerJoinToken(ctx, store.WorkerJoinToken{ + ID: "jt1", TokenHash: "hash1", Prefix: "sqiwjt_abcd", + ExpiresAt: now, CreatedAt: now.Add(-time.Hour), + }); err != nil { + t.Fatalf("CreateWorkerJoinToken: %v", err) + } + cred := store.WorkerCredential{ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: now} + + // Exactly at expiry counts as expired: the predicate is expires_at > now. + if _, err := s.RedeemWorkerJoinToken(ctx, "hash1", now, cred); !errors.Is(err, store.ErrNotFound) { + t.Errorf("RedeemWorkerJoinToken at the expiry instant: got %v, want store.ErrNotFound", err) + } + if _, err := s.RedeemWorkerJoinToken(ctx, "hash1", now.Add(time.Hour), cred); !errors.Is(err, store.ErrNotFound) { + t.Errorf("RedeemWorkerJoinToken after expiry: got %v, want store.ErrNotFound", err) + } + + // An expired token must not be marked used by the failed attempts, and + // no credential must have been created. + stored, err := s.GetWorkerJoinTokenByHash(ctx, "hash1") + if err != nil { + t.Fatalf("GetWorkerJoinTokenByHash: %v", err) + } + if stored.UsedAt != nil { + t.Error("a refused claim marked the token used") + } + if _, err := s.GetActiveWorkerCredentialByWorkerID(ctx, "w1"); !errors.Is(err, store.ErrNotFound) { + t.Errorf("expected no credential to have been created, got %v", err) + } +} + +func TestWorkerJoinToken_RedeemUnknown(t *testing.T) { + s := openTestStore(t) + cred := store.WorkerCredential{ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: time.Now().UTC()} + if _, err := s.RedeemWorkerJoinToken(context.Background(), "nope", time.Now().UTC(), cred); !errors.Is(err, store.ErrNotFound) { + t.Errorf("expected ErrNotFound, got %v", err) + } +} + +// TestWorkerJoinToken_RedeemConflictRollsBackClaim is the transactional +// guarantee this method exists for: a credential that cannot be created — +// here, a worker ID already bound to an active credential — must roll back +// the token claim too, so the token remains redeemable rather than being +// spent on a request that was always going to be rejected. +func TestWorkerJoinToken_RedeemConflictRollsBackClaim(t *testing.T) { + s := openTestStore(t) + ctx := context.Background() + now := time.Now().UTC().Truncate(time.Second) + + if _, err := s.CreateWorkerJoinToken(ctx, store.WorkerJoinToken{ + ID: "jt1", TokenHash: "hash1", Prefix: "sqiwjt_abcd", + ExpiresAt: now.Add(time.Hour), CreatedAt: now, + }); err != nil { + t.Fatalf("CreateWorkerJoinToken: %v", err) + } + if _, err := s.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: "wc-existing", WorkerID: "w1", PublicKey: "pub-existing", EnrolledAt: now, + }); err != nil { + t.Fatalf("seed CreateWorkerCredential: %v", err) + } + + _, err := s.RedeemWorkerJoinToken(ctx, "hash1", now.Add(time.Minute), store.WorkerCredential{ + ID: "wc1", WorkerID: "w1", PublicKey: "pub1", EnrolledAt: now, + }) + if !errors.Is(err, store.ErrConflict) { + t.Fatalf("RedeemWorkerJoinToken: got %v, want store.ErrConflict", err) + } + + storedTok, err := s.GetWorkerJoinTokenByHash(ctx, "hash1") + if err != nil { + t.Fatalf("GetWorkerJoinTokenByHash: %v", err) + } + if storedTok.UsedAt != nil { + t.Error("a conflicting redemption spent the token; the whole transaction must roll back") + } + + // The token is still redeemable, for a worker ID that does not conflict. + got, err := s.RedeemWorkerJoinToken(ctx, "hash1", now.Add(time.Minute), store.WorkerCredential{ + ID: "wc2", WorkerID: "w3", PublicKey: "pub3", EnrolledAt: now, + }) + if err != nil { + t.Fatalf("re-redemption after conflict: %v", err) + } + if got.WorkerID != "w3" { + t.Errorf("WorkerID = %q, want w3", got.WorkerID) + } +} + +// TestWorkerJoinToken_RedeemIsAtomic drives concurrent redemptions of one +// token through the real database, each for a distinct (non-conflicting) +// worker ID and public key — so the only thing that can make an attempt +// lose is the token claim itself. Exactly one may win, whatever the +// interleaving: the UPDATE's own WHERE clause is the check, so there is no +// window between checking used_at and setting it. +func TestWorkerJoinToken_RedeemIsAtomic(t *testing.T) { + s := openTestStore(t) + ctx := context.Background() + now := time.Now().UTC().Truncate(time.Second) + + if _, err := s.CreateWorkerJoinToken(ctx, store.WorkerJoinToken{ + ID: "jt1", TokenHash: "hash1", Prefix: "sqiwjt_abcd", + ExpiresAt: now.Add(time.Hour), CreatedAt: now, + }); err != nil { + t.Fatalf("CreateWorkerJoinToken: %v", err) + } + + const attempts = 8 + errs := make([]error, attempts) + var ready, done sync.WaitGroup + ready.Add(attempts) + done.Add(attempts) + start := make(chan struct{}) + for i := range attempts { + go func() { + defer done.Done() + ready.Done() + <-start + _, errs[i] = s.RedeemWorkerJoinToken(ctx, "hash1", now.Add(time.Minute), store.WorkerCredential{ + ID: fmt.Sprintf("wc%d", i), WorkerID: fmt.Sprintf("w%d", i), + PublicKey: fmt.Sprintf("pub%d", i), EnrolledAt: now, + }) + }() + } + ready.Wait() + close(start) + done.Wait() + + won := 0 + for i, err := range errs { + switch { + case err == nil: + won++ + case errors.Is(err, store.ErrNotFound): + default: + t.Errorf("attempt %d: unexpected error %v", i, err) + } + } + if won != 1 { + t.Errorf("%d of %d concurrent redemptions of one single-use token succeeded, want exactly 1", won, attempts) + } +} diff --git a/internal/store/store.go b/internal/store/store.go index 7624e087..f9308174 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -61,6 +61,7 @@ type Store interface { UserStore SessionStore APIKeyStore + WorkerCredentialStore io.Closer } diff --git a/internal/store/task_log.go b/internal/store/task_log.go index 7afb5ec5..3bb98e35 100644 --- a/internal/store/task_log.go +++ b/internal/store/task_log.go @@ -19,7 +19,7 @@ const ( // TaskLog is one chunk of output produced by a running task. // -// Workers publish log chunks to task.logs. as output is produced. +// Workers publish log chunks to task.logs.. as output is produced. // The server persists each chunk here with: // - A monotonic worker-side sequence number (SeqNum) per attempt, starting // at 1 and incrementing by 1. This allows the REST log-tail endpoint to diff --git a/internal/store/workercredential.go b/internal/store/workercredential.go new file mode 100644 index 00000000..5f75246f --- /dev/null +++ b/internal/store/workercredential.go @@ -0,0 +1,99 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package store + +import ( + "context" + "time" +) + +// WorkerCredential binds a worker's self-chosen ID to an Ed25519 nkey public +// key, issued during enrollment and presented on every broker connection +// thereafter. WorkerID is deliberately not a foreign key to a worker row: a +// credential is issued before the worker has ever registered, and on +// auth-off farms worker rows exist with no credential at all. +type WorkerCredential struct { + ID string + WorkerID string + PublicKey string + Name string + EnrolledAt time.Time + LastSeenAt *time.Time // nil = never seen since enrollment + RevokedAt *time.Time // nil = active +} + +// WorkerJoinToken is a single-use, time-limited token an operator issues so a +// worker can enroll and receive a [WorkerCredential]. Only TokenHash is +// stored; the raw token is shown to the operator exactly once, at creation. +type WorkerJoinToken struct { + ID string + TokenHash string // SHA-256 of the raw token; Go-side only + Prefix string // leading chars of the raw token, for list identification + Name string + ExpiresAt time.Time + UsedAt *time.Time // nil = not yet redeemed + CreatedBy string + CreatedAt time.Time +} + +// WorkerCredentialStore is the persistence interface for [WorkerCredential] +// and [WorkerJoinToken] records. +type WorkerCredentialStore interface { + // CreateWorkerCredential inserts a new credential. Returns [ErrConflict] + // on a worker-id, public-key, or id collision. + CreateWorkerCredential(ctx context.Context, c WorkerCredential) (WorkerCredential, error) + // GetActiveWorkerCredentialByWorkerID returns the ACTIVE credential for + // workerID — the one with a nil RevokedAt — or [ErrNotFound] if the + // worker has none. A revoked credential is never retrievable through + // this method, even if it is the only credential the worker has ever + // had: after a key rotation (revoke, then re-enroll with a new key) a + // worker can have both a revoked row and an active one, and this method + // exists specifically to resolve that ambiguity to the one row that + // matters for authentication. + // + // It has NO production callers — it is test-only API surface used to + // assert credential state directly, rather than through a code path that + // exercises it incidentally. + GetActiveWorkerCredentialByWorkerID(ctx context.Context, workerID string) (WorkerCredential, error) + // ListActiveWorkerCredentials returns every credential with a nil + // RevokedAt. It is what the broker's authorized-key set is rebuilt from, + // so a revoked credential must never appear in the result. + ListActiveWorkerCredentials(ctx context.Context) ([]WorkerCredential, error) + // RevokeWorkerCredential soft-revokes the credential for workerID. + // Returns [ErrNotFound] if no such worker has a credential. + RevokeWorkerCredential(ctx context.Context, workerID string, at time.Time) error + // TouchWorkerCredential sets LastSeenAt on workerID's ACTIVE credential. + // Returns [ErrNotFound] if the worker has no active credential. Callers + // that treat "seen" as best-effort bookkeeping (registration, not a + // security decision) should not fail on that error. + TouchWorkerCredential(ctx context.Context, workerID string, at time.Time) error + // CreateWorkerJoinToken inserts a new join token. Returns [ErrConflict] + // on a token-hash or id collision. + CreateWorkerJoinToken(ctx context.Context, t WorkerJoinToken) (WorkerJoinToken, error) + // GetWorkerJoinTokenByHash returns the token for hash, or [ErrNotFound] if + // no such token exists. + GetWorkerJoinTokenByHash(ctx context.Context, hash string) (WorkerJoinToken, error) + // MarkWorkerJoinTokenUsed sets UsedAt. Returns [ErrNotFound] if id is + // unknown. + MarkWorkerJoinTokenUsed(ctx context.Context, id string, at time.Time) error + // RedeemWorkerJoinToken atomically claims the single-use join token + // identified by hash AND creates cred, in one transaction. It succeeds + // only for a token that exists, has not been redeemed, and has not + // expired at now; every other case — including a token another request + // claimed a moment earlier — returns [ErrNotFound], and cred is not + // created. Returns [ErrConflict] if cred cannot be created (a worker ID + // already bound to an active credential, or a public key already + // enrolled to any worker); in that case the whole transaction rolls + // back, so the token is NOT consumed and remains redeemable. + // + // The token claim and the credential creation must be ONE transaction. + // Committing the claim before the credential is known to be creatable + // would burn a single-use token on a request that was always going to + // be rejected as a conflict — the caller gets nothing for it, and the + // operator has to issue a new one. Within the claim itself, check and + // set must still be a single statement: reading the token, inspecting + // UsedAt and marking it used separately is a check-then-act race, where + // two concurrent enrollments presenting the same single-use token both + // observe UsedAt as nil and both succeed. + RedeemWorkerJoinToken(ctx context.Context, hash string, now time.Time, cred WorkerCredential) (WorkerCredential, error) +} diff --git a/internal/worker/config/config.go b/internal/worker/config/config.go index ac5ef0ea..2492011a 100644 --- a/internal/worker/config/config.go +++ b/internal/worker/config/config.go @@ -24,6 +24,7 @@ import ( "gopkg.in/yaml.v3" + "github.com/uberware/sqi/internal/brokerauth" "github.com/uberware/sqi/internal/worker/capabilities" ) @@ -297,6 +298,34 @@ type NATSConfig struct { // Actual wait uses exponential backoff with jitter. // Env: SQI_WORKER_NATS_RECONNECT_WAIT ReconnectWait time.Duration `yaml:"reconnect_wait"` + + // CredentialFile is the path to this worker's nkey seed file. When empty + // it defaults to /worker.nk. The file is created by + // enrollment or by `sqi-worker keygen` and must be mode 0600. + // Env: SQI_WORKER_NATS_CREDENTIAL_FILE + CredentialFile string `yaml:"credential_file"` + + // JoinToken is a worker enrollment token. Used exactly once, on first + // start, to obtain a credential; ignored once CredentialFile exists. + // Prefer JoinTokenFile — a token in a config file is a secret at rest. + // Env: SQI_WORKER_NATS_JOIN_TOKEN + JoinToken string `yaml:"join_token"` + + // JoinTokenFile is a path to a file containing a join token. Takes + // precedence over JoinToken. + // Env: SQI_WORKER_NATS_JOIN_TOKEN_FILE + JoinTokenFile string `yaml:"join_token_file"` + + // ServerURL is the sqi-server HTTP base URL used for enrollment, e.g. + // "http://sqi-server.example:8080". Enrollment runs over REST, not over + // NATS: the broker's job is to reject unauthenticated connections, so it + // cannot also be the channel a worker gets its first credential over. + // Required whenever a join token is configured — it is NOT derived from + // mDNS discovery; a worker relying on mDNS with no explicit server_url + // fails enrollment with an actionable error naming this field rather + // than attempting a request with no host. + // Env: SQI_WORKER_NATS_SERVER_URL + ServerURL string `yaml:"server_url"` } // WorkerSettings controls the worker's identity and runtime behavior. @@ -381,6 +410,17 @@ type WorkerSettings struct { // list means the worker accepts assignments from all queues via a wildcard // JetStream consumer. Set this when running a heterogeneous farm where // some workers specialise in a subset of queues. + // + // Each entry becomes a NATS subject token in the work-lease subject + // work.lease.., so it must satisfy the same constraint + // as a worker ID (see brokerauth.ValidWorkerIDToken): non-empty, and free + // of '.', '*', '>' and whitespace. A queue ID that does not is rejected + // by [Validate] rather than silently producing a subject the broker will + // never route — with broker auth on, a queue ID containing a dot yields + // too many subject tokens for this worker's publish grant and every lease + // request for it is denied outright; with broker auth off it fails + // bus.ParseWorkerSubject on the server side and gets no reply either way, + // so the worker retries forever with nothing in its logs naming the cause. // Env: SQI_WORKER_QUEUE_IDS (comma-separated) QueueIDs []string `yaml:"queue_ids"` @@ -513,25 +553,58 @@ func defaultDataDir() string { return filepath.Join(home, ".sqi", "worker") } +// DefaultCredentialFile returns the default nkey seed path under dataDir, +// used whenever NATS.CredentialFile is left unset in [Load]'s resolution +// below. Exported so other worker-side entry points that write or look for +// the seed file — the "sqi-worker keygen" CLI, notably — derive the same +// path from the data directory rather than each hardcoding it. +func DefaultCredentialFile(dataDir string) string { + return filepath.Join(dataDir, "worker.nk") +} + // Load returns the effective WorkerConfig by merging layers in override order: // built-in defaults → YAML/JSON file → SQI_WORKER_* env vars → CLI flags. // // configFile may be empty, in which case Load searches the default paths. func Load(configFile string, flags FlagOverrides) (WorkerConfig, error) { + cfg, _, err := LoadWithSources(configFile, flags) + return cfg, err +} + +// Sources reports, for a subset of settings, whether the config file or +// environment layer explicitly decided the value before [Load]'s own +// default-fill step ran — as opposed to it being left empty for that step to +// fill in. Comparing the resolved value against the computed default (e.g. +// [DefaultCredentialFile]) is not sound: an operator can explicitly +// configure a value that happens to equal what the default-fill would have +// produced anyway, in which case value comparison cannot tell "explicitly +// configured" apart from "left unset". Sources answers from the file/env +// layers themselves, so it is not fooled by that. +type Sources struct { + // CredentialFile is true when nats.credential_file was set by the + // config file or by SQI_WORKER_NATS_CREDENTIAL_FILE, before [Load] + // resolves an empty value to DefaultCredentialFile(Worker.DataDir). + CredentialFile bool +} + +// LoadWithSources is [Load], additionally reporting which of a subset of +// settings (see [Sources]) were explicitly decided by the config file or +// environment layer. +func LoadWithSources(configFile string, flags FlagOverrides) (WorkerConfig, Sources, error) { cfg := Default() // ── Config file layer ───────────────────────────────────────────────── path, err := resolveConfigFile(configFile) if err != nil { - return WorkerConfig{}, fmt.Errorf("resolve config file: %w", err) + return WorkerConfig{}, Sources{}, fmt.Errorf("resolve config file: %w", err) } if path != "" { data, err := os.ReadFile(path) if err != nil { - return WorkerConfig{}, fmt.Errorf("read config file %s: %w", path, err) + return WorkerConfig{}, Sources{}, fmt.Errorf("read config file %s: %w", path, err) } if err := unmarshalConfig(data, &cfg); err != nil { - return WorkerConfig{}, fmt.Errorf("parse config file %s: %w", path, err) + return WorkerConfig{}, Sources{}, fmt.Errorf("parse config file %s: %w", path, err) } } @@ -549,7 +622,17 @@ func Load(configFile string, flags FlagOverrides) (WorkerConfig, error) { cfg.NATS.InsecureSkipVerify = true } - return cfg, nil + var src Sources + src.CredentialFile = cfg.NATS.CredentialFile != "" + + // CredentialFile defaults relative to Worker.DataDir, which any of the + // three layers above may have changed, so it is resolved last rather + // than at struct-literal time in Default. + if cfg.NATS.CredentialFile == "" { + cfg.NATS.CredentialFile = DefaultCredentialFile(cfg.Worker.DataDir) + } + + return cfg, src, nil } // resolveConfigFile returns the config file path to use. If explicit is @@ -700,6 +783,18 @@ func applyNATSEnv(c *NATSConfig) { c.ReconnectWait = d } } + if v := os.Getenv("SQI_WORKER_NATS_CREDENTIAL_FILE"); v != "" { + c.CredentialFile = v + } + if v := os.Getenv("SQI_WORKER_NATS_JOIN_TOKEN"); v != "" { + c.JoinToken = v + } + if v := os.Getenv("SQI_WORKER_NATS_JOIN_TOKEN_FILE"); v != "" { + c.JoinTokenFile = v + } + if v := os.Getenv("SQI_WORKER_NATS_SERVER_URL"); v != "" { + c.ServerURL = v + } } func applyWorkerEnv(c *WorkerSettings) { @@ -744,7 +839,14 @@ func applyWorkerEnv(c *WorkerSettings) { // applyWorkerEnv to keep each function under the cyclomatic-complexity limit. func applyWorkerPullEnv(c *WorkerSettings) { if v := os.Getenv("SQI_WORKER_QUEUE_IDS"); v != "" { - c.QueueIDs = splitTags(v) + // Deliberately NOT splitTags: splitTags silently drops empty entries + // (right for CapabilityTags/EnvPassthrough, which have no invalid + // shape), but a queue ID has one — see [Validate]'s validateQueueIDs. + // Splitting without filtering here means a blank entry from the env + // var is caught and reported the same way a blank entry in a YAML + // queue_ids list is, instead of being silently dropped on one path + // and rejected on the other. + c.QueueIDs = splitCSV(v) } if v := os.Getenv("SQI_WORKER_PULL_IDLE_BACKOFF"); v != "" { if d, err := time.ParseDuration(v); err == nil { @@ -805,7 +907,9 @@ func applyLogStreamerEnv(c *LogStreamerConfig) { } } -// splitTags splits a comma-separated tag list, trimming whitespace. +// splitTags splits a comma-separated tag list, trimming whitespace and +// dropping empty entries. Used for fields with no invalid shape of their +// own, where a stray comma should just be ignored. func splitTags(s string) []string { parts := strings.Split(s, ",") out := make([]string, 0, len(parts)) @@ -817,6 +921,20 @@ func splitTags(s string) []string { return out } +// splitCSV splits a comma-separated list, trimming whitespace around each +// entry but KEEPING empty entries — unlike splitTags. Used for fields whose +// entries have their own validity check (queue IDs), so that a blank entry +// surfaces as a validation error naming its position instead of being +// silently discarded. +func splitCSV(s string) []string { + parts := strings.Split(s, ",") + out := make([]string, len(parts)) + for i, p := range parts { + out[i] = strings.TrimSpace(p) + } + return out +} + // ValidationError describes a single configuration problem. type ValidationError struct { Field string @@ -898,6 +1016,7 @@ func Validate(cfg WorkerConfig) []ValidationError { errs = append(errs, validateIsolation(cfg)...) errs = append(errs, validateExpr(cfg.Expr)...) + errs = append(errs, validateQueueIDs(cfg.Worker.QueueIDs)...) return errs } @@ -996,6 +1115,34 @@ func validateIsolation(cfg WorkerConfig) []ValidationError { return errs } +// validateQueueIDs rejects any worker.queue_ids entry that cannot be used as +// a single NATS subject token, at load time rather than at first lease +// request. A queue ID becomes a token in the work-lease subject +// work.lease.. (internal/bus/subjects.go); an entry +// containing '.' splits it into extra tokens the server-side parser and this +// worker's own broker-auth publish grant both reject, and the failure mode +// is silent and total: the worker's lease requests are refused or simply go +// unanswered, and nothing in its logs says why. +// +// brokerauth.ValidWorkerIDToken implements exactly this predicate for worker +// IDs, which share the same constraint for the same reason; it is reused +// here rather than duplicated. +func validateQueueIDs(queueIDs []string) []ValidationError { + var errs []ValidationError + for i, q := range queueIDs { + if !brokerauth.ValidWorkerIDToken(q) { + errs = append(errs, ValidationError{ + Field: fmt.Sprintf("worker.queue_ids[%d]", i), + Message: fmt.Sprintf( + "%q is not a valid NATS subject token: it must be non-empty and must not contain '.', whitespace, '*' or '>'", + q, + ), + }) + } + } + return errs +} + // validateLogStreamer validates the LogStreamerConfig fields. // Zero values are self-corrected by logstreamer.Config.applyDefaults, so only // explicitly negative values (almost certainly typos) are rejected here. diff --git a/internal/worker/config/config_test.go b/internal/worker/config/config_test.go index e9cec273..2f9cfef5 100644 --- a/internal/worker/config/config_test.go +++ b/internal/worker/config/config_test.go @@ -442,6 +442,103 @@ func TestValidate_ValidEnvPassthroughGlobsAreFine(t *testing.T) { } } +// ── worker.queue_ids validation ────────────────────────────────────────────── + +func TestValidate_QueueIDRejectsEachInvalidShape(t *testing.T) { + tests := []struct { + name string + id string + }{ + {"empty", ""}, + {"dot", "queue.one"}, + {"star", "queue*"}, + {"gt", "queue>"}, + {"whitespace", "queue one"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg := Default() + cfg.NATS.URL = "nats://localhost:4222" + cfg.Worker.QueueIDs = []string{tt.id} + + errs := Validate(cfg) + if !containsField(errs, "worker.queue_ids[0]") { + t.Fatalf("expected worker.queue_ids[0] error for %q, got %v", tt.id, errs) + } + for _, e := range errs { + if e.Field != "worker.queue_ids[0]" { + continue + } + if tt.id == "" { + // tt.id is "" here, so strings.Contains(e.Message, tt.id) + // would hold against any message and verify nothing; the + // message must instead say the entry is empty. + if !strings.Contains(e.Message, "non-empty") { + t.Errorf("error message %q does not say the entry must be non-empty", e.Message) + } + continue + } + if !strings.Contains(e.Message, tt.id) { + t.Errorf("error message %q does not name the offending value %q", e.Message, tt.id) + } + } + }) + } +} + +func TestValidate_QueueIDAcceptsUUID(t *testing.T) { + cfg := Default() + cfg.NATS.URL = "nats://localhost:4222" + cfg.Worker.QueueIDs = []string{"3f2a9c9e-6b1a-4e2f-9c3d-8f1a2b3c4d5e"} + + errs := Validate(cfg) + if containsField(errs, "worker.queue_ids[0]") { + t.Errorf("expected no queue_ids error for a valid UUID, got %v", errs) + } +} + +func TestLoad_QueueIDsYAMLAndEnvAgreeOnEmptyEntries(t *testing.T) { + // A blank entry from a YAML list and a blank entry from a comma-separated + // env var must produce the same shape: preserved (not silently dropped) + // so Validate rejects it, naming its position, on both paths. + t.Run("yaml", func(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "sqi-worker.yaml") + yamlContent := "worker:\n queue_ids: [\"\", \"q1\"]\n" + if err := os.WriteFile(path, []byte(yamlContent), 0o600); err != nil { + t.Fatalf("write config: %v", err) + } + t.Setenv("SQI_WORKER_NATS_URL", "nats://x:4222") + + cfg, err := Load(path, FlagOverrides{}) + if err != nil { + t.Fatalf("Load: %v", err) + } + if len(cfg.Worker.QueueIDs) != 2 || cfg.Worker.QueueIDs[0] != "" || cfg.Worker.QueueIDs[1] != "q1" { + t.Fatalf("queue_ids = %v, want [\"\", \"q1\"] (unfiltered, so Validate can reject the empty entry)", cfg.Worker.QueueIDs) + } + if !containsField(Validate(cfg), "worker.queue_ids[0]") { + t.Errorf("expected worker.queue_ids[0] validation error for the blank YAML entry") + } + }) + + t.Run("env", func(t *testing.T) { + t.Setenv("SQI_WORKER_NATS_URL", "nats://x:4222") + t.Setenv("SQI_WORKER_QUEUE_IDS", ",q1") + + cfg, err := Load("", FlagOverrides{}) + if err != nil { + t.Fatalf("Load: %v", err) + } + if len(cfg.Worker.QueueIDs) != 2 || cfg.Worker.QueueIDs[0] != "" || cfg.Worker.QueueIDs[1] != "q1" { + t.Fatalf("queue_ids = %v, want [\"\", \"q1\"] (unfiltered, so Validate can reject the empty entry)", cfg.Worker.QueueIDs) + } + if !containsField(Validate(cfg), "worker.queue_ids[0]") { + t.Errorf("expected worker.queue_ids[0] validation error for the blank env entry") + } + }) +} + // ── Diagnostics: defaults and env overrides ────────────────────────────────── func TestDefault_DiagnosticsEnabledByDefault(t *testing.T) { @@ -615,3 +712,37 @@ func TestValidate_RejectsBadDetector(t *testing.T) { t.Errorf("expected validation error for detector with no checks") } } + +func TestDefaultCredentialFile(t *testing.T) { + got := DefaultCredentialFile("/tmp/sqi-worker-data") + want := filepath.Join("/tmp/sqi-worker-data", "worker.nk") + if got != want { + t.Errorf("DefaultCredentialFile = %q, want %q", got, want) + } +} + +func TestLoad_NATSCredentialFileDefaultsUnderDataDir(t *testing.T) { + t.Setenv("SQI_WORKER_NATS_URL", "nats://x:4222") // satisfy validation + t.Setenv("SQI_WORKER_DATA_DIR", "/tmp/sqi-worker-data") + + cfg, err := Load("", FlagOverrides{}) + if err != nil { + t.Fatalf("Load: %v", err) + } + want := filepath.Join("/tmp/sqi-worker-data", "worker.nk") + if cfg.NATS.CredentialFile != want { + t.Errorf("NATS.CredentialFile = %q, want %q", cfg.NATS.CredentialFile, want) + } +} + +func TestLoad_NATSCredentialFileExplicitValuePreserved(t *testing.T) { + body := "nats:\n credential_file: /etc/sqi/worker.nk\n" + f := writeTempFile(t, "worker.yaml", []byte(body)) + cfg, err := Load(f, FlagOverrides{}) + if err != nil { + t.Fatalf("Load: %v", err) + } + if cfg.NATS.CredentialFile != "/etc/sqi/worker.nk" { + t.Errorf("NATS.CredentialFile = %q, want /etc/sqi/worker.nk", cfg.NATS.CredentialFile) + } +} diff --git a/internal/worker/config/workerid.go b/internal/worker/config/workerid.go index 540914eb..e105f38a 100644 --- a/internal/worker/config/workerid.go +++ b/internal/worker/config/workerid.go @@ -15,6 +15,15 @@ import ( // workerIDFilename is the name of the file that persists the worker's UUID. const workerIDFilename = "worker.id" +// WorkerIDFilePath returns the path [LoadOrCreateWorkerID] reads and writes +// under dataDir. Exported so other worker-side entry points that need to +// check for an existing worker ID without creating one — the "sqi-worker +// keygen" CLI, notably — derive the same path from the data directory +// rather than each hardcoding the filename. +func WorkerIDFilePath(dataDir string) string { + return filepath.Join(dataDir, workerIDFilename) +} + // LoadOrCreateWorkerID returns the worker's persistent UUID. // // On the first call for a given dataDir the function: diff --git a/internal/worker/enroll/enroll.go b/internal/worker/enroll/enroll.go new file mode 100644 index 00000000..132de742 --- /dev/null +++ b/internal/worker/enroll/enroll.go @@ -0,0 +1,199 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +// Package enroll obtains the nkey broker credential sqi-worker authenticates +// with: loading an existing seed from disk when one is present, and +// otherwise enrolling with sqi-server over REST when a join token is +// configured. +// +// Enrollment runs over REST rather than NATS on purpose: the broker's entire +// job is to refuse unauthenticated connections, so it cannot also be the +// channel a worker gets its first credential over. mDNS already advertises +// the HTTP port, so zero-configuration discovery still works. +package enroll + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "log/slog" + "net/http" + "os" + "strings" + + "github.com/uberware/sqi/internal/brokerauth" +) + +// ErrNoCredential is returned when no credential file exists and no join +// token is configured to obtain one. It is not inherently fatal to a boot on +// a farm that does not require worker authentication — see +// [EnsureCredential]'s doc for how callers should treat it. +var ErrNoCredential = errors.New("worker: no credential and no join token") + +// Config holds everything EnsureCredential needs to load or obtain a +// credential. It is passed by value, deliberately separate from the broader +// worker configuration struct, so that a seed never has to travel through — +// and risk being logged by — general-purpose configuration plumbing. +type Config struct { + // WorkerID is this worker's stable, server-correlated identity. + WorkerID string + + // CredentialFile is the path to this worker's nkey seed file. + CredentialFile string + + // JoinToken is a worker enrollment token, used once to obtain a + // credential. + JoinToken string + + // JoinTokenFile is a path to a file containing a join token, and takes + // precedence over JoinToken when both are set. + JoinTokenFile string + + // ServerURL is the sqi-server HTTP base URL used for enrollment. + ServerURL string + + // HTTPClient performs the enrollment request. When nil, + // http.DefaultClient is used. + HTTPClient *http.Client +} + +// enrollRequest is the body of POST /api/v1/workers/enroll. +type enrollRequest struct { + JoinToken string `json:"join_token"` + WorkerID string `json:"worker_id"` + PublicKey string `json:"public_key"` +} + +// EnsureCredential returns this worker's nkey seed and public key: loading an +// existing credential file if one is present, and otherwise enrolling with +// sqi-server over REST using a configured join token. +// +// The order follows the spec exactly: an existing seed is loaded first — an +// already-enrolled worker must never re-enroll just because a join token is +// still configured — and only when no seed exists does enrollment happen, +// using whichever token is configured. When neither a seed nor a token is +// available, EnsureCredential returns ErrNoCredential rather than silently +// connecting with nothing: it is the caller's job to decide whether that is +// fatal (a farm that requires authentication) or fine (an auth-off farm, +// where the caller should proceed with no credential and let the broker +// itself decide whether to accept the connection). +// +// A credential is written to disk only after the server confirms enrollment +// succeeded — a failed enrollment must never leave a seed behind that a +// later boot would silently reuse. +func EnsureCredential(ctx context.Context, cfg Config, logger *slog.Logger) (seed []byte, publicKey string, err error) { + seed, loadErr := brokerauth.LoadSeed(cfg.CredentialFile) + switch { + case loadErr == nil: + publicKey, err = brokerauth.PublicKeyFromSeed(seed) + if err != nil { + return nil, "", err + } + logger.InfoContext(ctx, "enroll: loaded existing credential", slog.String("path", cfg.CredentialFile)) + return seed, publicKey, nil + case errors.Is(loadErr, os.ErrNotExist): + // Fall through to enrollment below. + default: + return nil, "", loadErr + } + + token, err := resolveJoinToken(cfg) + if err != nil { + return nil, "", err + } + if token == "" { + return nil, "", ErrNoCredential + } + // ServerURL is never derived from mDNS discovery — enrollment needs an + // HTTP base URL, and mDNS discovery here only ever resolves a NATS URL + // (see internal/worker/discovery). Fail before any HTTP attempt so the + // operator sees a config key to fix rather than a raw + // "unsupported protocol scheme" error. + if cfg.ServerURL == "" { + return nil, "", errors.New( + "worker: enrollment requires nats.server_url (env SQI_WORKER_NATS_SERVER_URL) to be set; it is not derived from mDNS discovery", + ) + } + + seed, publicKey, err = brokerauth.GenerateSeed() + if err != nil { + return nil, "", err + } + + if err := enrollWithServer(ctx, cfg, token, publicKey); err != nil { + return nil, "", err + } + + if err := brokerauth.SaveSeed(cfg.CredentialFile, seed); err != nil { + return nil, "", err + } + logger.InfoContext(ctx, "enroll: obtained new credential", slog.String("path", cfg.CredentialFile)) + return seed, publicKey, nil +} + +// resolveJoinToken returns the join token to enroll with, preferring +// JoinTokenFile over JoinToken when both are set. An empty return with a nil +// error means no token is configured at all. +func resolveJoinToken(cfg Config) (string, error) { + if cfg.JoinTokenFile != "" { + data, err := os.ReadFile(cfg.JoinTokenFile) + if err != nil { + return "", fmt.Errorf("worker: read join token file %s: %w", cfg.JoinTokenFile, err) + } + return strings.TrimSpace(string(data)), nil + } + return cfg.JoinToken, nil +} + +// enrollWithServer posts an enrollment request for publicKey to +// cfg.ServerURL and maps the response to an error naming both the cause and +// the remediation. A nil return means the server confirmed enrollment. +func enrollWithServer(ctx context.Context, cfg Config, token, publicKey string) error { + body, err := json.Marshal(enrollRequest{ + JoinToken: token, + WorkerID: cfg.WorkerID, + PublicKey: publicKey, + }) + if err != nil { + return fmt.Errorf("worker: encode enrollment request: %w", err) + } + + url := strings.TrimRight(cfg.ServerURL, "/") + "/api/v1/workers/enroll" + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body)) + if err != nil { + return fmt.Errorf("worker: build enrollment request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + + client := cfg.HTTPClient + if client == nil { + client = http.DefaultClient + } + + resp, err := client.Do(req) + if err != nil { + return fmt.Errorf("worker: enrollment request to %s: %w", url, err) + } + defer func() { + _, _ = io.Copy(io.Discard, resp.Body) //nolint:errcheck // draining a response body before close; nothing actionable on failure + _ = resp.Body.Close() + }() + + switch resp.StatusCode { + case http.StatusCreated, http.StatusOK: + return nil + case http.StatusUnauthorized, http.StatusForbidden: + return errors.New( + "worker: enrollment refused: the join token is unknown, expired, or already used — issue a new one with `sqi-server worker token issue`", + ) + case http.StatusConflict: + return fmt.Errorf( + "worker: worker id %s is already enrolled with a different key; revoke it with `sqi-server worker revoke %s` or clear this worker's data dir", + cfg.WorkerID, cfg.WorkerID, + ) + default: + return fmt.Errorf("worker: enrollment failed: server returned %s", resp.Status) + } +} diff --git a/internal/worker/enroll/enroll_test.go b/internal/worker/enroll/enroll_test.go new file mode 100644 index 00000000..b1c467ed --- /dev/null +++ b/internal/worker/enroll/enroll_test.go @@ -0,0 +1,264 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package enroll_test + +import ( + "context" + "encoding/json" + "errors" + "log/slog" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/uberware/sqi/internal/brokerauth" + "github.com/uberware/sqi/internal/worker/enroll" +) + +func discardLogger() *slog.Logger { + return slog.New(slog.DiscardHandler) +} + +func TestEnsureCredential_ExistingSeedIsLoadedWithNoHTTPRequest(t *testing.T) { + dir := t.TempDir() + credFile := filepath.Join(dir, "worker.nk") + + wantSeed, wantPub, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed: %v", err) + } + if err := brokerauth.SaveSeed(credFile, wantSeed); err != nil { + t.Fatalf("SaveSeed: %v", err) + } + + requested := false + srv := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) { + requested = true + })) + defer srv.Close() + + cfg := enroll.Config{ + WorkerID: "worker-a", + CredentialFile: credFile, + JoinToken: "should-not-be-used", + ServerURL: srv.URL, + } + seed, pub, err := enroll.EnsureCredential(context.Background(), cfg, discardLogger()) + if err != nil { + t.Fatalf("EnsureCredential: %v", err) + } + if string(seed) != string(wantSeed) { + t.Errorf("seed = %q, want %q", seed, wantSeed) + } + if pub != wantPub { + t.Errorf("public key = %q, want %q", pub, wantPub) + } + if requested { + t.Error("EnsureCredential made an HTTP request despite an existing seed file") + } +} + +func TestEnsureCredential_NoSeedWithTokenEnrollsAndWritesSeed(t *testing.T) { + dir := t.TempDir() + credFile := filepath.Join(dir, "worker.nk") + + var gotBody map[string]string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + t.Errorf("method = %s, want POST", r.Method) + } + if r.URL.Path != "/api/v1/workers/enroll" { + t.Errorf("path = %s, want /api/v1/workers/enroll", r.URL.Path) + } + if err := json.NewDecoder(r.Body).Decode(&gotBody); err != nil { + t.Fatalf("decode request body: %v", err) + } + w.WriteHeader(http.StatusCreated) + })) + defer srv.Close() + + cfg := enroll.Config{ + WorkerID: "worker-a", + CredentialFile: credFile, + JoinToken: "tok-123", + ServerURL: srv.URL, + } + seed, pub, err := enroll.EnsureCredential(context.Background(), cfg, discardLogger()) + if err != nil { + t.Fatalf("EnsureCredential: %v", err) + } + if len(seed) == 0 || pub == "" { + t.Fatalf("EnsureCredential returned empty seed/publicKey: seed=%d pub=%q", len(seed), pub) + } + + if gotBody["join_token"] != "tok-123" { + t.Errorf("join_token = %q, want tok-123", gotBody["join_token"]) + } + if gotBody["worker_id"] != "worker-a" { + t.Errorf("worker_id = %q, want worker-a", gotBody["worker_id"]) + } + if gotBody["public_key"] != pub { + t.Errorf("public_key = %q, want %q", gotBody["public_key"], pub) + } + + info, err := os.Stat(credFile) + if err != nil { + t.Fatalf("Stat credential file: %v", err) + } + if perm := info.Mode().Perm(); perm != 0o600 { + t.Errorf("credential file mode = %o, want 600", perm) + } + + // The credential file now contains the same seed that was returned, so a + // subsequent boot loads it without re-enrolling. + onDisk, err := brokerauth.LoadSeed(credFile) + if err != nil { + t.Fatalf("LoadSeed after enrollment: %v", err) + } + if string(onDisk) != string(seed) { + t.Errorf("seed on disk = %q, want %q", onDisk, seed) + } +} + +func TestEnsureCredential_NoSeedNoTokenReturnsErrNoCredential(t *testing.T) { + dir := t.TempDir() + credFile := filepath.Join(dir, "worker.nk") + + requested := false + srv := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) { + requested = true + })) + defer srv.Close() + + cfg := enroll.Config{ + WorkerID: "worker-a", + CredentialFile: credFile, + ServerURL: srv.URL, + } + _, _, err := enroll.EnsureCredential(context.Background(), cfg, discardLogger()) + if !errors.Is(err, enroll.ErrNoCredential) { + t.Fatalf("EnsureCredential error = %v, want ErrNoCredential", err) + } + if requested { + t.Error("EnsureCredential made an HTTP request with no seed and no token") + } +} + +func TestEnsureCredential_401MentionsTokenAndLeavesNoSeed(t *testing.T) { + dir := t.TempDir() + credFile := filepath.Join(dir, "worker.nk") + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + })) + defer srv.Close() + + cfg := enroll.Config{ + WorkerID: "worker-a", + CredentialFile: credFile, + JoinToken: "stale-token", + ServerURL: srv.URL, + } + _, _, err := enroll.EnsureCredential(context.Background(), cfg, discardLogger()) + if err == nil { + t.Fatal("EnsureCredential: want error for a 401 response, got nil") + } + if !strings.Contains(err.Error(), "token") { + t.Errorf("error %q does not mention the token", err.Error()) + } + + if _, statErr := os.Stat(credFile); !errors.Is(statErr, os.ErrNotExist) { + t.Errorf("Stat(credFile) error = %v, want os.ErrNotExist", statErr) + } +} + +func TestEnsureCredential_409MentionsWorkerIDAlreadyEnrolledAndLeavesNoSeed(t *testing.T) { + dir := t.TempDir() + credFile := filepath.Join(dir, "worker.nk") + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusConflict) + })) + defer srv.Close() + + cfg := enroll.Config{ + WorkerID: "worker-a", + CredentialFile: credFile, + JoinToken: "tok-123", + ServerURL: srv.URL, + } + _, _, err := enroll.EnsureCredential(context.Background(), cfg, discardLogger()) + if err == nil { + t.Fatal("EnsureCredential: want error for a 409 response, got nil") + } + if !strings.Contains(err.Error(), "worker-a") || !strings.Contains(err.Error(), "already enrolled") { + t.Errorf("error %q does not mention the worker id already being enrolled with a different key", err.Error()) + } + + if _, statErr := os.Stat(credFile); !errors.Is(statErr, os.ErrNotExist) { + t.Errorf("Stat(credFile) error = %v, want os.ErrNotExist", statErr) + } +} + +func TestEnsureCredential_TokenConfiguredButNoServerURLFailsBeforeHTTP(t *testing.T) { + dir := t.TempDir() + credFile := filepath.Join(dir, "worker.nk") + + cfg := enroll.Config{ + WorkerID: "worker-a", + CredentialFile: credFile, + JoinToken: "tok-123", + ServerURL: "", // not derived from mDNS — must fail before any HTTP attempt + } + _, _, err := enroll.EnsureCredential(context.Background(), cfg, discardLogger()) + if err == nil { + t.Fatal("EnsureCredential: want error for a join token with no server_url, got nil") + } + if !strings.Contains(err.Error(), "nats.server_url") { + t.Errorf("error %q does not name the nats.server_url config key", err.Error()) + } + if !strings.Contains(err.Error(), "SQI_WORKER_NATS_SERVER_URL") { + t.Errorf("error %q does not name the SQI_WORKER_NATS_SERVER_URL env var", err.Error()) + } + + if _, statErr := os.Stat(credFile); !errors.Is(statErr, os.ErrNotExist) { + t.Errorf("Stat(credFile) error = %v, want os.ErrNotExist", statErr) + } +} + +func TestEnsureCredential_JoinTokenFileTakesPrecedenceOverJoinToken(t *testing.T) { + dir := t.TempDir() + credFile := filepath.Join(dir, "worker.nk") + tokenFile := filepath.Join(dir, "token") + if err := os.WriteFile(tokenFile, []byte("file-token\n"), 0o600); err != nil { + t.Fatalf("write token file: %v", err) + } + + var gotToken string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body map[string]string + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Fatalf("decode request body: %v", err) + } + gotToken = body["join_token"] + w.WriteHeader(http.StatusCreated) + })) + defer srv.Close() + + cfg := enroll.Config{ + WorkerID: "worker-a", + CredentialFile: credFile, + JoinToken: "should-not-be-used", + JoinTokenFile: tokenFile, + ServerURL: srv.URL, + } + if _, _, err := enroll.EnsureCredential(context.Background(), cfg, discardLogger()); err != nil { + t.Fatalf("EnsureCredential: %v", err) + } + if gotToken != "file-token" { + t.Errorf("join_token = %q, want file-token", gotToken) + } +} diff --git a/internal/worker/heartbeat/heartbeat.go b/internal/worker/heartbeat/heartbeat.go index d33756ea..ee052077 100644 --- a/internal/worker/heartbeat/heartbeat.go +++ b/internal/worker/heartbeat/heartbeat.go @@ -6,7 +6,7 @@ // # Overview // // A [Publisher] ticks on a configurable interval and publishes a -// [protocol.HeartbeatMsg] to the [bus.SubjectWorkerHeartbeat] NATS subject. +// [protocol.HeartbeatMsg] to this worker's [bus.WorkerHeartbeatSubject]. // Each message carries the worker's current runtime state (active task count, // active task IDs, uptime, last assignment time) so the server can detect // stale assignments without additional store queries. @@ -108,8 +108,8 @@ type Registrar interface { LastRegisteredAt() time.Time } -// Publisher periodically publishes [protocol.HeartbeatMsg] to -// [bus.SubjectWorkerHeartbeat] and runs an internal watchdog goroutine that +// Publisher periodically publishes [protocol.HeartbeatMsg] to this worker's +// [bus.WorkerHeartbeatSubject] and runs an internal watchdog goroutine that // triggers re-registration when the NATS connection is restored after a drop // and the reconnect callback did not already succeed. // @@ -225,7 +225,7 @@ func (p *Publisher) publish(ctx context.Context) { return } - if err := p.nc.Publish(bus.SubjectWorkerHeartbeat, data); err != nil { + if err := p.nc.Publish(bus.WorkerHeartbeatSubject(p.workerID), data); err != nil { p.logger.WarnContext( ctx, "heartbeat: publish failed", slog.String("worker_id", p.workerID), diff --git a/internal/worker/heartbeat/heartbeat_test.go b/internal/worker/heartbeat/heartbeat_test.go index e3721c56..3edbf407 100644 --- a/internal/worker/heartbeat/heartbeat_test.go +++ b/internal/worker/heartbeat/heartbeat_test.go @@ -15,6 +15,7 @@ import ( "testing" "time" + "github.com/uberware/sqi/internal/bus" "github.com/uberware/sqi/internal/worker/protocol" ) @@ -125,8 +126,9 @@ func TestPublish_CorrectFields(t *testing.T) { p := New(nc, workerID, maxTasks, 15*time.Second, NoopStateSource{}, reg, discardLogger()) p.publish(context.Background()) - if publishedSubj != "worker.heartbeat" { - t.Errorf("published to subject %q, want %q", publishedSubj, "worker.heartbeat") + wantSubj := bus.WorkerHeartbeatSubject(workerID) + if publishedSubj != wantSubj { + t.Errorf("published to subject %q, want %q", publishedSubj, wantSubj) } var msg protocol.HeartbeatMsg diff --git a/internal/worker/lease/lease.go b/internal/worker/lease/lease.go index 5ae46c12..de3a078f 100644 --- a/internal/worker/lease/lease.go +++ b/internal/worker/lease/lease.go @@ -24,8 +24,10 @@ type Config struct { } // Transport sends a lease request and returns the server's reply bytes. +// workerID identifies the requesting worker; it is a token of the subject the +// request travels on, which is how the server attributes the request. type Transport interface { - RequestLease(ctx context.Context, queueID string, data []byte, timeout time.Duration) ([]byte, error) + RequestLease(ctx context.Context, workerID, queueID string, data []byte, timeout time.Duration) ([]byte, error) } // Dispatcher executes one assignment. @@ -85,7 +87,7 @@ func (l *Loop) runQueue(ctx context.Context, queueID string) { if ctx.Err() != nil { return } - data, err := l.transport.RequestLease(ctx, queueID, reqBytes, l.cfg.RequestTimeout) + data, err := l.transport.RequestLease(ctx, l.cfg.WorkerID, queueID, reqBytes, l.cfg.RequestTimeout) if err != nil { // No server / timeout / transient: brief backoff, then re-request. select { diff --git a/internal/worker/lease/lease_test.go b/internal/worker/lease/lease_test.go index 1975f7d6..1cf9635f 100644 --- a/internal/worker/lease/lease_test.go +++ b/internal/worker/lease/lease_test.go @@ -19,7 +19,7 @@ type fakeTransport struct { calls int } -func (f *fakeTransport) RequestLease(_ context.Context, _ string, _ []byte, _ time.Duration) ([]byte, error) { +func (f *fakeTransport) RequestLease(_ context.Context, _, _ string, _ []byte, _ time.Duration) ([]byte, error) { f.mu.Lock() defer f.mu.Unlock() f.calls++ diff --git a/internal/worker/logstreamer/logstreamer.go b/internal/worker/logstreamer/logstreamer.go index 1e81c84b..97a52d62 100644 --- a/internal/worker/logstreamer/logstreamer.go +++ b/internal/worker/logstreamer/logstreamer.go @@ -4,7 +4,7 @@ // // A [Publisher] implements [executor.OutputHandler] and accumulates process // output lines from stdout and stderr into [protocol.LogChunkMsg] batches, -// publishing them to the task.logs. NATS JetStream subject. +// publishing them to the task.logs.. NATS JetStream subject. // // # Chunking // @@ -167,19 +167,25 @@ type Publisher struct { logger *slog.Logger cfg Config + // workerID is this worker's stable identity. It is a subject token on + // every chunk published, which is how the server attributes log output to + // the worker that produced it. + workerID string + // mu protects the attempts map. mu sync.Mutex attempts map[string]*attemptBuf // keyed by attemptID } -// New returns a ready-to-use Publisher. +// New returns a ready-to-use Publisher publishing as workerID. // Any zero or negative fields in cfg are replaced with defaults before use. -func New(nc natsPublisher, cfg Config, logger *slog.Logger) *Publisher { +func New(nc natsPublisher, workerID string, cfg Config, logger *slog.Logger) *Publisher { cfg.applyDefaults() return &Publisher{ nc: nc, logger: logger, cfg: cfg, + workerID: workerID, attempts: make(map[string]*attemptBuf), } } @@ -322,7 +328,7 @@ func (p *Publisher) drainBuf(ctx context.Context, buf *attemptBuf, stream string } // publishChunk assigns the next sequence number, encodes a [protocol.LogChunkMsg], -// and publishes it to the task.logs. NATS subject. +// and publishes it to the task.logs.. NATS subject. // It is a no-op if lines is empty. // // SeqNum assignment and NATS publish are not atomic. When the flushLoop @@ -363,7 +369,7 @@ func (p *Publisher) publishChunk(ctx context.Context, buf *attemptBuf, stream st return } - subj := bus.TaskLogsSubject(buf.taskID) + subj := bus.TaskLogsSubject(p.workerID, buf.taskID) if err := p.nc.Publish(subj, data); err != nil { p.logger.WarnContext( ctx, "logstreamer: publish log chunk failed", diff --git a/internal/worker/logstreamer/logstreamer_test.go b/internal/worker/logstreamer/logstreamer_test.go index ffe86366..1005360b 100644 --- a/internal/worker/logstreamer/logstreamer_test.go +++ b/internal/worker/logstreamer/logstreamer_test.go @@ -18,6 +18,10 @@ import ( "github.com/uberware/sqi/internal/worker/protocol" ) +// testWorkerID is the publishing worker identity every Publisher under test is +// built with; it is a token of every subject those tests observe. +const testWorkerID = "w-test" + // ── Mock NATS ───────────────────────────────────────────────────────────────── // mockNATS captures published messages for assertion in tests. @@ -104,7 +108,7 @@ func TestSequenceNumberMonotonicity(t *testing.T) { nc := &mockNATS{} // MaxLinesPerChunk=3 so every 3rd line triggers an immediate flush. cfg := slowFlushCfg(3, 1<<20) - p := logstreamer.New(nc, cfg, discard()) + p := logstreamer.New(nc, testWorkerID, cfg, discard()) ctx := context.Background() const totalLines = 10 @@ -152,7 +156,7 @@ func TestChunkBoundaryLines(t *testing.T) { nc := &mockNATS{} const maxLines = 3 cfg := slowFlushCfg(maxLines, 1<<20) - p := logstreamer.New(nc, cfg, discard()) + p := logstreamer.New(nc, testWorkerID, cfg, discard()) ctx := context.Background() // Write exactly maxLines lines — should trigger one immediate flush. @@ -192,7 +196,7 @@ func TestChunkBoundaryBytes(t *testing.T) { // Each line is 10 chars; byte accounting = 10+1=11 per line. // MaxBytesPerChunk=20: after line 2 (22 bytes), flush triggers. cfg := slowFlushCfg(1000, 20) - p := logstreamer.New(nc, cfg, discard()) + p := logstreamer.New(nc, testWorkerID, cfg, discard()) ctx := context.Background() const line = "0123456789" // 10 bytes @@ -233,7 +237,7 @@ func TestFlushBeforeTerminalStatus(t *testing.T) { nc := &mockNATS{} // Large thresholds so no immediate flush fires during HandleLine calls. cfg := slowFlushCfg(1000, 1<<20) - p := logstreamer.New(nc, cfg, discard()) + p := logstreamer.New(nc, testWorkerID, cfg, discard()) ctx := context.Background() const nLines = 5 @@ -285,7 +289,7 @@ func TestFlushInterval(t *testing.T) { MaxBytesPerChunk: 1 << 20, FlushInterval: interval, } - p := logstreamer.New(nc, cfg, discard()) + p := logstreamer.New(nc, testWorkerID, cfg, discard()) ctx := context.Background() p.HandleLine(ctx, "task-tick", "attempt-tick", "sess", "stdout", "hello") @@ -319,7 +323,7 @@ func TestFlushInterval(t *testing.T) { func TestStdoutStderrSeparateChunks(t *testing.T) { nc := &mockNATS{} cfg := slowFlushCfg(1000, 1<<20) - p := logstreamer.New(nc, cfg, discard()) + p := logstreamer.New(nc, testWorkerID, cfg, discard()) ctx := context.Background() p.HandleLine(ctx, "task-streams", "attempt-streams", "sess", "stdout", "out-line") @@ -352,7 +356,7 @@ func TestStdoutStderrSeparateChunks(t *testing.T) { func TestMultipleAttemptsIndependentSequences(t *testing.T) { nc := &mockNATS{} cfg := slowFlushCfg(3, 1<<20) - p := logstreamer.New(nc, cfg, discard()) + p := logstreamer.New(nc, testWorkerID, cfg, discard()) ctx := context.Background() // Write 6 lines for attempt-A (triggers two immediate flushes). @@ -408,7 +412,7 @@ func TestMultipleAttemptsIndependentSequences(t *testing.T) { func TestFlushLogsIdempotent(t *testing.T) { nc := &mockNATS{} cfg := slowFlushCfg(100, 1<<20) - p := logstreamer.New(nc, cfg, discard()) + p := logstreamer.New(nc, testWorkerID, cfg, discard()) ctx := context.Background() p.HandleLine(ctx, "task-idem", "attempt-idem", "sess", "stdout", "hello") @@ -431,7 +435,7 @@ func TestFlushLogsIdempotent(t *testing.T) { // received any lines returns nil and publishes nothing. func TestFlushLogsNoLines(t *testing.T) { nc := &mockNATS{} - p := logstreamer.New(nc, logstreamer.DefaultConfig(), discard()) + p := logstreamer.New(nc, testWorkerID, logstreamer.DefaultConfig(), discard()) ctx := context.Background() if err := p.FlushLogs(ctx, "task-empty", "attempt-empty"); err != nil { @@ -443,11 +447,11 @@ func TestFlushLogsNoLines(t *testing.T) { } // TestNATSSubject verifies that log chunks are published to the correct -// task.logs. NATS subject. +// task.logs.. NATS subject. func TestNATSSubject(t *testing.T) { nc := &mockNATS{} cfg := slowFlushCfg(1, 1<<20) // flush on the first line - p := logstreamer.New(nc, cfg, discard()) + p := logstreamer.New(nc, testWorkerID, cfg, discard()) ctx := context.Background() const taskID = "my-task-id" @@ -457,7 +461,7 @@ func TestNATSSubject(t *testing.T) { if len(msgs) == 0 { t.Fatal("no messages published") } - want := bus.TaskLogsSubject(taskID) + want := bus.TaskLogsSubject(testWorkerID, taskID) for _, m := range msgs { if m.subject != want { t.Errorf("subject = %q, want %q", m.subject, want) @@ -470,7 +474,7 @@ func TestNATSSubject(t *testing.T) { func TestPublishErrorDoesNotPanic(t *testing.T) { nc := &mockNATS{pubErr: errors.New("nats: connection closed")} cfg := slowFlushCfg(1, 1<<20) // flush on the first line - p := logstreamer.New(nc, cfg, discard()) + p := logstreamer.New(nc, testWorkerID, cfg, discard()) ctx := context.Background() // Should not panic even though NATS publish returns an error. diff --git a/internal/worker/natsclient/natsclient.go b/internal/worker/natsclient/natsclient.go index a72ee879..24f284c6 100644 --- a/internal/worker/natsclient/natsclient.go +++ b/internal/worker/natsclient/natsclient.go @@ -15,7 +15,7 @@ // // Typical usage: // -// nc, closedCh, err := natsclient.Connect(ctx, cfg.NATS, logger) +// nc, closedCh, err := natsclient.Connect(ctx, cfg.NATS, workerID, seed, publicKey, logger) // if err != nil { // return fmt.Errorf("nats connect: %w", err) // } @@ -36,6 +36,7 @@ import ( nats "github.com/nats-io/nats.go" + "github.com/uberware/sqi/internal/brokerauth" workerconfig "github.com/uberware/sqi/internal/worker/config" ) @@ -66,18 +67,38 @@ const ( // // Callers should select on closedCh to detect permanent disconnects that occur // outside of a planned shutdown sequence. -func Connect(ctx context.Context, cfg workerconfig.NATSConfig, logger *slog.Logger) (*nats.Conn, <-chan struct{}, error) { +// +// seed and publicKey are the worker's nkey broker credential, obtained +// separately (see internal/worker/enroll). They are taken as parameters +// rather than folded into cfg so that a seed never sits in a config struct +// that might get logged elsewhere. An empty seed connects with no +// credential, which is correct on a farm that does not require worker +// authentication; see buildOptions for how that case is distinguished from a +// broker that actively rejects the connection. +// +// workerID scopes this connection's reply inboxes to a per-worker prefix +// (see [brokerauth.InboxPrefix]). It is required whether or not a credential +// is present, so the connect path is identical in both modes, and it must be +// a single NATS subject token. +func Connect(ctx context.Context, cfg workerconfig.NATSConfig, workerID string, seed []byte, publicKey string, logger *slog.Logger) (*nats.Conn, <-chan struct{}, error) { // closedCh is closed by the ClosedHandler callback when the NATS connection // permanently closes (MaxReconnects exhausted or explicit nc.Close() call). closedCh := make(chan struct{}) - opts, err := buildOptions(ctx, cfg, logger, closedCh) + opts, err := buildOptions(ctx, cfg, workerID, seed, publicKey, logger, closedCh) if err != nil { return nil, nil, fmt.Errorf("natsclient: build options: %w", err) } nc, err := nats.Connect(cfg.URL, opts...) if err != nil { + // An authorization failure is FATAL and must never enter the + // reconnect-backoff loop: retrying a rejected credential produces a + // worker that never appears, with the reason buried in backoff. An + // unreachable server is the opposite — that is what backoff is for. + if wrapped, ok := credentialRejectedError(err); ok { + return nil, nil, wrapped + } return nil, nil, fmt.Errorf("natsclient: connect %q: %w", cfg.URL, err) } @@ -90,6 +111,27 @@ func Connect(ctx context.Context, cfg workerconfig.NATSConfig, logger *slog.Logg return nc, closedCh, nil } +// credentialRejectedError reports whether err is the broker rejecting this +// connection's nkey credential — nats.ErrAuthorization (unknown or wrong +// key) or nats.ErrAuthExpired (a key the broker no longer accepts, notably +// after a live revocation: internal/bus.Broker.ReloadCredentials drops the +// key and synchronously disconnects any client using it). When it is, the +// second return is true and the first is the single crafted message every +// credential-rejection path uses — the initial dial in [Connect] and a later +// live revocation observed via [nats.Conn.LastError] in the ClosedHandler +// below both call this, so the two paths cannot drift apart. +func credentialRejectedError(err error) (error, bool) { + if err == nil { + return nil, false + } + if !errors.Is(err, nats.ErrAuthorization) && !errors.Is(err, nats.ErrAuthExpired) { + return nil, false + } + return fmt.Errorf( + "natsclient: the broker rejected this worker's credential — it may have been revoked; re-enroll with a new join token: %w", err, + ), true +} + // Drain gracefully closes nc by draining in-flight subscriptions and flushing // any pending publishes before closing the connection. It blocks until the // drain completes or gracePeriod elapses. If the grace period expires first, @@ -119,8 +161,19 @@ func Drain(nc *nats.Conn, gracePeriod time.Duration, logger *slog.Logger) { // buildOptions assembles the nats.Option slice from WorkerNATSConfig. // closedCh is closed by the ClosedHandler when the connection permanently -// closes so callers can detect unexpected disconnects. -func buildOptions(ctx context.Context, cfg workerconfig.NATSConfig, logger *slog.Logger, closedCh chan struct{}) ([]nats.Option, error) { +// closes so callers can detect unexpected disconnects. seed and publicKey, +// when non-empty, add an nkey signing option so the connection authenticates +// as this worker's broker credential. workerID scopes the connection's reply +// inboxes to this worker's own prefix and is required. +func buildOptions( + ctx context.Context, + cfg workerconfig.NATSConfig, + workerID string, + seed []byte, + publicKey string, + logger *slog.Logger, + closedCh chan struct{}, +) ([]nats.Option, error) { opts := []nats.Option{ nats.MaxReconnects(cfg.MaxReconnectAttempts), @@ -152,8 +205,21 @@ func buildOptions(ctx context.Context, cfg workerconfig.NATSConfig, logger *slog // state: either MaxReconnects was exhausted or the connection was // explicitly closed. Closing closedCh signals any goroutine that is // watching for unexpected permanent disconnects. - nats.ClosedHandler(func(_ *nats.Conn) { - logger.InfoContext(ctx, "natsclient: connection closed") + // + // A live credential revocation (internal/bus.Broker.ReloadCredentials + // disconnects the client synchronously via authViolation) lands here, + // not at the initial dial in [Connect] — the connection was already + // established when the credential stopped being valid. nats.Conn's + // own doc for LastError says it "can be used reliably within + // ClosedCB in order to find out reason why connection was closed", + // so this is the one place that case can be classified and named, + // rather than surfacing only as a generic closure to the operator. + nats.ClosedHandler(func(nc *nats.Conn) { + if wrapped, ok := credentialRejectedError(nc.LastError()); ok { + logger.ErrorContext(ctx, "natsclient: connection closed", slog.Any("error", wrapped)) + } else { + logger.InfoContext(ctx, "natsclient: connection closed") + } close(closedCh) }), nats.ErrorHandler(func(_ *nats.Conn, _ *nats.Subscription, err error) { @@ -161,6 +227,29 @@ func buildOptions(ctx context.Context, cfg workerconfig.NATSConfig, logger *slog }), } + // ── Per-worker reply inbox ─────────────────────────────────── + // + // Without this the connection takes nats.go's process-global "_INBOX" + // prefix, and the only permission that could then cover a lease reply + // is "_INBOX.>" — which covers every OTHER client's reply inbox on the + // same broker too, handing any enrolled worker the assignment batches + // of work it never leased. See [brokerauth.InboxPrefix]. + // + // Applied whether or not a credential is present, so the connect path + // does not diverge between the auth-on and auth-off modes. An absent + // worker ID, or one that is not a single subject token, would leave the + // connection on the shared prefix or silently widen the subtree the + // matching grant covers — so it is rejected rather than trusted: + // LoadOrCreateWorkerID writes a UUID, but the file it writes is one an + // operator can edit. + if !brokerauth.ValidWorkerIDToken(workerID) { + return nil, fmt.Errorf( + "natsclient: worker id %q is not a valid NATS subject token — it must be non-empty and must not contain '.', whitespace, '*' or '>'", + workerID, + ) + } + opts = append(opts, nats.CustomInboxPrefix(brokerauth.InboxPrefix(workerID))) + // ── TLS ────────────────────────────────────────────────────── tlsOpts, err := buildTLSOptions(cfg) if err != nil { @@ -168,6 +257,20 @@ func buildOptions(ctx context.Context, cfg workerconfig.NATSConfig, logger *slog } opts = append(opts, tlsOpts...) + // ── Broker credential ──────────────────────────────────────── + // + // A non-empty seed authenticates this connection as the worker's + // enrolled nkey. An empty seed adds no option at all, which is correct + // on a farm that does not require worker authentication — the broker + // accepts the anonymous connection exactly as it does today. When the + // broker DOES require authentication, connecting with no credential (or + // a rejected one) fails at nats.Connect above with nats.ErrAuthorization + // or nats.ErrAuthExpired, which is classified as fatal rather than + // retried. + if len(seed) > 0 { + opts = append(opts, brokerauth.NkeyOption(publicKey, seed)) + } + return opts, nil } diff --git a/internal/worker/natsclient/natsclient_test.go b/internal/worker/natsclient/natsclient_test.go index dafd6773..c2bcfdc6 100644 --- a/internal/worker/natsclient/natsclient_test.go +++ b/internal/worker/natsclient/natsclient_test.go @@ -3,11 +3,20 @@ package natsclient import ( + "bytes" "context" + "errors" "log/slog" + "net" + "strings" "testing" "time" + nats "github.com/nats-io/nats.go" + "github.com/nats-io/nkeys" + + "github.com/uberware/sqi/internal/brokerauth" + "github.com/uberware/sqi/internal/bus" workerconfig "github.com/uberware/sqi/internal/worker/config" ) @@ -80,8 +89,279 @@ func TestConnect_DialErrorIsReturned(t *testing.T) { MaxReconnectAttempts: 0, ReconnectWait: 10 * time.Millisecond, } - _, _, err := Connect(context.Background(), cfg, logger) + _, _, err := Connect(context.Background(), cfg, "worker-a", nil, "", logger) if err == nil { t.Fatal("Connect to dead port: want error, got nil") } } + +func TestBuildOptions_AddsNkeyOptionWhenSeedPresent(t *testing.T) { + logger := slog.New(slog.DiscardHandler) + seed, pub, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed: %v", err) + } + + opts, err := buildOptions(context.Background(), workerconfig.NATSConfig{}, "worker-a", seed, pub, logger, make(chan struct{})) + if err != nil { + t.Fatalf("buildOptions: %v", err) + } + + applied := &nats.Options{} + for _, opt := range opts { + if err := opt(applied); err != nil { + t.Fatalf("apply option: %v", err) + } + } + if applied.Nkey != pub { + t.Errorf("Options.Nkey = %q, want %q", applied.Nkey, pub) + } + if applied.SignatureCB == nil { + t.Fatal("Options.SignatureCB is nil; want a signing callback") + } + + // The callback must actually sign with the given seed, not just be + // present — verify a nonce signs against the matching public key. + nonce := []byte("test-nonce") + sig, err := applied.SignatureCB(nonce) + if err != nil { + t.Fatalf("SignatureCB: %v", err) + } + kp, err := nkeys.FromPublicKey(pub) + if err != nil { + t.Fatalf("FromPublicKey: %v", err) + } + if err := kp.Verify(nonce, sig); err != nil { + t.Errorf("signature does not verify against %s: %v", pub, err) + } +} + +func TestBuildOptions_NoNkeyOptionWhenSeedEmpty(t *testing.T) { + logger := slog.New(slog.DiscardHandler) + opts, err := buildOptions(context.Background(), workerconfig.NATSConfig{}, "worker-a", nil, "", logger, make(chan struct{})) + if err != nil { + t.Fatalf("buildOptions: %v", err) + } + + applied := &nats.Options{} + for _, opt := range opts { + if err := opt(applied); err != nil { + t.Fatalf("apply option: %v", err) + } + } + if applied.Nkey != "" { + t.Errorf("Options.Nkey = %q, want empty", applied.Nkey) + } + if applied.SignatureCB != nil { + t.Error("Options.SignatureCB is set; want nil when no seed is configured") + } +} + +// freePort asks the OS for an unused loopback TCP port, for booting a +// throwaway embedded broker. +func freePort(t *testing.T) int { + t.Helper() + var lc net.ListenConfig + l, err := lc.Listen(context.Background(), "tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("freePort: listen: %v", err) + } + defer func() { _ = l.Close() }() + addr, ok := l.Addr().(*net.TCPAddr) + if !ok { + t.Fatalf("freePort: listener address is %T, want *net.TCPAddr", l.Addr()) + } + return addr.Port +} + +// startTestBroker boots a real embedded NATS broker with auth enabled and +// exactly one enrolled worker credential, on a throwaway loopback port and +// JetStream dir. It exists to prove [Connect]'s classification of a rejected +// credential against a REAL broker, not a mock — an nkey auth handshake is +// exactly the kind of wire behavior a fake cannot reproduce faithfully. +func startTestBroker(t *testing.T, enrolled bus.WorkerCredentialRef) *bus.Broker { + t.Helper() + b := bus.New(bus.BrokerConfig{ + Addr: net.JoinHostPort("127.0.0.1", itoa(freePort(t))), + DataDir: t.TempDir() + "/nats", + Auth: bus.BrokerAuthConfig{ + Enabled: true, + Credentials: []bus.WorkerCredentialRef{enrolled}, + }, + }, slog.New(slog.DiscardHandler)) + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + if err := b.Start(ctx); err != nil { + t.Fatalf("startTestBroker: Start: %v", err) + } + t.Cleanup(b.Shutdown) + return b +} + +func itoa(p int) string { + return (&net.TCPAddr{Port: p}).String()[1:] // ":"[1:] == "" +} + +func TestConnect_ClassifiesRejectedCredentialAsFatal(t *testing.T) { + _, enrolledPub, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed: %v", err) + } + b := startTestBroker(t, bus.WorkerCredentialRef{WorkerID: "worker-a", PublicKey: enrolledPub}) + + // A freshly generated keypair that was never enrolled with the broker. + strangerSeed, _, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed: %v", err) + } + strangerPub, err := brokerauth.PublicKeyFromSeed(strangerSeed) + if err != nil { + t.Fatalf("PublicKeyFromSeed: %v", err) + } + + cfg := workerconfig.NATSConfig{ + URL: b.ClientURL(), + MaxReconnectAttempts: 0, + ReconnectWait: 10 * time.Millisecond, + } + logger := slog.New(slog.DiscardHandler) + _, _, connErr := Connect(context.Background(), cfg, "worker-a", strangerSeed, strangerPub, logger) + if connErr == nil { + t.Fatal("Connect with an unenrolled nkey: want error, got nil") + } + if !errors.Is(connErr, nats.ErrAuthorization) && !errors.Is(connErr, nats.ErrAuthExpired) { + t.Errorf("Connect error = %v, want it to wrap nats.ErrAuthorization or nats.ErrAuthExpired", connErr) + } +} + +func TestConnect_AuthOffFarmConnectsWithNoCredential(t *testing.T) { + b := bus.New(bus.BrokerConfig{ + Addr: net.JoinHostPort("127.0.0.1", itoa(freePort(t))), + DataDir: t.TempDir() + "/nats", + Auth: bus.BrokerAuthConfig{Enabled: false}, + }, slog.New(slog.DiscardHandler)) + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + if err := b.Start(ctx); err != nil { + t.Fatalf("Start: %v", err) + } + t.Cleanup(b.Shutdown) + + cfg := workerconfig.NATSConfig{ + URL: b.ClientURL(), + MaxReconnectAttempts: 0, + ReconnectWait: 10 * time.Millisecond, + } + logger := slog.New(slog.DiscardHandler) + nc, _, err := Connect(context.Background(), cfg, "worker-a", nil, "", logger) + if err != nil { + t.Fatalf("Connect with no credential against an auth-off broker: %v", err) + } + defer nc.Close() + if !nc.IsConnected() { + t.Error("connection is not in the connected state") + } +} + +// TestConnect_LiveRevocationNamesCauseAndRemediation covers the common +// revocation shape in this project, not just a rejected initial dial: +// internal/bus.Broker.ReloadCredentials drops an enrolled key and +// synchronously disconnects the client that was using it (authViolation on +// the server side). The worker is already connected when this happens, so +// the classification has to happen in the ClosedHandler, not at the initial +// nats.Connect call — this test drives that path against a real broker and +// asserts the SAME crafted cause+remediation text reaches the logger, not a +// generic "connection closed" line. +func TestConnect_LiveRevocationNamesCauseAndRemediation(t *testing.T) { + enrolledSeed, enrolledPub, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed: %v", err) + } + b := startTestBroker(t, bus.WorkerCredentialRef{WorkerID: "worker-a", PublicKey: enrolledPub}) + + var buf bytes.Buffer + logger := slog.New(slog.NewTextHandler(&buf, &slog.HandlerOptions{Level: slog.LevelDebug})) + + cfg := workerconfig.NATSConfig{ + URL: b.ClientURL(), + // Reconnect enabled (unlimited, a short wait) deliberately — + // matching a real worker's config, NOT nats.NoReconnect(). nats.go + // only lands nats.ErrAuthorization in nc.LastError() by way of the + // RECONNECT handshake itself hitting the same rejected credential + // twice (see nats.go's processConnectInit / nc.ar): the very first + // disconnect races an async "-ERR authorization violation" against + // the raw TCP close, and that race's loser is whichever error + // happened to land in nc.err last — with reconnect enabled, that + // race is irrelevant because a fresh CONNECT attempt reproduces the + // auth error deterministically, which is what makes this + // classification reliable in production rather than a coin flip. + MaxReconnectAttempts: -1, + ReconnectWait: 5 * time.Millisecond, + } + nc, closedCh, err := Connect(context.Background(), cfg, "worker-a", enrolledSeed, enrolledPub, logger) + if err != nil { + t.Fatalf("Connect with an enrolled credential: %v", err) + } + defer nc.Close() + + // Revoke the only enrolled credential. Revocation is synchronous on the + // broker side; the deadline below tolerates client-side scheduling + // jitter in observing the resulting close, not a slow revocation. + if err := b.ReloadCredentials(nil); err != nil { + t.Fatalf("ReloadCredentials: %v", err) + } + + select { + case <-closedCh: + case <-time.After(5 * time.Second): + t.Fatal("connection was not closed after credential revocation") + } + + logs := buf.String() + if !strings.Contains(logs, "the broker rejected this worker's credential") { + t.Errorf("logs do not name the cause; got:\n%s", logs) + } + if !strings.Contains(logs, "re-enroll with a new join token") { + t.Errorf("logs do not name the remediation; got:\n%s", logs) + } +} + +// TestBuildOptions_SetsPerWorkerInboxPrefix pins the option that keeps a +// worker's lease replies out of every other worker's reach: without it the +// connection takes nats.go's process-global "_INBOX", and the only grant +// that could cover a reply would also cover every other client's inbox on +// the same broker. The matching permission is asserted in +// internal/brokerauth; the end-to-end consequence in internal/bus. +func TestBuildOptions_SetsPerWorkerInboxPrefix(t *testing.T) { + logger := slog.New(slog.DiscardHandler) + const workerID = "0f1d2c3b-4a59-6879-8a9b-0c1d2e3f4a5b" + + opts, err := buildOptions(context.Background(), workerconfig.NATSConfig{}, workerID, nil, "", logger, make(chan struct{})) + if err != nil { + t.Fatalf("buildOptions: %v", err) + } + applied := &nats.Options{} + for _, opt := range opts { + if err := opt(applied); err != nil { + t.Fatalf("apply option: %v", err) + } + } + if want := brokerauth.InboxPrefix(workerID); applied.InboxPrefix != want { + t.Errorf("Options.InboxPrefix = %q, want %q", applied.InboxPrefix, want) + } +} + +// TestBuildOptions_RejectsWorkerIDThatIsNotASubjectToken covers the +// defensive check: worker.id holds a UUID, but it is a file an operator can +// edit. An absent ID would leave the connection on nats.go's shared "_INBOX" +// prefix, and one carrying a "." or a wildcard would widen the subtree the +// broker grant built from it covers. +func TestBuildOptions_RejectsWorkerIDThatIsNotASubjectToken(t *testing.T) { + logger := slog.New(slog.DiscardHandler) + for _, workerID := range []string{"", "a.b", "a b", "*", ">"} { + if _, err := buildOptions(context.Background(), workerconfig.NATSConfig{}, workerID, nil, "", logger, make(chan struct{})); err == nil { + t.Errorf("buildOptions with worker id %q: want error, got nil", workerID) + } + } +} diff --git a/internal/worker/openjd/interceptor.go b/internal/worker/openjd/interceptor.go index 01c3bd49..f56afc4a 100644 --- a/internal/worker/openjd/interceptor.go +++ b/internal/worker/openjd/interceptor.go @@ -327,7 +327,7 @@ func (i *Interceptor) handleStatus(ctx context.Context, attemptID, line string) // Intermediate openjd_status publishes are best-effort: they are live UI // updates, not terminal state transitions, so a single transient failure is // logged and dropped rather than retried (unlike the status.Publisher path). - if err := i.nc.Publish(bus.TaskStatusSubject(st.jobID), data); err != nil { + if err := i.nc.Publish(bus.TaskStatusSubject(i.workerID, st.jobID), data); err != nil { i.logger.WarnContext( ctx, "openjd: publish status update failed", slog.String("attempt_id", attemptID), diff --git a/internal/worker/openjd/interceptor_test.go b/internal/worker/openjd/interceptor_test.go index 030300ce..2ff68081 100644 --- a/internal/worker/openjd/interceptor_test.go +++ b/internal/worker/openjd/interceptor_test.go @@ -188,7 +188,7 @@ func TestStatusDirectivePublishesNATS(t *testing.T) { if len(msgs) != 1 { t.Fatalf("expected 1 NATS message, got %d", len(msgs)) } - wantSubj := bus.TaskStatusSubject("job-1") + wantSubj := bus.TaskStatusSubject("test-worker", "job-1") if msgs[0].subject != wantSubj { t.Errorf("NATS subject = %q, want %q", msgs[0].subject, wantSubj) } diff --git a/internal/worker/protocol/diag.go b/internal/worker/protocol/diag.go index 87477aa3..7bb9a8a9 100644 --- a/internal/worker/protocol/diag.go +++ b/internal/worker/protocol/diag.go @@ -7,7 +7,7 @@ import "time" // DiagLogMsg is a single diagnostic (operational) log record published by // sqi-worker to the core-NATS subject worker.diag.. It carries the // worker's own slog output — distinct from task process output, which flows via -// LogChunkMsg on task.logs.. +// LogChunkMsg on task.logs... // // Published with core NATS (not JetStream): delivery is best-effort and nothing // is retained on the broker. The server holds a bounded in-memory ring buffer. diff --git a/internal/worker/protocol/protocol.go b/internal/worker/protocol/protocol.go index e6659209..7c6f329f 100644 --- a/internal/worker/protocol/protocol.go +++ b/internal/worker/protocol/protocol.go @@ -5,13 +5,18 @@ // // # Message flow // -// Each message type flows in a specific direction over NATS JetStream: +// Each message type flows in a specific direction over NATS. Five of the six +// ride JetStream streams for at-least-once delivery; AssignMsg is the +// exception, carried as the reply half of a core-NATS request/reply exchange +// with no stream behind it (see internal/bus for the full subject/stream +// split): // -// RegisterMsg worker → server (worker.register) -// HeartbeatMsg worker → server (worker.heartbeat) -// AssignMsg server → worker (work.assign.) -// TaskStatusMsg worker → server (task.status.) -// LogChunkMsg worker → server (task.logs.) +// RegisterMsg worker → server JetStream (worker.register.) +// DeregisterMsg worker → server JetStream (worker.deregister.) +// HeartbeatMsg worker → server JetStream (worker.heartbeat.) +// AssignMsg server → worker core NATS (reply to work.lease..) +// TaskStatusMsg worker → server JetStream (task.status..) +// LogChunkMsg worker → server JetStream (task.logs..) // // # Versioning // @@ -63,7 +68,15 @@ import "time" // (internal/worker/lease.decodeAssignment) keys off, and why workers must // be upgraded AFTER the server -- a "2" worker rejects every assignment a // "1" server offers it and the tasks churn through reclaim. -const ProtocolVersion = "2" +// +// "3" moves the publishing worker's identity into every worker → server NATS +// subject: task.status.., task.logs.., +// work.lease.., and worker.{register,heartbeat,deregister}. +// (see the subject table in internal/bus). No message body changes, but the +// routing does: a worker publishing on the shorter subjects reaches no +// consumer at all, so this bump is not merely advisory. As with "2", upgrade +// the server first. +const ProtocolVersion = "3" // ── Message type constants ──────────────────────────────────────────────────── @@ -88,7 +101,7 @@ const ( // ── RegisterMsg ─────────────────────────────────────────────────────────────── -// RegisterMsg is the JSON payload workers publish to worker.register. +// RegisterMsg is the JSON payload workers publish to worker.register.. // Workers MUST publish this on first connect and on every reconnect so the // server always has a current view of their capabilities. // @@ -233,7 +246,7 @@ type GPUInfo struct { // ── DeregisterMsg ───────────────────────────────────────────────────────────── -// DeregisterMsg is the JSON payload workers publish to worker.deregister on +// DeregisterMsg is the JSON payload workers publish to worker.deregister. on // graceful shutdown. The server marks the worker offline immediately upon // receipt, rather than waiting for the heartbeat timeout sweep. // @@ -256,7 +269,7 @@ type DeregisterMsg struct { // ── HeartbeatMsg ───────────────────────────────────────────────────────────── -// HeartbeatMsg is the JSON payload workers publish to worker.heartbeat on a +// HeartbeatMsg is the JSON payload workers publish to worker.heartbeat. on a // regular interval. The server uses these to track worker liveness; workers // that stop sending heartbeats are marked offline by the heartbeat sweep after // [scheduler.Config.WorkerTimeout] elapses. @@ -368,8 +381,8 @@ type AssignMsg struct { // user, which is the pre-isolation behavior. // // This carries a USERNAME ONLY and never a credential. Worker↔server - // transport authentication does not exist yet (deferred to Phase 4), so - // nothing secret may travel on this channel. + // transport authentication is opt-in (nats.auth.enabled) and off by + // default, so nothing secret may travel on this channel regardless. Isolation *IsolationSpec `json:"isolation,omitempty"` // ── Parameter space ─────────────────────────────────────────────────── @@ -593,7 +606,7 @@ type StageEntry struct { // ── TaskStatusMsg ───────────────────────────────────────────────────────────── -// TaskStatusMsg is the JSON payload workers publish to task.status. +// TaskStatusMsg is the JSON payload workers publish to task.status.. // to report a task-state transition. // // A worker MUST publish: @@ -658,7 +671,7 @@ type TaskStatusMsg struct { // ── LogChunkMsg ─────────────────────────────────────────────────────────────── -// LogChunkMsg is the JSON payload workers publish to task.logs. as a +// LogChunkMsg is the JSON payload workers publish to task.logs.. as a // running task emits output. // // Workers SHOULD publish log chunks continuously as output is produced rather diff --git a/internal/worker/registration/registration.go b/internal/worker/registration/registration.go index 03dd386e..3370471e 100644 --- a/internal/worker/registration/registration.go +++ b/internal/worker/registration/registration.go @@ -6,9 +6,9 @@ // // A [Registrar] is created once at worker startup, after the NATS connection // is established. Its [Registrar.Register] method publishes a [protocol.RegisterMsg] -// to the worker.register JetStream subject. On graceful shutdown, +// to its own worker.register. JetStream subject. On graceful shutdown, // [Registrar.Deregister] publishes a [protocol.DeregisterMsg] to -// worker.deregister so the server marks the worker offline immediately. +// worker.deregister. so the server marks the worker offline immediately. // // # Re-registration // @@ -23,7 +23,7 @@ // returns only once the SQI_WORKER stream has durably stored the message. A // plain core-NATS publish would be silently discarded when no stream is behind // the subject — which strands the worker permanently, since the server only -// learns of it via worker.register and NAKs the heartbeats of workers it has no +// learns of it via worker.register. and NAKs the heartbeats of workers it has no // record of. // // That "no stream yet" window is real and routinely hit: a worker started @@ -33,7 +33,7 @@ // absent, up to [RetryBudget] or the caller's context deadline, whichever is // sooner. // -// The server processes worker.register via a JetStream push consumer and does +// The server processes worker.register. via a JetStream push consumer and does // not send a reply of its own; the stream ack is the acknowledgment. Explicit // application-level accept/reject is a planned protocol enhancement. // @@ -146,7 +146,7 @@ func New( }, nil } -// Register publishes a RegisterMsg to worker.register and returns once the +// Register publishes a RegisterMsg to worker.register. and returns once the // stream has acked it. It is safe to call multiple times (at boot and on NATS // reconnect). // @@ -229,8 +229,9 @@ func (r *Registrar) Register(ctx context.Context) error { return nil } -// publishRegister publishes data to worker.register and waits for the stream to -// ack it, retrying while the SQI_WORKER stream does not exist yet. +// publishRegister publishes data to this worker's worker.register subject and +// waits for the stream to ack it, retrying while the SQI_WORKER stream does not +// exist yet. // // Only a missing stream is retried. Every other failure (a closed connection, // say) is returned immediately: it will not resolve by waiting. @@ -238,14 +239,15 @@ func (r *Registrar) publishRegister(ctx context.Context, data []byte) error { ctx, cancel := context.WithTimeout(ctx, RetryBudget) defer cancel() + subj := bus.WorkerRegisterSubject(r.workerID) backoff := retryInitial for attempt := 1; ; attempt++ { - _, err := r.js.Publish(ctx, bus.SubjectWorkerRegister, data) + _, err := r.js.Publish(ctx, subj, data) if err == nil { return nil } if !errors.Is(err, jetstream.ErrNoStreamResponse) { - return fmt.Errorf("registration: publish to %s: %w", bus.SubjectWorkerRegister, err) + return fmt.Errorf("registration: publish to %s: %w", subj, err) } r.logger.DebugContext( @@ -259,7 +261,7 @@ func (r *Registrar) publishRegister(ctx context.Context, data []byte) error { case <-ctx.Done(): return fmt.Errorf( "registration: publish to %s: no stream after %d attempts (is the server provisioning JetStream?): %w", - bus.SubjectWorkerRegister, attempt, ctx.Err(), + subj, attempt, ctx.Err(), ) case <-time.After(backoff): } @@ -267,7 +269,7 @@ func (r *Registrar) publishRegister(ctx context.Context, data []byte) error { } } -// Deregister publishes a DeregisterMsg to worker.deregister on graceful +// Deregister publishes a DeregisterMsg to worker.deregister. on graceful // shutdown so the server marks this worker offline immediately rather than // waiting for the heartbeat timeout sweep. // @@ -292,7 +294,7 @@ func (r *Registrar) Deregister(reason string) { return } - if err := r.nc.Publish(bus.SubjectWorkerDeregister, data); err != nil { + if err := r.nc.Publish(bus.WorkerDeregisterSubject(r.workerID), data); err != nil { r.logger.WarnContext(ctx, "registration: deregister publish failed — server will detect absence via heartbeat timeout", slog.String("worker_id", r.workerID), slog.Any("error", err)) diff --git a/internal/worker/registration/registration_extra_test.go b/internal/worker/registration/registration_extra_test.go index 24e344db..eccc4127 100644 --- a/internal/worker/registration/registration_extra_test.go +++ b/internal/worker/registration/registration_extra_test.go @@ -74,12 +74,12 @@ func connectReconnect(tb testing.TB, url string) *nats.Conn { // ── Deregister ────────────────────────────────────────────────────────────── // TestDeregister_PublishesMessage asserts the worker publishes a well-formed -// DeregisterMsg to worker.deregister. +// DeregisterMsg to worker.deregister.. func TestDeregister_PublishesMessage(t *testing.T) { url := startTestNATS(t) nc := connectNATS(t, url) - sub, err := nc.SubscribeSync(bus.SubjectWorkerDeregister) + sub, err := nc.SubscribeSync(bus.WorkerDeregisterSubject("worker-bye")) if err != nil { t.Fatalf("SubscribeSync: %v", err) } @@ -151,7 +151,7 @@ func TestRegister_SingleQueueAndCapabilities(t *testing.T) { Tags: map[string]string{"role": "gpu"}, } - sub, err := nc.SubscribeSync(bus.SubjectWorkerRegister) + sub, err := nc.SubscribeSync(bus.WorkerRegisterSubject("worker-q")) if err != nil { t.Fatalf("SubscribeSync: %v", err) } @@ -225,7 +225,7 @@ func TestRegister_MultiQueue_LeavesQueueIDEmpty(t *testing.T) { cfg := minimalCfg() cfg.QueueIDs = []string{"q1", "q2"} - sub, err := nc.SubscribeSync(bus.SubjectWorkerRegister) + sub, err := nc.SubscribeSync(bus.WorkerRegisterSubject("worker-multi")) if err != nil { t.Fatalf("SubscribeSync: %v", err) } @@ -305,7 +305,7 @@ func TestSetupReconnectHook_ReregistersOnReconnect(t *testing.T) { reg := newRegistrar(t, wnc, "worker-recon", minimalCfg(), capabilities.Capabilities{OS: "linux", CPUCount: 8}) - sub, err := wnc.SubscribeSync(bus.SubjectWorkerRegister) + sub, err := wnc.SubscribeSync(bus.WorkerRegisterSubject("worker-recon")) if err != nil { t.Fatalf("SubscribeSync: %v", err) } diff --git a/internal/worker/registration/registration_test.go b/internal/worker/registration/registration_test.go index 41780f2e..4c36eebc 100644 --- a/internal/worker/registration/registration_test.go +++ b/internal/worker/registration/registration_test.go @@ -89,9 +89,9 @@ func createWorkerStream(url string) error { if _, err := js.CreateStream(ctx, jetstream.StreamConfig{ Name: bus.StreamWorker, Subjects: []string{ - bus.SubjectWorkerRegister, - bus.SubjectWorkerHeartbeat, - bus.SubjectWorkerDeregister, + bus.SubjectWorkerRegisterPrefix + ".>", + bus.SubjectWorkerHeartbeatPrefix + ".>", + bus.SubjectWorkerDeregisterPrefix + ".>", }, Retention: jetstream.WorkQueuePolicy, Storage: jetstream.MemoryStorage, diff --git a/internal/worker/status/publisher.go b/internal/worker/status/publisher.go index 230d3ccf..38dbe817 100644 --- a/internal/worker/status/publisher.go +++ b/internal/worker/status/publisher.go @@ -3,7 +3,7 @@ // Package status implements the typed task-status publisher for sqi-worker. // // A [Publisher] publishes task state-transition messages to the -// task.status. NATS subject. It covers the four transitions the worker +// task.status.. NATS subject. It covers the four transitions the worker // is responsible for: "running", "succeeded", "failed", and "canceled". // // # Fields @@ -76,7 +76,7 @@ type Config struct { // ── Publisher ───────────────────────────────────────────────────────────────── -// Publisher publishes typed task-status messages to task.status.. +// Publisher publishes typed task-status messages to task.status... // // It injects worker_id and last_progress into every message and // retries transient NATS publish failures with exponential backoff. @@ -208,7 +208,7 @@ func (p *Publisher) ShutdownFailed(ctx context.Context, tasks []ShutdownTask) { // ── Internal helpers ────────────────────────────────────────────────────────── -// publishWithRetry marshals msg and publishes it to task.status. with up +// publishWithRetry marshals msg and publishes it to task.status.. with up // to cfg.MaxRetries retries on failure, using exponential backoff. // // On a transient NATS failure the method waits cfg.RetryDelay (doubling each @@ -229,7 +229,7 @@ func (p *Publisher) publishWithRetry(ctx context.Context, msg protocol.TaskStatu return } - subj := bus.TaskStatusSubject(msg.JobID) + subj := bus.TaskStatusSubject(p.cfg.WorkerID, msg.JobID) delay := p.cfg.RetryDelay for attempt := 0; attempt <= p.cfg.MaxRetries; attempt++ { diff --git a/scripts/auth-demo.sh b/scripts/auth-demo.sh index 0e47e3fc..1ee00b3d 100644 --- a/scripts/auth-demo.sh +++ b/scripts/auth-demo.sh @@ -391,9 +391,11 @@ assert_eq "PATCH /auth/me ignored role escalation" "user" \ step "start sqi-worker — note it is given NO credentials" -# Worker<->server transport auth is deliberately out of scope for Phase 3 -# (deferred to Phase 4 hardening). The worker connects to the embedded NATS -# broker unauthenticated; enabling auth changes the REST/UI surface only. +# Broker authentication (nats.auth.enabled) is a separate, opt-in gate from +# the auth.enabled this demo exercises — see docs/auth.md's transport +# section. This demo leaves it off (the default), so the worker connects to +# the embedded NATS broker with no credential; auth.enabled changes the +# REST/UI surface only, not the worker transport. SQI_WORKER_NATS_URL="nats://${NATS_ADDR}" \ SQI_WORKER_DISCOVERY_ENABLE_MDNS="false" \ SQI_WORKER_FARM_ID="$FARM_ID" \ diff --git a/scripts/smoke.sh b/scripts/smoke.sh index 40f82c2c..d61de17f 100644 --- a/scripts/smoke.sh +++ b/scripts/smoke.sh @@ -14,8 +14,16 @@ # the VALUE its expressions resolved to — see the "EXPR job" section below for # why that value can only have been produced by the worker at phase 3. # +# The whole flow above runs TWICE, back to back: once with broker +# authentication left off (the default an operator gets with no nats.auth +# configuration at all — this run is untouched by the second), and once with +# it turned on and the worker enrolling itself via a join token before it can +# connect. The auth-off run always goes first, so a failure immediately says +# which mode broke: a failure before the "MODE 2/2" banner is the default +# path regressing, which is the more serious of the two by a wide margin. +# # Usage: -# bash scripts/smoke.sh # builds binaries if missing, runs the flow +# bash scripts/smoke.sh # builds binaries if missing, runs both modes # make smoke # same, via the Makefile # # Environment overrides: @@ -23,8 +31,8 @@ # SQI_WORKER_BIN path to a prebuilt sqi-worker (default: /bin/sqi-worker) # SQI_SMOKE_PYTHON python interpreter for the WS check (auto-detected otherwise) # -# Exit status: 0 only if every assertion passed; non-zero with a clear message -# (and the relevant server/worker log tail) otherwise. +# Exit status: 0 only if every assertion passed in both modes; non-zero with a +# clear message (and the relevant server/worker log tail) otherwise. set -euo pipefail @@ -156,6 +164,22 @@ else [ -x "$WORKER_BIN" ] || fail "worker binary not found after build: $WORKER_BIN" fi +# ── The smoke flow, run once per broker-auth mode ───────────────────────────── +# +# run_smoke_flow MODE boots a fresh server+worker pair and drives the whole +# assertion set described at the top of this file against it. MODE is +# "noauth" (broker authentication left off — every env block below behaves +# exactly as this script always has) or "brokerauth" (nats.auth.enabled=true, +# with a join token minted before the server starts and handed to the worker +# instead of nothing). +# +# Called as ( run_smoke_flow MODE ) — in a subshell — once per mode, from the +# bottom of this file. That gives each run its own temp workspace, ports, +# PIDs, and EXIT/INT/TERM trap, and means this function's own "exit 0" on +# success only ends that one subshell, not the whole script. +run_smoke_flow() { + local mode="$1" + # ── Temp workspace + teardown trap ──────────────────────────────────────────── TMP_DIR="$(mktemp -d "${TMPDIR:-/tmp}/sqi-smoke.XXXXXX")" @@ -203,15 +227,48 @@ HTTP_ADDR="127.0.0.1:${HTTP_PORT}" NATS_ADDR="127.0.0.1:${NATS_PORT}" BASE_URL="http://${HTTP_ADDR}" -log "starting sqi-server (http=${HTTP_ADDR}, nats=${NATS_ADDR})" -SQI_HTTP_ADDR="$HTTP_ADDR" \ -SQI_NATS_ADDR="$NATS_ADDR" \ -SQI_NATS_DATA_DIR="${TMP_DIR}/nats" \ -SQI_STORE_SQLITE_PATH="${TMP_DIR}/sqi.db" \ -SQI_DISCOVERY_ENABLED="false" \ -SQI_SCHEDULER_TICK_INTERVAL="100ms" \ -SQI_LOG_LEVEL="warn" \ - "$SERVER_BIN" serve >"$SERVER_LOG" 2>&1 & +# In broker-auth mode, mint a join token BEFORE the server starts — the CLI +# operates directly on the SQLite file (no running server required), exactly +# the offline path an operator would use. The worker gets the raw token +# through a file, never an argv or an env var visible in a process listing. +JOIN_TOKEN_FILE="" +if [ "$mode" = "brokerauth" ]; then + log "creating the SQLite database (worker subcommands never create one themselves)" + "$SERVER_BIN" migrate up --db "${TMP_DIR}/sqi.db" >/dev/null || fail "sqi-server migrate up failed" + + # Resolve via SQI_STORE_SQLITE_PATH rather than --db here, so this run + # exercises the config-layer resolution path (the one an operator who set + # up their deployment through the environment actually uses) rather than + # only the explicit-flag path already covered above. + log "minting a worker join token (broker auth mode)" + JOIN_TOKEN="$(SQI_STORE_SQLITE_PATH="${TMP_DIR}/sqi.db" "$SERVER_BIN" worker token issue --name smoke-brokerauth)" + [ -n "$JOIN_TOKEN" ] || fail "sqi-server worker token issue produced no token" + JOIN_TOKEN_FILE="${TMP_DIR}/join-token" + printf '%s' "$JOIN_TOKEN" >"$JOIN_TOKEN_FILE" + chmod 600 "$JOIN_TOKEN_FILE" +fi + +log "starting sqi-server (http=${HTTP_ADDR}, nats=${NATS_ADDR}, mode=${mode})" +if [ "$mode" = "brokerauth" ]; then + SQI_HTTP_ADDR="$HTTP_ADDR" \ + SQI_NATS_ADDR="$NATS_ADDR" \ + SQI_NATS_DATA_DIR="${TMP_DIR}/nats" \ + SQI_STORE_SQLITE_PATH="${TMP_DIR}/sqi.db" \ + SQI_DISCOVERY_ENABLED="false" \ + SQI_SCHEDULER_TICK_INTERVAL="100ms" \ + SQI_LOG_LEVEL="warn" \ + SQI_NATS_AUTH_ENABLED="true" \ + "$SERVER_BIN" serve >"$SERVER_LOG" 2>&1 & +else + SQI_HTTP_ADDR="$HTTP_ADDR" \ + SQI_NATS_ADDR="$NATS_ADDR" \ + SQI_NATS_DATA_DIR="${TMP_DIR}/nats" \ + SQI_STORE_SQLITE_PATH="${TMP_DIR}/sqi.db" \ + SQI_DISCOVERY_ENABLED="false" \ + SQI_SCHEDULER_TICK_INTERVAL="100ms" \ + SQI_LOG_LEVEL="warn" \ + "$SERVER_BIN" serve >"$SERVER_LOG" 2>&1 & +fi SERVER_PID=$! # Poll /readyz until 200 (bounded). Fail fast if the process exits early. @@ -249,19 +306,39 @@ log "created queue ${QUEUE_ID}" # ── Start the worker ────────────────────────────────────────────────────────── -log "starting sqi-worker (nats=nats://${NATS_ADDR}, farm=${FARM_ID}, queue=${QUEUE_ID})" -SQI_WORKER_NATS_URL="nats://${NATS_ADDR}" \ -SQI_WORKER_DISCOVERY_ENABLE_MDNS="false" \ -SQI_WORKER_FARM_ID="$FARM_ID" \ -SQI_WORKER_QUEUE_IDS="$QUEUE_ID" \ -SQI_WORKER_DATA_DIR="${TMP_DIR}/worker-data" \ -SQI_WORKER_ALLOW_ROOT="true" \ -SQI_WORKER_LOG_LEVEL="warn" \ -SQI_WORKER_LOG_FORMAT="text" \ -SQI_WORKER_HEARTBEAT_INTERVAL="1s" \ -SQI_WORKER_PULL_IDLE_BACKOFF="300ms" \ -SQI_WORKER_METRICS_ADDR="127.0.0.1:$(free_port)" \ - "$WORKER_BIN" start >"$WORKER_LOG" 2>&1 & +log "starting sqi-worker (nats=nats://${NATS_ADDR}, farm=${FARM_ID}, queue=${QUEUE_ID}, mode=${mode})" +if [ "$mode" = "brokerauth" ]; then + # No SQI_WORKER_NATS_CREDENTIAL_FILE: it defaults under + # SQI_WORKER_DATA_DIR, which is what makes this worker's enrolled + # credential land in its own fresh, per-run data directory. + SQI_WORKER_NATS_URL="nats://${NATS_ADDR}" \ + SQI_WORKER_DISCOVERY_ENABLE_MDNS="false" \ + SQI_WORKER_FARM_ID="$FARM_ID" \ + SQI_WORKER_QUEUE_IDS="$QUEUE_ID" \ + SQI_WORKER_DATA_DIR="${TMP_DIR}/worker-data" \ + SQI_WORKER_ALLOW_ROOT="true" \ + SQI_WORKER_LOG_LEVEL="warn" \ + SQI_WORKER_LOG_FORMAT="text" \ + SQI_WORKER_HEARTBEAT_INTERVAL="1s" \ + SQI_WORKER_PULL_IDLE_BACKOFF="300ms" \ + SQI_WORKER_METRICS_ADDR="127.0.0.1:$(free_port)" \ + SQI_WORKER_NATS_JOIN_TOKEN_FILE="$JOIN_TOKEN_FILE" \ + SQI_WORKER_NATS_SERVER_URL="$BASE_URL" \ + "$WORKER_BIN" start >"$WORKER_LOG" 2>&1 & +else + SQI_WORKER_NATS_URL="nats://${NATS_ADDR}" \ + SQI_WORKER_DISCOVERY_ENABLE_MDNS="false" \ + SQI_WORKER_FARM_ID="$FARM_ID" \ + SQI_WORKER_QUEUE_IDS="$QUEUE_ID" \ + SQI_WORKER_DATA_DIR="${TMP_DIR}/worker-data" \ + SQI_WORKER_ALLOW_ROOT="true" \ + SQI_WORKER_LOG_LEVEL="warn" \ + SQI_WORKER_LOG_FORMAT="text" \ + SQI_WORKER_HEARTBEAT_INTERVAL="1s" \ + SQI_WORKER_PULL_IDLE_BACKOFF="300ms" \ + SQI_WORKER_METRICS_ADDR="127.0.0.1:$(free_port)" \ + "$WORKER_BIN" start >"$WORKER_LOG" 2>&1 & +fi WORKER_PID=$! # Poll GET /api/v1/workers until our worker is online. @@ -570,7 +647,7 @@ log "EXPR assertion PASSED: phase-3 resolved text found in task logs" # ── Summary ─────────────────────────────────────────────────────────────────── log "==============================================" -log "SMOKE TEST PASSED" +log "SMOKE TEST PASSED (mode=${mode})" log " REST log assertion: PASSED" case "$WS_OK" in pass) log " WS log assertion: PASSED" ;; @@ -579,3 +656,25 @@ esac log " EXPR phase-3 assertion: PASSED" log "==============================================" exit 0 +} + +# ── Run both modes ───────────────────────────────────────────────────────────── +# +# Each call runs in its own subshell so run_smoke_flow's internal "exit 0" +# ends only that run, and its trap, PIDs, and temp workspace never leak into +# the other. Auth-off goes first, unconditionally: if it fails, mode 2 never +# starts, and the last banner printed is the one that broke. + +log "==================================================================" +log "MODE 1/2: broker auth OFF -- the default path, unmodified" +log "==================================================================" +( run_smoke_flow noauth ) + +log "==================================================================" +log "MODE 2/2: broker auth ON, worker enrolled via a join token" +log "==================================================================" +( run_smoke_flow brokerauth ) + +log "==================================================================" +log "SMOKE TEST PASSED IN BOTH MODES" +log "==================================================================" diff --git a/test/integration/broker_auth_test.go b/test/integration/broker_auth_test.go new file mode 100644 index 00000000..ce79ad61 --- /dev/null +++ b/test/integration/broker_auth_test.go @@ -0,0 +1,595 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +//go:build integration + +package integration + +// TestRevocation_DisconnectsAndReclaims proves revocation end to end: DELETE +// /api/v1/workers/{id}/credential revokes in the store, then reloads the +// running broker's authorized-key set, which disconnects the revoked +// worker's live NATS connection SYNCHRONOUSLY — inside the reload call, +// because nats-server's reloadAuthorization re-runs isClientAuthorized over +// every connected client and calls authViolation() on any that no longer +// pass. The disconnected worker's in-flight task is then returned to ready +// by the EXISTING heartbeat-sweep/reclaim path (internal/scheduler), not by +// anything this test reimplements, and a second, unrevoked +// worker is left completely unaffected and able to lease the reclaimed +// task. +// +// This is deliberately distinct from "sqi-server worker revoke" (the CLI), +// which writes the same store row from a separate process holding no broker +// handle and only takes effect at the running server's next start — that +// path is covered by cmd/sqi-server's own tests, not here. + +import ( + "context" + "encoding/json" + "fmt" + "log/slog" + "net/http" + "os" + "testing" + "time" + + nats "github.com/nats-io/nats.go" + "github.com/nats-io/nats.go/jetstream" + "github.com/nats-io/nkeys" + + "github.com/uberware/sqi/internal/auth/jointoken" + "github.com/uberware/sqi/internal/brokerauth" + "github.com/uberware/sqi/internal/config" + "github.com/uberware/sqi/internal/scheduler" + "github.com/uberware/sqi/internal/server" + "github.com/uberware/sqi/internal/store" + "github.com/uberware/sqi/internal/store/sqlite" +) + +// ── Server boot (broker auth on, auth off) ────────────────────────────────── + +// startBrokerAuthServer boots a full sqi-server with broker authentication +// enabled and a short heartbeat-sweep timing so a revoked worker's task is +// reclaimed within seconds rather than this suite's default 30s. Auth +// (session/API-key) is left off: DELETE /api/v1/workers/{id}/credential is +// mounted unconditionally regardless of AuthEnabled (the anonymous +// superuser principal is granted it, same as most permission-gated routes), +// so nothing about the revocation path this test exercises needs a login +// flow. +// +// sqlitePath is the caller's choice, not a generated temp path, so a caller +// can open its own store.Store on the same file and seed rows (e.g. an +// enrolled worker credential) BEFORE calling this — see +// seedWorkerCredential. The store's write pool is a single connection; +// seeding must complete and close its handle before this server opens its +// own, not run concurrently with it. +func startBrokerAuthServer(t *testing.T, sqlitePath string, mutate func(*server.Config)) *testServer { + t.Helper() + + httpAddr := fmt.Sprintf("127.0.0.1:%d", freePort(t)) + natsAddr := fmt.Sprintf("127.0.0.1:%d", freePort(t)) + tmpDir := t.TempDir() + + cfg := server.Config{ + HTTPAddr: httpAddr, + CORSOrigins: []string{"*"}, + NATSAddr: natsAddr, + NATSDataDir: tmpDir + "/nats", + NATSMaxStoreMB: 64, + SQLitePath: sqlitePath, + CheckpointInterval: time.Minute, + Scheduler: scheduler.Config{ + AssignInterval: 100 * time.Millisecond, + AssignBatchSize: 10, + AssignWorkers: 2, + // Short enough that the heartbeat-sweep reclaim this test + // verifies against completes in a few seconds, not this + // suite's usual 30s/15s (tuned for tests that never exercise + // offline detection at all). + WorkerTimeout: 2 * time.Second, + HeartbeatSweepInterval: 300 * time.Millisecond, + }, + DiscoveryEnabled: false, + + // The enrollment endpoint is off by default: most callers enroll by + // seeding the store directly before boot (see seedWorkerCredential) + // rather than through POST /workers/enroll. A caller that needs the + // real REST enrollment surface turns it on via mutate — see + // TestEnrollment_ConnectsToRunningBrokerWithoutRestart. + NATSAuthEnabled: true, + } + if mutate != nil { + mutate(&cfg) + } + + logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelWarn})) + srv := server.New(cfg, logger, nil) + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { done <- srv.Run(ctx) }() + + select { + case err := <-done: + cancel() + t.Fatalf("startBrokerAuthServer: server exited during startup: %v", err) + case <-time.After(200 * time.Millisecond): + } + + ts := &testServer{HTTPAddr: httpAddr, NATSAddr: natsAddr, cancel: cancel, done: done} + t.Cleanup(func() { + cancel() + select { + case <-done: + case <-time.After(15 * time.Second): + t.Logf("warning: server did not stop within 15s after context cancel") + } + }) + + if !waitForTCP(t, httpAddr, 10*time.Second) { + t.Fatal("startBrokerAuthServer: HTTP server did not start listening") + } + if !waitForReadyz(t, httpAddr, 10*time.Second) { + t.Fatal("startBrokerAuthServer: server did not become ready") + } + return ts +} + +// seedWorkerCredential inserts an already-enrolled [store.WorkerCredential] +// row directly into the SQLite database at dbPath, mirroring what +// "sqi-server worker enroll" (the offline CLI enrollment path) writes — a +// direct store.CreateWorkerCredential call from a process holding no broker +// handle, so (unlike POST /workers/enroll, see enrollWorker below) it never +// reaches a running broker's authorized-key set on its own. Used by +// TestRevocation_DisconnectsAndReclaims, which wants both its workers +// connectable from the moment the server boots and has no other reason to +// exercise the REST enrollment surface. Must run before the server opens +// its own connection to the same file — see startBrokerAuthServer. +func seedWorkerCredential(t *testing.T, dbPath, workerID, publicKey string) { + t.Helper() + ctx := context.Background() + st, err := sqlite.Open(ctx, dbPath, sqlite.DefaultOptions()) + if err != nil { + t.Fatalf("seedWorkerCredential: sqlite.Open: %v", err) + } + defer func() { _ = st.Close() }() + + if _, err := st.CreateWorkerCredential(ctx, store.WorkerCredential{ + ID: workerID + "-cred", + WorkerID: workerID, + PublicKey: publicKey, + EnrolledAt: time.Now().UTC(), + }); err != nil { + t.Fatalf("seedWorkerCredential: CreateWorkerCredential: %v", err) + } +} + +// seedJoinToken inserts a join token row directly into the SQLite database +// at dbPath and returns the raw token, bypassing POST /workers/join-tokens +// (which requires session auth this suite does not otherwise need — minting +// is not what either enrollment or revocation testing is about here). Must +// run before the server opens its own connection to the same file, same +// rule as seedWorkerCredential. +func seedJoinToken(t *testing.T, dbPath, name string) string { + t.Helper() + raw, hash, prefix, err := jointoken.Generate() + if err != nil { + t.Fatalf("seedJoinToken: jointoken.Generate: %v", err) + } + + ctx := context.Background() + st, err := sqlite.Open(ctx, dbPath, sqlite.DefaultOptions()) + if err != nil { + t.Fatalf("seedJoinToken: sqlite.Open: %v", err) + } + defer func() { _ = st.Close() }() + + now := time.Now().UTC() + if _, err := st.CreateWorkerJoinToken(ctx, store.WorkerJoinToken{ + ID: name + "-token", + TokenHash: hash, + Prefix: prefix, + Name: name, + ExpiresAt: now.Add(time.Hour), + CreatedAt: now, + }); err != nil { + t.Fatalf("seedJoinToken: CreateWorkerJoinToken: %v", err) + } + return raw +} + +// ── REST helpers for the synchronous enroll and revoke paths under test ──── + +// workerCredentialWireResp is the subset of POST /workers/enroll's response +// this suite needs (mirrors internal/api/workerenroll.go's +// workerCredentialResponse). +type workerCredentialWireResp struct { + ID string `json:"id"` + WorkerID string `json:"worker_id"` +} + +// enrollWorker exchanges joinToken for a broker credential over the real, +// unauthenticated POST /api/v1/workers/enroll wire protocol — the path that +// reloads a RUNNING broker's authorized-key set rather than only taking +// effect at the server's next start. +func enrollWorker(t *testing.T, ts *testServer, joinToken, workerID, publicKey string) { + t.Helper() + body, err := json.Marshal(map[string]string{ + "join_token": joinToken, + "worker_id": workerID, + "public_key": publicKey, + }) + if err != nil { + t.Fatalf("enrollWorker: marshal: %v", err) + } + var resp workerCredentialWireResp + mustDoJSON(t, http.MethodPost, apiURL(ts, "/api/v1/workers/enroll"), body, "application/json", http.StatusCreated, &resp) + if resp.WorkerID != workerID { + t.Fatalf("enrollWorker: response worker_id = %q, want %q", resp.WorkerID, workerID) + } +} + +// revokeWorkerCredential calls the synchronous revocation endpoint under +// test: DELETE /api/v1/workers/{id}/credential. +func revokeWorkerCredential(t *testing.T, ts *testServer, workerID string) { + t.Helper() + mustDoJSON(t, http.MethodDelete, apiURL(ts, "/api/v1/workers/"+workerID+"/credential"), nil, "", http.StatusNoContent, nil) +} + +// pollTaskStatus polls GET /api/v1/tasks/{id} (via getTaskDetail, defined in +// retry_test.go) until its status is one of targets, or timeout elapses. +func pollTaskStatus(t *testing.T, ts *testServer, taskID string, targets []string, timeout time.Duration) string { + t.Helper() + targetSet := make(map[string]bool, len(targets)) + for _, s := range targets { + targetSet[s] = true + } + var last string + deadline := time.Now().Add(timeout) + for time.Now().Before(deadline) { + last = getTaskDetail(t, ts, taskID).Status + if targetSet[last] { + return last + } + time.Sleep(100 * time.Millisecond) + } + t.Fatalf("pollTaskStatus: task %s stuck at %q; wanted one of %v after %s", taskID, last, targets, timeout) + return "" +} + +// ── nkey-authenticated mock worker ────────────────────────────────────────── + +// newMockWorkerWithNkey connects to natsURL as the given enrolled nkey +// credential and returns a *mockWorker built around that connection. +// extraOpts is appended after the nkey and NoReconnect options — NoReconnect +// so a revoked worker's disconnect is observed directly rather than masked +// by nats.go retrying (and getting the same auth error) for some time first, +// mirroring internal/bus's own revocation tests. +// +// The per-worker inbox prefix is not optional decoration: an enrolled +// worker is granted "_INBOX_.>" and nothing wider, so a connection left +// on nats.go's process-global "_INBOX" cannot even subscribe to its own +// reply inbox — every JetStream publish then fails waiting for a PubAck it +// is not allowed to hear. internal/worker/natsclient sets the same option +// on every real worker connection. +func newMockWorkerWithNkey(t *testing.T, natsURL, workerID, farmID, queueID string, seed []byte, pub string, extraOpts ...nats.Option) *mockWorker { + t.Helper() + + opts := append([]nats.Option{ + nats.Nkey(pub, func(nonce []byte) ([]byte, error) { + kp, err := nkeys.FromSeed(seed) + if err != nil { + return nil, err + } + return kp.Sign(nonce) + }), + nats.CustomInboxPrefix(brokerauth.InboxPrefix(workerID)), + nats.NoReconnect(), + }, extraOpts...) + + nc, err := nats.Connect(natsURL, opts...) + if err != nil { + t.Fatalf("newMockWorkerWithNkey(%s): Connect: %v", workerID, err) + } + + js, err := jetstream.New(nc) + if err != nil { + nc.Close() + t.Fatalf("newMockWorkerWithNkey(%s): jetstream.New: %v", workerID, err) + } + + t.Cleanup(func() { + if !nc.IsClosed() { + nc.Close() + } + }) + + return &mockWorker{t: t, id: workerID, farmID: farmID, queueID: queueID, nc: nc, js: js} +} + +// ── The test ───────────────────────────────────────────────────────────────── + +func TestRevocation_DisconnectsAndReclaims(t *testing.T) { + sqlitePath := t.TempDir() + "/sqi-broker-auth.db" + + seedA, pubA, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed(A): %v", err) + } + seedB, pubB, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed(B): %v", err) + } + + // Both workers are enrolled BEFORE the server starts, so the broker's + // initial authorized-key set (built once, at Start, from + // ListActiveWorkerCredentials) includes both — see seedWorkerCredential's + // doc comment for why this test does not use POST /workers/enroll. + seedWorkerCredential(t, sqlitePath, "worker-a", pubA) + seedWorkerCredential(t, sqlitePath, "worker-b", pubB) + + ts := startBrokerAuthServer(t, sqlitePath, nil) + + farmID, queueID := seedFarmAndQueue(t, ts) + + natsURL := "nats://" + ts.NATSAddr + + closedA := make(chan struct{}) + workerA := newMockWorkerWithNkey(t, natsURL, "worker-a", farmID, queueID, seedA, pubA, + nats.ClosedHandler(func(*nats.Conn) { close(closedA) })) + workerB := newMockWorkerWithNkey(t, natsURL, "worker-b", farmID, queueID, seedB, pubB) + + workerA.register() + workerB.register() + // Heartbeat well under the 2s WorkerTimeout configured above. Worker A's + // heartbeat loop stops on its own the moment revocation closes its NATS + // connection (the publish fails and the loop returns) — nothing here + // needs to stop it explicitly. + workerA.startHeartbeat(200 * time.Millisecond) + workerB.startHeartbeat(200 * time.Millisecond) + + pollWorkerOnline(t, ts, "worker-a", 5*time.Second) + pollWorkerOnline(t, ts, "worker-b", 5*time.Second) + + jobID := submitJob(t, ts, farmID, queueID) + assign := workerA.pullAssignment(15 * time.Second) + if assign.JobID != jobID { + t.Fatalf("assignment job ID: got %q, want %q", assign.JobID, jobID) + } + taskID := assign.TaskID + + workerA.publishStatus(assign, "running", nil) + pollTaskStatus(t, ts, taskID, []string{"running"}, 5*time.Second) + + // ── Revoke worker A: the synchronous path under test ──────────────────── + revokeWorkerCredential(t, ts, "worker-a") + + // 1. A's NATS connection closes. Revocation is synchronous — nats-server + // re-authorizes every connected client inside the broker's + // ReloadOptions call, which DELETE /workers/{id}/credential's handler + // waits on — so this is not a poll for a generous timeout, only + // tolerance for scheduling jitter in observing an event that already + // happened. + select { + case <-closedA: + case <-time.After(2 * time.Second): + t.Fatal("worker A's NATS connection was not closed by revocation") + } + + // 2. A's task returns to ready via a legal store.ValidateTaskTransition + // arrow (running -> ready), applied by the EXISTING heartbeat-sweep and + // reclaim path once A's missed heartbeats exceed WorkerTimeout — not by + // anything this test or RevokeWorker itself does directly. + pollTaskStatus(t, ts, taskID, []string{"ready"}, 8*time.Second) + + // 3. B is unaffected and can still lease — proven by actually leasing + // the reclaimed task, which simultaneously confirms the reclaim resulted + // in a real reassignment rather than a status stuck at "ready". + assignB := workerB.pullAssignment(10 * time.Second) + if assignB.TaskID != taskID { + t.Fatalf("worker B's assignment: got task %q, want the reclaimed task %q", assignB.TaskID, taskID) + } + if !workerB.nc.IsConnected() { + t.Fatal("worker B's connection was disturbed by A's revocation") + } +} + +// pollWorkerOffline polls GET /api/v1/workers until workerID is visible with +// status "offline", or timeout elapses. Used to wait for the heartbeat-sweep +// to notice a worker that has stopped heartbeating, without depending on the +// sweep's exact timing. +func pollWorkerOffline(t *testing.T, ts *testServer, workerID string, timeout time.Duration) { + t.Helper() + deadline := time.Now().Add(timeout) + for time.Now().Before(deadline) { + var resp workerListResp + mustDoJSON(t, http.MethodGet, apiURL(ts, "/api/v1/workers"), nil, "", http.StatusOK, &resp) + for _, w := range resp.Items { + if w.ID == workerID && w.Status == "offline" { + return + } + } + time.Sleep(200 * time.Millisecond) + } + t.Fatalf("pollWorkerOffline: worker %s did not go offline within %s", workerID, timeout) +} + +// TestWorkerDeletion_RevokesCredentialAndDisconnects proves the cascade DELETE +// /api/v1/workers/{id} is now expected to perform: removing a worker record +// also revokes its broker credential, through the SAME synchronous +// store-write-then-broker-reload path DELETE /workers/{id}/credential uses. +// Without it, a machine an operator has just decommissioned from the farm +// keeps live broker access — able to connect, lease work and execute job +// code — because WorkersManage (what deleting a worker requires) does not +// imply WorkersEnroll (what revoking a credential directly requires). +// +// worker-a never heartbeats after registering, so the heartbeat sweep marks +// it offline on its own — the ONLY status DELETE /workers/{id} accepts +// without an extra disable step — while its NATS connection stays live +// (nothing here closes it). The test then deletes the worker over REST and +// asserts the still-open connection is closed by the same +// ReloadCredentials call TestRevocation_DisconnectsAndReclaims already +// proves is synchronous for the dedicated credential-revoke endpoint. +func TestWorkerDeletion_RevokesCredentialAndDisconnects(t *testing.T) { + sqlitePath := t.TempDir() + "/sqi-broker-auth-delete.db" + + seedA, pubA, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed(A): %v", err) + } + seedWorkerCredential(t, sqlitePath, "worker-a", pubA) + + ts := startBrokerAuthServer(t, sqlitePath, nil) + farmID, queueID := seedFarmAndQueue(t, ts) + natsURL := "nats://" + ts.NATSAddr + + closedA := make(chan struct{}) + workerA := newMockWorkerWithNkey(t, natsURL, "worker-a", farmID, queueID, seedA, pubA, + nats.ClosedHandler(func(*nats.Conn) { close(closedA) })) + workerA.register() + // Deliberately no startHeartbeat: the worker must go offline in the + // store (the only status DELETE /workers/{id} accepts here) while its + // NATS connection stays open, so the disconnect this test asserts on + // can only be explained by the revoke-on-delete cascade, not by the + // worker's own connection dying of neglect. + pollWorkerOffline(t, ts, "worker-a", 5*time.Second) + + mustDoJSON(t, http.MethodDelete, apiURL(ts, "/api/v1/workers/worker-a"), nil, "", http.StatusNoContent, nil) + + // The credential revoke inside the delete handler must have disconnected + // worker A's still-open NATS connection, synchronously. + select { + case <-closedA: + case <-time.After(2 * time.Second): + t.Fatal("worker A's NATS connection was not closed by deleting its worker record") + } + + // The worker row itself is gone. + mustDoJSON(t, http.MethodGet, apiURL(ts, "/api/v1/workers/worker-a"), nil, "", http.StatusNotFound, nil) +} + +// TestEnrollment_ConnectsToRunningBrokerWithoutRestart guards against the +// broker's authorized-key set ever again going unreloaded after POST +// /workers/enroll creates a credential. loadBrokerAuthConfig only ever runs +// once, at Start, so without an explicit reload a worker enrolled against a +// RUNNING server could not actually connect: nats-server would refuse it +// with "Authorization Violation", and the real sqi-worker binary exits +// fatally naming that rejection. This enrolls a worker over the real REST +// wire protocol AFTER the server is already up, with no restart in between, +// and asserts it connects and registers successfully. +func TestEnrollment_ConnectsToRunningBrokerWithoutRestart(t *testing.T) { + sqlitePath := t.TempDir() + "/sqi-broker-auth-enroll.db" + + rawToken := seedJoinToken(t, sqlitePath, "worker-c") + + ts := startBrokerAuthServer(t, sqlitePath, func(cfg *server.Config) { + cfg.NATSAuthEnrollmentEndpointEnabled = true + }) + + farmID, queueID := seedFarmAndQueue(t, ts) + + seedC, pubC, err := brokerauth.GenerateSeed() + if err != nil { + t.Fatalf("GenerateSeed(C): %v", err) + } + + // Enroll AFTER the server has already booted and become ready — the + // broker's initial authorized-key set (built once, at Start) could not + // possibly contain this credential. + enrollWorker(t, ts, rawToken, "worker-c", pubC) + + natsURL := "nats://" + ts.NATSAddr + workerC := newMockWorkerWithNkey(t, natsURL, "worker-c", farmID, queueID, seedC, pubC) + workerC.register() + workerC.startHeartbeat(200 * time.Millisecond) + + // pollWorkerOnline itself has no generous fixed sleep baked in beyond its + // own timeout — a rejected connection here would mean register()'s + // underlying nats.Connect (inside newMockWorkerWithNkey) already failed + // the test outright with "Authorization Violation", so reaching this + // point at all already proves the enrolled worker could connect. + pollWorkerOnline(t, ts, "worker-c", 5*time.Second) +} + +// ── The default-config regression ─────────────────────────────────────────── + +// listActiveWorkerCredentials opens a second, independent connection to the +// SQLite database a running server was started against and calls +// ListActiveWorkerCredentials. This is safe to run WHILE that server is +// still up: the store opens every database in WAL mode, which allows a +// reader to run alongside the server's own single-connection writer. +// AutoMigrate is left off — the server that owns dbPath has already applied +// every migration, and re-running goose's version check here would be pure +// overhead against a live file for no benefit. +func listActiveWorkerCredentials(t *testing.T, dbPath string) []store.WorkerCredential { + t.Helper() + ctx := context.Background() + st, err := sqlite.Open(ctx, dbPath, sqlite.Options{AutoMigrate: false}) + if err != nil { + t.Fatalf("listActiveWorkerCredentials: sqlite.Open: %v", err) + } + defer func() { _ = st.Close() }() + + creds, err := st.ListActiveWorkerCredentials(ctx) + if err != nil { + t.Fatalf("listActiveWorkerCredentials: %v", err) + } + return creds +} + +// TestDefaultConfig_NoBrokerAuth is the load-bearing regression for the +// whole component: a server and worker started with NO broker-auth +// configuration must behave exactly as they did before broker +// authentication existed. Every other test in this file proves the new +// capability works; this one proves it costs nothing when unused. +// +// If this test fails, the default path has regressed. Fix the cause — never +// adjust the test to accommodate the regression. Every operator who has +// never heard of broker authentication meets this exact path on day one. +// +// This is one half of a two-part proof. The other half — that +// internal/config's own zero-configuration default carries nats.auth.enabled +// = false through to the value the server actually runs with — is pinned at +// the unit level by TestServerConfig_DefaultsAreTheConfigDefaults and +// TestServerConfig_CarriesTheBrokerAuthSettings in cmd/sqi-server, which +// exercise the exact config.Config -> server.Config mapping function the +// serve subcommand uses. This test starts from that already-proven +// conclusion (asserted below as a guard, not re-derived) and proves the +// RUNTIME behavior it implies: a real sqi-worker subprocess with no +// credential file and no join token registers, is leased a task, runs it to +// completion, and the worker-credential table -- which only broker auth +// ever writes to -- stays empty throughout. +func TestDefaultConfig_NoBrokerAuth(t *testing.T) { + // Guard the claim the rest of this test depends on: an operator's own + // zero-configuration default leaves broker authentication off. + if config.DefaultConfig().NATS.Auth.Enabled { + t.Fatal("config.DefaultConfig().NATS.Auth.Enabled = true, want false -- " + + "broker authentication must default to off") + } + + // startServer boots server.Config with every NATSAuth* field left at its + // zero value: the same "nobody has ever configured this" state an + // operator gets with no nats.auth section in their config file, no + // SQI_NATS_AUTH_* environment variable, and no --nats-auth-* flag -- + // matching the config default asserted above. + ts := startServer(t) + farmID, queueID := seedFarmAndQueue(t, ts) + + // A real sqi-worker subprocess. startRealWorker sets no join-token, + // credential-file, or server-url environment variable at all -- exactly + // as every worker invocation looked before broker auth existed -- and + // blocks until the worker is visible online, which is this test's proof + // of criterion 1: the worker registers with no credential. + startRealWorker(t, ts, farmID, queueID) + + // Criterion 2: it is leased a task and runs it to completion. + jobID := submitJob(t, ts, farmID, queueID) + if got := pollJobStatus(t, ts, jobID, []string{"completed", "failed", "canceled"}, 30*time.Second); got != "completed" { + t.Fatalf("job on the default no-broker-auth path ended %q, want completed", got) + } + + // Criterion 3: nothing enrolled anything. Broker auth was never engaged + // on this path, so no row should ever have been written to the + // worker-credential table. + if creds := listActiveWorkerCredentials(t, ts.DBPath); len(creds) != 0 { + t.Errorf("default path created %d worker credentials, want 0: %+v", len(creds), creds) + } +} diff --git a/test/integration/failure_reason_test.go b/test/integration/failure_reason_test.go index eb07d329..b04a95ce 100644 --- a/test/integration/failure_reason_test.go +++ b/test/integration/failure_reason_test.go @@ -90,7 +90,7 @@ func getTaskAttempts(t *testing.T, ts *testServer, taskID string) taskAttemptsRe } // publishStatusWithMessage publishes a [protocol.TaskStatusMsg] to -// task.status. carrying a Message, exercising the same field a real +// task.status.. carrying a Message, exercising the same field a real // worker's failure report populates (see internal/worker's failPreExec / // process-exit reporting). The harness's publishStatus helper (used by the // auto-retry tests) never sets Message, so this test publishes directly @@ -117,7 +117,7 @@ func publishStatusWithMessage(t *testing.T, w *mockWorker, assign protocol.Assig ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - if _, err := w.js.Publish(ctx, bus.TaskStatusSubject(assign.JobID), data); err != nil { + if _, err := w.js.Publish(ctx, bus.TaskStatusSubject(w.id, assign.JobID), data); err != nil { t.Fatalf("publishStatusWithMessage(%s): publish: %v", status, err) } } diff --git a/test/integration/harness_test.go b/test/integration/harness_test.go index a9c727bb..8af7eb93 100644 --- a/test/integration/harness_test.go +++ b/test/integration/harness_test.go @@ -75,6 +75,13 @@ type testServer struct { // NATSAddr is the full "host:port" address the embedded NATS server // is listening on. Workers connect to "nats://" + NATSAddr. NATSAddr string + // DBPath is the SQLite file this server was started against. Set by + // constructors that know it, empty otherwise. A caller that needs to + // inspect store state the REST API does not expose (e.g. confirming no + // row was ever written to a table) can open a second, independent + // connection to this same path — safe because the store runs in WAL + // mode, which allows a reader to run alongside the server's own writer. + DBPath string cancel context.CancelFunc done chan error @@ -160,6 +167,7 @@ func startServer(t *testing.T) *testServer { ts := &testServer{ HTTPAddr: httpAddr, NATSAddr: natsAddr, + DBPath: sqlitePath, cancel: cancel, done: done, } @@ -269,7 +277,7 @@ func newMockWorker(t *testing.T, natsURL, workerID, farmID, queueID string) *moc } } -// register publishes a [protocol.RegisterMsg] to worker.register so the +// register publishes a [protocol.RegisterMsg] to worker.register. so the // server records this worker as online and eligible for task assignment. func (w *mockWorker) register() { w.t.Helper() @@ -293,7 +301,7 @@ func (w *mockWorker) register() { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - if _, err := w.js.Publish(ctx, bus.SubjectWorkerRegister, data); err != nil { + if _, err := w.js.Publish(ctx, bus.WorkerRegisterSubject(w.id), data); err != nil { w.t.Fatalf("mockWorker.register: publish: %v", err) } } @@ -329,7 +337,7 @@ func (w *mockWorker) startHeartbeat(interval time.Duration) { pubCtx, pubCancel := context.WithTimeout(ctx, 2*time.Second) // Ignore publish errors in the heartbeat loop — the connection // may be closing during test cleanup. - if _, pubErr := w.js.Publish(pubCtx, bus.SubjectWorkerHeartbeat, data); pubErr != nil { + if _, pubErr := w.js.Publish(pubCtx, bus.WorkerHeartbeatSubject(w.id), data); pubErr != nil { pubCancel() return } @@ -339,7 +347,7 @@ func (w *mockWorker) startHeartbeat(interval time.Duration) { }() } -// pullAssignment requests a work lease on work.lease. and blocks until +// pullAssignment requests a work lease on work.lease.. and blocks until // the server returns a non-empty batch or timeout expires. It returns the // first decoded [protocol.AssignMsg] in the reply. The request timeout exceeds // the server's long-poll hold so a parked request is never abandoned (which @@ -358,7 +366,7 @@ func (w *mockWorker) pullAssignment(timeout time.Duration) protocol.AssignMsg { for time.Now().Before(deadline) { reqTimeout := min(time.Until(deadline), 35*time.Second) reqCtx, cancel := context.WithTimeout(context.Background(), reqTimeout) - msg, reqErr := w.nc.RequestWithContext(reqCtx, bus.WorkLeaseSubject(w.queueID), reqBytes) + msg, reqErr := w.nc.RequestWithContext(reqCtx, bus.WorkLeaseSubject(w.id, w.queueID), reqBytes) cancel() if reqErr != nil { if !errors.Is(reqErr, context.DeadlineExceeded) && !errors.Is(reqErr, nats.ErrTimeout) { @@ -385,7 +393,7 @@ func (w *mockWorker) pullAssignment(timeout time.Duration) protocol.AssignMsg { return protocol.AssignMsg{} // unreachable } -// publishStatus publishes a [protocol.TaskStatusMsg] to task.status.. +// publishStatus publishes a [protocol.TaskStatusMsg] to task.status... func (w *mockWorker) publishStatus(assign protocol.AssignMsg, status string, exitCode *int) { w.t.Helper() @@ -407,12 +415,12 @@ func (w *mockWorker) publishStatus(assign protocol.AssignMsg, status string, exi ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - if _, err := w.js.Publish(ctx, bus.TaskStatusSubject(assign.JobID), data); err != nil { + if _, err := w.js.Publish(ctx, bus.TaskStatusSubject(w.id, assign.JobID), data); err != nil { w.t.Fatalf("mockWorker.publishStatus(%s): publish: %v", status, err) } } -// publishLogChunk publishes a [protocol.LogChunkMsg] to task.logs.. +// publishLogChunk publishes a [protocol.LogChunkMsg] to task.logs... func (w *mockWorker) publishLogChunk(assign protocol.AssignMsg, seqNum int64, data string) { w.t.Helper() @@ -433,7 +441,7 @@ func (w *mockWorker) publishLogChunk(assign protocol.AssignMsg, seqNum int64, da ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - if _, err := w.js.Publish(ctx, bus.TaskLogsSubject(assign.TaskID), raw); err != nil { + if _, err := w.js.Publish(ctx, bus.TaskLogsSubject(w.id, assign.TaskID), raw); err != nil { w.t.Fatalf("mockWorker.publishLogChunk: publish: %v", err) } } diff --git a/test/integration/load_test.go b/test/integration/load_test.go index 023ed906..144560f1 100644 --- a/test/integration/load_test.go +++ b/test/integration/load_test.go @@ -190,7 +190,7 @@ type loadWorker struct { } // newLoadWorker dials NATS and returns a load worker that requests work leases -// on work.lease.. Multiple load workers request independently; the +// on work.lease... Multiple load workers request independently; the // server's atomic LeaseReadyTask distributes ready tasks among them. func newLoadWorker(tb testing.TB, natsURL, workerID, farmID, queueID string) *loadWorker { tb.Helper() @@ -246,7 +246,7 @@ func (w *loadWorker) register(tb testing.TB) { } ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - if _, err := w.js.Publish(ctx, bus.SubjectWorkerRegister, data); err != nil { + if _, err := w.js.Publish(ctx, bus.WorkerRegisterSubject(w.id), data); err != nil { tb.Fatalf("loadWorker.register %s: publish: %v", w.id, err) } } @@ -273,7 +273,7 @@ func (w *loadWorker) heartbeatLoop(ctx context.Context, interval time.Duration, return } pubCtx, pubCancel := context.WithTimeout(ctx, 2*time.Second) - if _, pubErr := w.js.Publish(pubCtx, bus.SubjectWorkerHeartbeat, data); pubErr != nil { + if _, pubErr := w.js.Publish(pubCtx, bus.WorkerHeartbeatSubject(w.id), data); pubErr != nil { pubCancel() return } @@ -303,7 +303,7 @@ func (w *loadWorker) drainLoop(ctx context.Context, wg *sync.WaitGroup, assigned return } reqCtx, cancel := context.WithTimeout(ctx, 35*time.Second) - msg, reqErr := w.nc.RequestWithContext(reqCtx, bus.WorkLeaseSubject(w.queueID), reqBytes) + msg, reqErr := w.nc.RequestWithContext(reqCtx, bus.WorkLeaseSubject(w.id, w.queueID), reqBytes) cancel() if reqErr != nil { if ctx.Err() != nil { @@ -333,7 +333,7 @@ func (w *loadWorker) drainLoop(ctx context.Context, wg *sync.WaitGroup, assigned } } -// publishStatus publishes a TaskStatusMsg to task.status.. +// publishStatus publishes a TaskStatusMsg to task.status... func (w *loadWorker) publishStatus(assign protocol.AssignMsg, status string, exitCode *int) { msg := protocol.TaskStatusMsg{ Version: protocol.ProtocolVersion, @@ -352,7 +352,7 @@ func (w *loadWorker) publishStatus(assign protocol.AssignMsg, status string, exi } ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) defer cancel() - if _, err := w.js.Publish(ctx, bus.TaskStatusSubject(assign.JobID), data); err != nil { + if _, err := w.js.Publish(ctx, bus.TaskStatusSubject(w.id, assign.JobID), data); err != nil { // Best-effort in load loops; callers detect missed completions via // the waitAllJobsTerminal timeout. _ = err diff --git a/web/src/auth/policy.test.ts b/web/src/auth/policy.test.ts index 196c9715..927bb296 100644 --- a/web/src/auth/policy.test.ts +++ b/web/src/auth/policy.test.ts @@ -88,6 +88,7 @@ describe('permission union stays in lockstep with the server', () => { 'apikeys.self', 'apikeys.admin', 'isolation.manage', + 'workers.enroll', ] it('ALL_PERMISSIONS matches the server-declared permission set exactly', () => { diff --git a/web/src/auth/policy.ts b/web/src/auth/policy.ts index d939d1f5..f0cef238 100644 --- a/web/src/auth/policy.ts +++ b/web/src/auth/policy.ts @@ -27,6 +27,7 @@ export type Permission = | 'apikeys.self' | 'apikeys.admin' | 'isolation.manage' + | 'workers.enroll' // PERMISSION_SET exists only for compile-time exhaustiveness and to give a // test something with a runtime representation to check: the Permission @@ -55,6 +56,7 @@ const PERMISSION_SET: Record = { 'apikeys.self': true, 'apikeys.admin': true, 'isolation.manage': true, + 'workers.enroll': true, } /** Every permission the server can grant, sorted. See PERMISSION_SET. */