From 3a1461315f02440dbc978ceca9c0c6b353038fb9 Mon Sep 17 00:00:00 2001 From: Efremenko Arhip Date: Mon, 27 Jul 2026 16:11:07 +0300 Subject: [PATCH] feat: self-service worker enrollment bound to a user account Let a signed-in user turn their own machine into a worker without the shared token. The coordinator already binds a JWT-authenticated registration to owner_id as untrusted; this adds the missing pieces. userservice: long-lived worker keys (scimesh_wk_live_*, hash-at-rest) with create/list/revoke and a public /worker-tokens/exchange that trades a key for a short-lived JWT carrying the owner current role/verified. python worker: SCIMESH_WORKER_KEY + SCIMESH_USERSERVICE_URL; a token provider exchanges the key and refreshes the JWT proactively and on 401, so a long-running worker survives token expiry. Static bearer token path is unchanged. coordinator UI: an "add your machine" page that mints a key and shows a ready-to-run command, proxying key management to the userservice; the dashboard gains an owner-scoped "my machines" section. docs: how to run a worker from your account, plus the untrusted/quorum/ verified trust model. --- coordinator/cmd/coordinator/main.go | 2 +- coordinator/docker-compose.users.yml | 5 + coordinator/internal/infra/config.go | 43 ++-- coordinator/internal/memstore/ui_read.go | 36 ++- .../internal/storage/postgres/ui_read_repo.go | 26 +++ coordinator/internal/transport/http/server.go | 47 +++- .../transport/http/templates/add-worker.html | 55 +++++ .../transport/http/templates/dashboard.html | 9 +- .../internal/transport/http/ui_workers.go | 118 ++++++++++ .../http/ui_workers_internal_test.go | 106 +++++++++ coordinator/internal/usecase/ui.go | 34 ++- .../internal/usecase/ui_workers_scope_test.go | 66 ++++++ docs/building-workers.md | 55 ++++- scimesh/worker/artifacts.py | 50 ++++- scimesh/worker/auth.py | 128 +++++++++++ scimesh/worker/cli.py | 30 ++- scimesh/worker/config.py | 21 ++ scimesh/worker/coordinator.py | 31 ++- tests/test_worker_auth.py | 205 ++++++++++++++++++ users/cmd/userservice/main.go | 15 +- users/internal/domain/errors.go | 2 + users/internal/domain/workerkey.go | 84 +++++++ users/internal/domain/workerkey_test.go | 61 ++++++ .../storage/postgres/workerkey_repo.go | 123 +++++++++++ users/internal/transport/http/dto.go | 50 +++++ users/internal/transport/http/errors.go | 6 + users/internal/transport/http/handlers.go | 100 ++++++++- users/internal/transport/http/server.go | 38 +++- users/internal/usecase/errors.go | 8 + users/internal/usecase/ports.go | 19 ++ users/internal/usecase/workerkey.go | 112 ++++++++++ users/internal/usecase/workerkey_test.go | 186 ++++++++++++++++ users/migrations/0003_worker_keys.down.sql | 5 + users/migrations/0003_worker_keys.up.sql | 31 +++ 34 files changed, 1829 insertions(+), 78 deletions(-) create mode 100644 coordinator/internal/transport/http/templates/add-worker.html create mode 100644 coordinator/internal/transport/http/ui_workers.go create mode 100644 coordinator/internal/transport/http/ui_workers_internal_test.go create mode 100644 coordinator/internal/usecase/ui_workers_scope_test.go create mode 100644 scimesh/worker/auth.py create mode 100644 tests/test_worker_auth.py create mode 100644 users/internal/domain/workerkey.go create mode 100644 users/internal/domain/workerkey_test.go create mode 100644 users/internal/storage/postgres/workerkey_repo.go create mode 100644 users/internal/usecase/workerkey.go create mode 100644 users/internal/usecase/workerkey_test.go create mode 100644 users/migrations/0003_worker_keys.down.sql create mode 100644 users/migrations/0003_worker_keys.up.sql diff --git a/coordinator/cmd/coordinator/main.go b/coordinator/cmd/coordinator/main.go index dd6e33d..af83019 100644 --- a/coordinator/cmd/coordinator/main.go +++ b/coordinator/cmd/coordinator/main.go @@ -121,7 +121,7 @@ func run() error { // pool.Ping backs /health: readiness means the database answers, not just // that the process is alive. - api := httptransport.NewServer(useCases, log, cfg.RequestTimeout, cfg.HeartbeatInterval, cfg.MaxUploadBytes, cfg.JWTSecret, cfg.UserserviceURL, m, pool.Ping) + api := httptransport.NewServer(useCases, log, cfg.RequestTimeout, cfg.HeartbeatInterval, cfg.MaxUploadBytes, cfg.JWTSecret, cfg.UserserviceURL, m, pool.Ping, cfg.PublicCoordinatorURL, cfg.PublicUserserviceURL) err = infra.RunServer(ctx, log, cfg.Addr, api.Handler(cfg.Token, cfg.UIToken)) // Shutdown order matters, and defers alone cannot express it (they run diff --git a/coordinator/docker-compose.users.yml b/coordinator/docker-compose.users.yml index 30e5483..2b27262 100644 --- a/coordinator/docker-compose.users.yml +++ b/coordinator/docker-compose.users.yml @@ -64,3 +64,8 @@ services: environment: JWT_SECRET: ${JWT_SECRET} USERSERVICE_URL: http://userservice:8081 + # Browser/host-facing URLs for the "add your machine" command. A user's + # worker runs on the host, so it reaches the published ports on localhost, + # not the in-cluster service names. + PUBLIC_COORDINATOR_URL: http://localhost:${COORDINATOR_PORT:-8080} + PUBLIC_USERSERVICE_URL: http://localhost:${USERSERVICE_PORT:-8081} diff --git a/coordinator/internal/infra/config.go b/coordinator/internal/infra/config.go index 5a414bd..94ca185 100644 --- a/coordinator/internal/infra/config.go +++ b/coordinator/internal/infra/config.go @@ -37,6 +37,13 @@ type Config struct { // login/registration (cookie session) instead of the static UI_AUTH_TOKEN // basic auth. Empty keeps the basic-auth UI. UserserviceURL string + // Browser-facing base URLs used to render the "add your machine" command on + // the UI. They must be reachable from a user's own machine, which is not + // necessarily the in-cluster address the coordinator uses for UserserviceURL. + // PublicCoordinatorURL empty lets the page fall back to its own origin; + // PublicUserserviceURL empty falls back to UserserviceURL. + PublicCoordinatorURL string + PublicUserserviceURL string // Minimum log level: debug, info, warn, error. LogLevel string @@ -92,23 +99,25 @@ func LoadConfig() (Config, error) { DatabaseURL: os.Getenv("DATABASE_URL"), // COORDINATOR_TOKEN is the contract name; WORKER_AUTH_TOKEN is the // former name, still honoured so existing .env files keep working. - Token: getEnv("COORDINATOR_TOKEN", os.Getenv("WORKER_AUTH_TOKEN")), - UIToken: os.Getenv("UI_AUTH_TOKEN"), - JWTSecret: os.Getenv("JWT_SECRET"), - UserserviceURL: os.Getenv("USERSERVICE_URL"), - LogLevel: getEnv("LOG_LEVEL", "info"), - LogFile: os.Getenv("LOG_FILE"), - StorageDir: getEnv("COORDINATOR_STORAGE_DIR", "./data"), - MaxUploadBytes: 1 << 30, // 1 GiB - DBMaxConns: 10, - DBConnectTimeout: 30 * time.Second, - RequestTimeout: 15 * time.Second, - HeartbeatInterval: 15 * time.Second, - LeaseDuration: 2 * time.Minute, - DefaultMaxAttempts: 3, - QuorumSize: 2, - ReaperInterval: 30 * time.Second, - WorkerOfflineAfter: 1 * time.Minute, + Token: getEnv("COORDINATOR_TOKEN", os.Getenv("WORKER_AUTH_TOKEN")), + UIToken: os.Getenv("UI_AUTH_TOKEN"), + JWTSecret: os.Getenv("JWT_SECRET"), + UserserviceURL: os.Getenv("USERSERVICE_URL"), + PublicCoordinatorURL: os.Getenv("PUBLIC_COORDINATOR_URL"), + PublicUserserviceURL: getEnv("PUBLIC_USERSERVICE_URL", os.Getenv("USERSERVICE_URL")), + LogLevel: getEnv("LOG_LEVEL", "info"), + LogFile: os.Getenv("LOG_FILE"), + StorageDir: getEnv("COORDINATOR_STORAGE_DIR", "./data"), + MaxUploadBytes: 1 << 30, // 1 GiB + DBMaxConns: 10, + DBConnectTimeout: 30 * time.Second, + RequestTimeout: 15 * time.Second, + HeartbeatInterval: 15 * time.Second, + LeaseDuration: 2 * time.Minute, + DefaultMaxAttempts: 3, + QuorumSize: 2, + ReaperInterval: 30 * time.Second, + WorkerOfflineAfter: 1 * time.Minute, } if cfg.DatabaseURL == "" { diff --git a/coordinator/internal/memstore/ui_read.go b/coordinator/internal/memstore/ui_read.go index cc2c681..8f3f191 100644 --- a/coordinator/internal/memstore/ui_read.go +++ b/coordinator/internal/memstore/ui_read.go @@ -87,16 +87,44 @@ func (r *UIReadRepo) ListWorkers(_ context.Context, limit int) ([]domain.Worker, copy.Capabilities = append([]string(nil), worker.Capabilities...) out = append(out, copy) } + sortWorkers(out) + if len(out) > limit { + out = out[:limit] + } + return out, nil +} + +func (r *UIReadRepo) ListWorkersByOwner(_ context.Context, owner uuid.UUID, limit int) ([]domain.Worker, error) { + if limit < 1 || limit > 100 { + return nil, domain.ErrInvalidInput + } + r.workers.mu.Lock() + defer r.workers.mu.Unlock() + out := []domain.Worker{} + for _, worker := range r.workers.workers { + if worker.OwnerID == nil || *worker.OwnerID != owner { + continue + } + copy := *worker + copy.Capabilities = append([]string(nil), worker.Capabilities...) + out = append(out, copy) + } + sortWorkers(out) + if len(out) > limit { + out = out[:limit] + } + return out, nil +} + +// sortWorkers orders workers most-recently-seen first, breaking ties on id so +// the order is deterministic across calls. +func sortWorkers(out []domain.Worker) { sort.Slice(out, func(i, j int) bool { if out[i].LastHeartbeatAt.Equal(out[j].LastHeartbeatAt) { return out[i].ID.String() > out[j].ID.String() } return out[i].LastHeartbeatAt.After(out[j].LastHeartbeatAt) }) - if len(out) > limit { - out = out[:limit] - } - return out, nil } func (r *UIReadRepo) ListArtifactsByJob(_ context.Context, jobID uuid.UUID) ([]domain.Artifact, error) { r.artifacts.mu.Lock() diff --git a/coordinator/internal/storage/postgres/ui_read_repo.go b/coordinator/internal/storage/postgres/ui_read_repo.go index 5c8a6a6..1382ee4 100644 --- a/coordinator/internal/storage/postgres/ui_read_repo.go +++ b/coordinator/internal/storage/postgres/ui_read_repo.go @@ -128,6 +128,32 @@ func (r *UIReadRepo) ListWorkers(ctx context.Context, limit int) ([]domain.Worke return workers, rows.Err() } +func (r *UIReadRepo) ListWorkersByOwner(ctx context.Context, owner uuid.UUID, limit int) ([]domain.Worker, error) { + if limit < 1 || limit > 100 { + return nil, domain.ErrInvalidInput + } + sql, args, err := psql.Select(workerColumns...).From("workers"). + Where(sq.Eq{"owner_id": owner}). + OrderBy("last_heartbeat_at DESC", "id DESC").Limit(uint64(limit)).ToSql() + if err != nil { + return nil, err + } + rows, err := conn(ctx, r.pool).Query(ctx, sql, args...) + if err != nil { + return nil, fmt.Errorf("list workers by owner: %w", err) + } + defer rows.Close() + workers := make([]domain.Worker, 0) + for rows.Next() { + worker, err := scanWorker(rows) + if err != nil { + return nil, err + } + workers = append(workers, *worker) + } + return workers, rows.Err() +} + func (r *UIReadRepo) ListArtifactsByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Artifact, error) { sql, args, err := psql.Select(artifactColumns...).From("artifacts").Where(sq.Eq{"job_id": jobID}).OrderBy("created_at ASC", "id ASC").ToSql() if err != nil { diff --git a/coordinator/internal/transport/http/server.go b/coordinator/internal/transport/http/server.go index 6439da1..a2967e7 100644 --- a/coordinator/internal/transport/http/server.go +++ b/coordinator/internal/transport/http/server.go @@ -49,6 +49,11 @@ type Server struct { // userserviceURL is the base URL the UI proxies login/registration to. Empty // keeps the static basic-auth UI. userserviceURL string + // publicCoordinatorURL / publicUserserviceURL are the browser-facing URLs + // rendered into the worker-enrollment command. Either may be empty; the + // template falls back (own origin / userserviceURL respectively). + publicCoordinatorURL string + publicUserserviceURL string // httpClient makes the login/register calls to the userservice. httpClient *http.Client // metrics holds the Prometheus registry and HTTP instrumentation. @@ -59,21 +64,33 @@ type Server struct { } func NewServer(uc UseCases, log *slog.Logger, requestTimeout, heartbeatInterval time.Duration, - maxUploadBytes int64, jwtSecret, userserviceURL string, m *metrics.Metrics, ready func(context.Context) error) *Server { + maxUploadBytes int64, jwtSecret, userserviceURL string, m *metrics.Metrics, ready func(context.Context) error, + publicURLs ...string) *Server { if m == nil { m = metrics.New() } + // publicURLs is variadic so existing callers/tests need no change: [0] is the + // public coordinator URL, [1] the public userservice URL; both optional. + var publicCoordinatorURL, publicUserserviceURL string + if len(publicURLs) > 0 { + publicCoordinatorURL = strings.TrimRight(publicURLs[0], "/") + } + if len(publicURLs) > 1 { + publicUserserviceURL = strings.TrimRight(publicURLs[1], "/") + } return &Server{ - uc: uc, - log: log, - requestTimeout: requestTimeout, - heartbeatInterval: heartbeatInterval, - maxUploadBytes: maxUploadBytes, - verifier: tokenpkg.NewVerifier(jwtSecret), - userserviceURL: strings.TrimRight(userserviceURL, "/"), - httpClient: &http.Client{Timeout: 10 * time.Second}, - metrics: m, - ready: ready, + uc: uc, + log: log, + requestTimeout: requestTimeout, + heartbeatInterval: heartbeatInterval, + maxUploadBytes: maxUploadBytes, + verifier: tokenpkg.NewVerifier(jwtSecret), + userserviceURL: strings.TrimRight(userserviceURL, "/"), + publicCoordinatorURL: publicCoordinatorURL, + publicUserserviceURL: publicUserserviceURL, + httpClient: &http.Client{Timeout: 10 * time.Second}, + metrics: m, + ready: ready, } } @@ -139,6 +156,14 @@ func (s *Server) Handler(token string, uiToken ...string) http.Handler { ui.Handle(rt.pattern, gate(rt.handler)) } ui.Handle("GET /ui/profile", gate(http.HandlerFunc(s.handleUIProfile))) + // Worker enrollment: a user creates/lists/revokes their own worker keys + // and copies a ready-to-run command. Session-only — it proxies to the + // userservice with the caller's token, so it has no meaning under basic + // auth (which has no userservice). + ui.Handle("GET /ui/workers/new", gate(http.HandlerFunc(s.handleUIAddWorker))) + ui.Handle("GET /ui/api/worker-keys", gate(http.HandlerFunc(s.handleUIWorkerKeysList))) + ui.Handle("POST /ui/api/worker-keys", gate(http.HandlerFunc(s.handleUIWorkerKeyCreate))) + ui.Handle("POST /ui/api/worker-keys/{id}/revoke", gate(http.HandlerFunc(s.handleUIWorkerKeyRevoke))) // Admin panel: session + admin role. ui.Handle("GET /ui/admin", chain(http.HandlerFunc(s.handleUIAdmin), gate, requireAdmin)) ui.Handle("POST /ui/admin/user-action", chain(http.HandlerFunc(s.handleUIAdminUserAction), gate, requireAdmin)) diff --git a/coordinator/internal/transport/http/templates/add-worker.html b/coordinator/internal/transport/http/templates/add-worker.html new file mode 100644 index 0000000..e0de231 --- /dev/null +++ b/coordinator/internal/transport/http/templates/add-worker.html @@ -0,0 +1,55 @@ +{{define "add-worker.html"}} + + + + + + Add your machine · SciMesh + + + +
+ ← Back to control room

Contribute compute

Turn this computer into a worker

Create a key, install the worker, and run one command. The worker binds to your account and pulls tasks whenever it is online.

+
+
+

Your worker keys

+

A key is long-lived and does not expire like a login. The worker trades it for short-lived tokens automatically. Revoke a key to stop its machines.

+
+ + + + +
+ +
+
+ +
+
+ + + +{{end}} diff --git a/coordinator/internal/transport/http/templates/dashboard.html b/coordinator/internal/transport/http/templates/dashboard.html index f4ee300..cbbd118 100644 --- a/coordinator/internal/transport/http/templates/dashboard.html +++ b/coordinator/internal/transport/http/templates/dashboard.html @@ -13,7 +13,7 @@

Local scientific compute

SciMesh control room

Follow the real path from a molecular TSV to a globally reduced similarity result—without reading coordinator logs.

Live overview · refreshes every 2 seconds
-
{{if .Session}}Signed in · {{.Session.Role}}{{end}}{{if .Session}}Profile{{end}}{{if and .Session (eq .Session.Role "admin")}}Admin{{end}}＋ New similarity search{{if .Session}}
{{end}}
+
{{if .Session}}Signed in · {{.Session.Role}}{{end}}{{if .Session}}Profile{{end}}{{if and .Session (eq .Session.Role "admin")}}Admin{{end}}{{if .Session}}🖥 Add your machine{{end}}＋ New similarity search{{if .Session}}
{{end}}
How a search becomes a result
01Upload TSVThe coordinator validates and slices the dataset.
02Run shardsWorkers fingerprint molecules and return shard top-k CSVs.
03Merge exactlyThe coordinator ranks retained candidates deterministically.
04Download CSVA checksum-protected global result is ready.
@@ -23,6 +23,7 @@

Recent computations

{{len .Jobs}} shown · newest first

{{range .Jobs}}
{{workloadLabel .Workload}}
{{.ID}}
{{statusLabel .Status}}
{{statusHint .Status}}
{{.Completed}} / {{.Total}} shards complete{{if gt .Failed 0}} · {{.Failed}} failed{{end}}
{{else}}
No computations yet.
Start a small similarity search, then keep one or more workers running to watch this dashboard come alive.
{{end}}
+ {{if and .Session (ne .Session.Role "admin")}}

My machines

Workers you registered. Add your machine →

{{range .MyWorkers}}
{{.Name}}{{workerStatusLabel .Status}}
{{.ID}}

{{range .Capabilities}}{{.}}{{end}}

Last signal · {{time .LastHeartbeatAt}}

{{else}}
No machine of yours is connected.
Turn this computer into a worker →
{{end}}
{{end}}

Worker fleet

Workers register themselves; this page never controls their processes.

{{range .Workers}}
{{.Name}}{{workerStatusLabel .Status}}
{{.ID}}

{{range .Capabilities}}{{.}}{{end}}

Last signal · {{time .LastHeartbeatAt}}

{{else}}
No worker is registered.
Start scimesh-worker in another terminal, then return here.
{{end}}
diff --git a/coordinator/internal/transport/http/ui_workers.go b/coordinator/internal/transport/http/ui_workers.go new file mode 100644 index 0000000..0db4b0f --- /dev/null +++ b/coordinator/internal/transport/http/ui_workers.go @@ -0,0 +1,118 @@ +package http + +import ( + "bytes" + "context" + "io" + "net/http" + + "github.com/google/uuid" + + "github.com/emil28092005/SciMesh/coordinator/internal/authctx" + "github.com/emil28092005/SciMesh/coordinator/internal/usecase" +) + +// handleUIAddWorker renders the "add your machine" page: instructions, the +// user's existing worker keys, and a ready-to-run command carrying a freshly +// minted key. All key operations happen client-side against the JSON endpoints +// below; this handler only supplies the browser-facing URLs. +func (s *Server) handleUIAddWorker(w http.ResponseWriter, r *http.Request) { + data := map[string]any{ + "CoordinatorURL": s.publicCoordinatorURL, + "UserserviceURL": s.publicUserserviceURL, + } + if req, ok := authctx.From(r.Context()); ok { + data["Session"] = &usecase.SessionView{Role: req.Role, Verified: req.Verified} + } + s.renderUI(w, "add-worker.html", data) +} + +// handleUIWorkerKeysList proxies the caller's live worker keys from the +// userservice, forwarding their session token. +func (s *Server) handleUIWorkerKeysList(w http.ResponseWriter, r *http.Request) { + c, err := r.Cookie(sessionCookie) + if err != nil { + writeJSON(w, http.StatusUnauthorized, map[string]string{"error": "not signed in"}) + return + } + status, body, err := s.callUserserviceAuthed(r.Context(), http.MethodGet, "/worker-keys", c.Value) + if err != nil { + s.log.Error("worker-keys list proxy", "err", err) + writeJSON(w, http.StatusBadGateway, map[string]string{"error": "userservice unavailable"}) + return + } + proxyJSON(w, status, body) +} + +// handleUIWorkerKeyCreate mints a new worker key via the userservice and returns +// its response — including the one-time plaintext key — straight to the browser. +func (s *Server) handleUIWorkerKeyCreate(w http.ResponseWriter, r *http.Request) { + c, err := r.Cookie(sessionCookie) + if err != nil { + writeJSON(w, http.StatusUnauthorized, map[string]string{"error": "not signed in"}) + return + } + body, _ := io.ReadAll(io.LimitReader(r.Body, 1<<12)) + if len(body) == 0 { + body = []byte("{}") + } + status, respBody, err := s.callUserserviceAuthedBody(r.Context(), http.MethodPost, "/worker-keys", c.Value, body) + if err != nil { + s.log.Error("worker-key create proxy", "err", err) + writeJSON(w, http.StatusBadGateway, map[string]string{"error": "userservice unavailable"}) + return + } + proxyJSON(w, status, respBody) +} + +// handleUIWorkerKeyRevoke retires one of the caller's keys via the userservice. +// The id is validated as a UUID so the proxied path can never be attacker-shaped. +func (s *Server) handleUIWorkerKeyRevoke(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + if _, err := uuid.Parse(id); err != nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid worker key id"}) + return + } + c, err := r.Cookie(sessionCookie) + if err != nil { + writeJSON(w, http.StatusUnauthorized, map[string]string{"error": "not signed in"}) + return + } + status, _, err := s.callUserserviceAuthed(r.Context(), http.MethodDelete, "/worker-keys/"+id, c.Value) + if err != nil { + s.log.Error("worker-key revoke proxy", "err", err) + writeJSON(w, http.StatusBadGateway, map[string]string{"error": "userservice unavailable"}) + return + } + w.WriteHeader(status) +} + +// proxyJSON forwards a userservice JSON response verbatim, preserving its status. +func proxyJSON(w http.ResponseWriter, status int, body []byte) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + _, _ = w.Write(body) +} + +// callUserserviceAuthedBody is callUserserviceAuthed with a JSON request body, +// used for the create call. Kept separate so the bodyless admin/profile callers +// stay unchanged. +func (s *Server) callUserserviceAuthedBody(ctx context.Context, method, path, bearer string, body []byte) (int, []byte, error) { + req, err := http.NewRequestWithContext(ctx, method, s.userserviceURL+path, bytes.NewReader(body)) //nolint:gosec // G704: path is a fixed literal, host is config + if err != nil { + return 0, nil, err + } + req.Header.Set("Authorization", "Bearer "+bearer) + req.Header.Set("Content-Type", "application/json") + + resp, err := s.httpClient.Do(req) //nolint:gosec // G704: see above + if err != nil { + return 0, nil, err + } + defer func() { _ = resp.Body.Close() }() + respBody, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + if err != nil { + return 0, nil, err + } + return resp.StatusCode, respBody, nil +} diff --git a/coordinator/internal/transport/http/ui_workers_internal_test.go b/coordinator/internal/transport/http/ui_workers_internal_test.go new file mode 100644 index 0000000..7a63f68 --- /dev/null +++ b/coordinator/internal/transport/http/ui_workers_internal_test.go @@ -0,0 +1,106 @@ +package http + +import ( + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func TestWorkerKeyCreateProxiesWithBody(t *testing.T) { + var gotAuth, gotPath, gotMethod, gotBody string + stub := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuth, gotPath, gotMethod = r.Header.Get("Authorization"), r.URL.Path, r.Method + b, _ := io.ReadAll(r.Body) + gotBody = string(b) + w.WriteHeader(http.StatusCreated) + _, _ = w.Write([]byte(`{"id":"11111111-1111-1111-1111-111111111111","name":"box","prefix":"scimesh_wk_live_ab","created_at":"2026-07-26T00:00:00Z","key":"scimesh_wk_live_secret"}`)) + })) + defer stub.Close() + s := newLoginServer(stub) + + req := newReq(http.MethodPost, "/ui/api/worker-keys", strings.NewReader(`{"name":"box"}`)) + req.AddCookie(&http.Cookie{Name: sessionCookie, Value: "my.jwt"}) + rec := httptest.NewRecorder() + s.handleUIWorkerKeyCreate(rec, req) + + if gotAuth != "Bearer my.jwt" || gotPath != "/worker-keys" || gotMethod != http.MethodPost { + t.Fatalf("proxy: auth=%q path=%q method=%q", gotAuth, gotPath, gotMethod) + } + if !strings.Contains(gotBody, `"name":"box"`) { + t.Errorf("request body not forwarded: %q", gotBody) + } + if rec.Code != http.StatusCreated || !strings.Contains(rec.Body.String(), "scimesh_wk_live_secret") { + t.Errorf("response not passed through: %d %s", rec.Code, rec.Body.String()) + } +} + +func TestWorkerKeysListProxies(t *testing.T) { + stub := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet || r.URL.Path != "/worker-keys" { + t.Errorf("unexpected upstream call %s %s", r.Method, r.URL.Path) + } + _, _ = w.Write([]byte(`{"worker_keys":[{"id":"1","name":"box","prefix":"scimesh_wk_live_ab","created_at":"2026-07-26T00:00:00Z"}]}`)) + })) + defer stub.Close() + s := newLoginServer(stub) + + req := newReq(http.MethodGet, "/ui/api/worker-keys", nil) + req.AddCookie(&http.Cookie{Name: sessionCookie, Value: "my.jwt"}) + rec := httptest.NewRecorder() + s.handleUIWorkerKeysList(rec, req) + + if rec.Code != http.StatusOK || !strings.Contains(rec.Body.String(), "worker_keys") { + t.Errorf("list not passed through: %d %s", rec.Code, rec.Body.String()) + } +} + +func TestWorkerKeyRevokeProxiesDelete(t *testing.T) { + const id = "22222222-2222-2222-2222-222222222222" + var gotPath, gotMethod string + stub := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath, gotMethod = r.URL.Path, r.Method + w.WriteHeader(http.StatusNoContent) + })) + defer stub.Close() + s := newLoginServer(stub) + + req := newReq(http.MethodPost, "/ui/api/worker-keys/"+id+"/revoke", nil) + req.SetPathValue("id", id) + req.AddCookie(&http.Cookie{Name: sessionCookie, Value: "my.jwt"}) + rec := httptest.NewRecorder() + s.handleUIWorkerKeyRevoke(rec, req) + + if gotMethod != http.MethodDelete || gotPath != "/worker-keys/"+id { + t.Fatalf("proxy: method=%q path=%q", gotMethod, gotPath) + } + if rec.Code != http.StatusNoContent { + t.Errorf("revoke status = %d, want 204", rec.Code) + } +} + +func TestWorkerKeyRevokeRejectsBadID(t *testing.T) { + s := newLoginServer(httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + t.Fatal("must not call userservice for an invalid id") + }))) + req := newReq(http.MethodPost, "/ui/api/worker-keys/not-a-uuid/revoke", nil) + req.SetPathValue("id", "not-a-uuid") + req.AddCookie(&http.Cookie{Name: sessionCookie, Value: "my.jwt"}) + rec := httptest.NewRecorder() + s.handleUIWorkerKeyRevoke(rec, req) + if rec.Code != http.StatusBadRequest { + t.Errorf("bad id: got %d, want 400", rec.Code) + } +} + +func TestWorkerKeysRequireSession(t *testing.T) { + s := newLoginServer(httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + t.Fatal("must not call userservice without a session cookie") + }))) + rec := httptest.NewRecorder() + s.handleUIWorkerKeysList(rec, newReq(http.MethodGet, "/ui/api/worker-keys", nil)) + if rec.Code != http.StatusUnauthorized { + t.Errorf("no cookie: got %d, want 401", rec.Code) + } +} diff --git a/coordinator/internal/usecase/ui.go b/coordinator/internal/usecase/ui.go index a8624ab..f26d102 100644 --- a/coordinator/internal/usecase/ui.go +++ b/coordinator/internal/usecase/ui.go @@ -21,6 +21,9 @@ type UIReadRepository interface { ListTasksByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Task, error) ListTasksByJobs(ctx context.Context, jobIDs []uuid.UUID) (map[uuid.UUID][]domain.Task, error) ListWorkers(ctx context.Context, limit int) ([]domain.Worker, error) + // ListWorkersByOwner returns the most recent workers registered by one user, + // for the "my machines" section of the dashboard. + ListWorkersByOwner(ctx context.Context, owner uuid.UUID, limit int) ([]domain.Worker, error) ListArtifactsByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Artifact, error) } @@ -83,8 +86,11 @@ type WorkerCard struct { } type DashboardView struct { - Jobs []JobCard `json:"jobs"` - Workers []WorkerCard `json:"workers"` + Jobs []JobCard `json:"jobs"` + Workers []WorkerCard `json:"workers"` + // MyWorkers is the signed-in user's own registered workers. Empty for an + // admin or a basic-auth operator, who instead see the whole fleet in Workers. + MyWorkers []WorkerCard `json:"my_workers"` ActiveJobs int `json:"active_jobs"` FinishedJobs int `json:"finished_jobs"` OnlineWorkers int `json:"online_workers"` @@ -152,15 +158,37 @@ func (d *Dashboard) Overview(ctx context.Context, limit int) (DashboardView, err } } for _, worker := range workers { - out.Workers = append(out.Workers, WorkerCard{ID: worker.ID.String(), Name: worker.Name, Status: string(worker.Status), Capabilities: worker.Capabilities, LastHeartbeatAt: worker.LastHeartbeatAt}) + out.Workers = append(out.Workers, workerCard(worker)) if worker.Status == domain.WorkerOnline || worker.Status == domain.WorkerBusy { out.OnlineWorkers++ } } + // A plain user also gets a dedicated "my machines" list scoped to their own + // registrations; an admin/operator sees only the fleet above. + if owner := uiOwnerFilter(ctx); owner != nil { + mine, err := d.read.ListWorkersByOwner(ctx, *owner, limit) + if err != nil { + return DashboardView{}, err + } + out.MyWorkers = make([]WorkerCard, 0, len(mine)) + for _, worker := range mine { + out.MyWorkers = append(out.MyWorkers, workerCard(worker)) + } + } out.Session = sessionViewFrom(ctx) return out, nil } +func workerCard(w domain.Worker) WorkerCard { + return WorkerCard{ + ID: w.ID.String(), + Name: w.Name, + Status: string(w.Status), + Capabilities: w.Capabilities, + LastHeartbeatAt: w.LastHeartbeatAt, + } +} + func (d *Dashboard) JobDetail(ctx context.Context, jobID uuid.UUID) (JobDetailView, error) { job, err := d.read.GetJob(ctx, jobID) if err != nil { diff --git a/coordinator/internal/usecase/ui_workers_scope_test.go b/coordinator/internal/usecase/ui_workers_scope_test.go new file mode 100644 index 0000000..2da6547 --- /dev/null +++ b/coordinator/internal/usecase/ui_workers_scope_test.go @@ -0,0 +1,66 @@ +package usecase_test + +import ( + "context" + "testing" + "time" + + "github.com/google/uuid" + + "github.com/emil28092005/SciMesh/coordinator/internal/domain" + "github.com/emil28092005/SciMesh/coordinator/internal/memstore" + "github.com/emil28092005/SciMesh/coordinator/internal/usecase" +) + +func newDashboardWithWorkers() (*usecase.Dashboard, *memstore.WorkerRepo) { + jobs := memstore.NewJobRepo() + tasks := memstore.NewTaskRepo() + workers := memstore.NewWorkerRepo() + artifacts := memstore.NewArtifactRepo() + return usecase.NewDashboard(memstore.NewUIReadRepo(jobs, tasks, workers, artifacts)), workers +} + +func seedWorker(t *testing.T, workers *memstore.WorkerRepo, owner *uuid.UUID, name string) { + t.Helper() + w := &domain.Worker{ + ID: uuid.New(), + Name: name, + Capabilities: []string{"similarity-search"}, + Status: domain.WorkerOnline, + OwnerID: owner, + LastHeartbeatAt: time.Now().UTC(), + } + if err := workers.Insert(context.Background(), w); err != nil { + t.Fatalf("insert worker: %v", err) + } +} + +func TestOverviewSplitsMyWorkers(t *testing.T) { + dash, workers := newDashboardWithWorkers() + alice, bob := uuid.New(), uuid.New() + seedWorker(t, workers, &alice, "alice-box") + seedWorker(t, workers, &bob, "bob-box") + seedWorker(t, workers, nil, "lab-shared") // owner-less shared-token worker + + // A plain user sees the whole fleet, but MyWorkers holds only their own. + v, err := dash.Overview(userCtx(alice, "user"), 20) + if err != nil { + t.Fatal(err) + } + if len(v.Workers) != 3 { + t.Errorf("fleet shows %d workers, want 3", len(v.Workers)) + } + if len(v.MyWorkers) != 1 || v.MyWorkers[0].Name != "alice-box" { + t.Errorf("MyWorkers = %+v, want only alice-box", v.MyWorkers) + } + + // An admin is not owner-scoped: they get the fleet and no personal list. + if av, _ := dash.Overview(userCtx(uuid.New(), "admin"), 20); len(av.MyWorkers) != 0 || len(av.Workers) != 3 { + t.Errorf("admin MyWorkers=%d Workers=%d, want 0 and 3", len(av.MyWorkers), len(av.Workers)) + } + + // A basic-auth operator (no requester) also gets no personal list. + if ov, _ := dash.Overview(context.Background(), 20); len(ov.MyWorkers) != 0 { + t.Errorf("operator MyWorkers=%d, want 0", len(ov.MyWorkers)) + } +} diff --git a/docs/building-workers.md b/docs/building-workers.md index 956d302..8530f7d 100644 --- a/docs/building-workers.md +++ b/docs/building-workers.md @@ -33,14 +33,61 @@ Everything below fills in the details. ## 0. Auth -Every request except `GET /health` carries a shared bearer token: +Every request except `GET /health` carries a bearer token: + +``` +Authorization: Bearer +``` + +There are two ways to obtain that token. + +### Shared coordinator token (lab / operator workers) ``` Authorization: Bearer ``` The token is handed to you out of band (env var / secret) — the same string the -coordinator was started with. Never log it, never send it in an error body. +coordinator was started with. A worker using it registers **owner-less and +trusted**: its results are accepted without quorum. Never log it, never send it +in an error body. + +### Worker key (run a worker bound to your own account) + +Any signed-in user can turn their machine into a worker without the shared +secret: + +1. In the web UI, open **“Add your machine”** (`/ui/workers/new`) and create a + **worker key** (`scimesh_wk_live_…`). It is shown once — copy it. +2. Install and run the reference worker with the copied command: + + ``` + git clone https://github.com/emil28092005/SciMesh.git + cd SciMesh + python -m venv .venv + source .venv/bin/activate + pip install -e . + + SCIMESH_COORDINATOR_URL= \ + SCIMESH_USERSERVICE_URL= \ + SCIMESH_WORKER_KEY=scimesh_wk_live_xxx \ + scimesh-worker --worker-name my-machine + ``` + + The worker ships in this repository, not on PyPI, so it is installed from a + clone (`pip install -e .`) rather than `pip install scimesh`. + +Under the hood the worker trades the key at `POST /worker-tokens/exchange` for a +short-lived JWT and refreshes it automatically before it expires — so unlike a +raw login token, a worker key keeps a long-running worker authenticated. Revoke +the key in the UI to cut a machine off. + +**Trust and quorum.** A worker registered with a plain user's key is +**untrusted**: its result is quarantined and only accepted once a second, +independent worker (a different owner) computes the same answer — the quorum +(default 2). If an admin marks your account **verified**, your workers become +trusted and their results count immediately; re-register the worker after being +verified so it picks up the upgraded trust. ## 1. Register (once, at startup) @@ -191,7 +238,9 @@ Per the worker contract, at minimum: - `SCIMESH_COORDINATOR_URL` (e.g. `http://coordinator:8080`) - worker name (the coordinator returns its `worker_id` at registration; `SCIMESH_WORKER_ID` is only a legacy/test override) -- the bearer token +- the credential — either `SCIMESH_BEARER_TOKEN` (shared token or a raw JWT) or + `SCIMESH_WORKER_KEY` together with `SCIMESH_USERSERVICE_URL` (a worker key the + worker exchanges and refreshes; see §0) - poll interval and request timeout - a working directory for downloaded inputs and generated outputs diff --git a/scimesh/worker/artifacts.py b/scimesh/worker/artifacts.py index a16ba43..3f2c7f8 100644 --- a/scimesh/worker/artifacts.py +++ b/scimesh/worker/artifacts.py @@ -7,9 +7,11 @@ import http.client import json from pathlib import Path from typing import Protocol +from urllib.error import HTTPError from urllib.parse import quote, urljoin, urlsplit from urllib.request import Request, build_opener +from .auth import StaticTokenProvider, TokenProvider from .coordinator import CoordinatorConflictError from .models import ClaimedTask, ProducedArtifact, UploadedArtifact from .transport import SameOriginAuthRedirectHandler, origin @@ -29,23 +31,51 @@ class ArtifactClient(Protocol): class HttpArtifactClient: """Transfers artifacts through the coordinator without leaking credentials.""" - def __init__(self, coordinator_url: str, timeout: float, bearer_token: str | None = None) -> None: + def __init__( + self, + coordinator_url: str, + timeout: float, + bearer_token: str | None = None, + *, + token_provider: TokenProvider | None = None, + ) -> None: self.coordinator_url = coordinator_url.rstrip("/") self.timeout = timeout - self.bearer_token = bearer_token + self._tokens: TokenProvider = token_provider or StaticTokenProvider(bearer_token) self.coordinator_origin = origin(coordinator_url) self._opener = build_opener(SameOriginAuthRedirectHandler(self.coordinator_origin)) + @property + def bearer_token(self) -> str | None: + return self._tokens.token() + def download(self, uri: str, destination: Path) -> None: destination.parent.mkdir(parents=True, exist_ok=True) resolved_uri = urljoin(f"{self.coordinator_url}/", uri) + self._download_once(resolved_uri, destination, allow_refresh=True) + + def _download_once(self, resolved_uri: str, destination: Path, *, allow_refresh: bool) -> None: request = Request(resolved_uri, headers=self._auth_headers_for(resolved_uri)) - with self._opener.open(request, timeout=self.timeout) as response, destination.open("wb") as target: - while chunk := response.read(1024 * 1024): - target.write(chunk) + try: + with self._opener.open(request, timeout=self.timeout) as response, destination.open("wb") as target: + while chunk := response.read(1024 * 1024): + target.write(chunk) + except HTTPError as error: + # Refresh an expired token and retry once, mirroring the coordinator + # client, so a token that lapses mid-task does not fail the download. + if error.code == 401 and allow_refresh: + self._tokens.refresh() + self._download_once(resolved_uri, destination, allow_refresh=False) + return + raise def upload( self, task: ClaimedTask, worker_id: str, artifact: ProducedArtifact + ) -> UploadedArtifact: + return self._upload_once(task, worker_id, artifact, allow_refresh=True) + + def _upload_once( + self, task: ClaimedTask, worker_id: str, artifact: ProducedArtifact, *, allow_refresh: bool ) -> UploadedArtifact: """Stream an artifact and require durable coordinator-owned metadata.""" url = ( @@ -76,6 +106,11 @@ class HttpArtifactClient: connection.send(chunk) response = connection.getresponse() body = response.read() + if response.status == 401 and allow_refresh: + # Token lapsed mid-task: refresh and retry the upload once. + self._tokens.refresh() + connection.close() + return self._upload_once(task, worker_id, artifact, allow_refresh=False) if response.status == 409: raise CoordinatorConflictError("artifact upload rejected because the task lease was lost") if response.status != 200: @@ -93,8 +128,9 @@ class HttpArtifactClient: def _auth_headers_for(self, uri: str) -> dict[str, str]: """Only coordinator-owned URLs receive the coordinator bearer token.""" - if self.bearer_token and origin(uri) == self.coordinator_origin: - return {"Authorization": f"Bearer {self.bearer_token}"} + token = self._tokens.token() + if token and origin(uri) == self.coordinator_origin: + return {"Authorization": f"Bearer {token}"} return {} def sha256_file(path: Path) -> str: diff --git a/scimesh/worker/auth.py b/scimesh/worker/auth.py new file mode 100644 index 0000000..a447867 --- /dev/null +++ b/scimesh/worker/auth.py @@ -0,0 +1,128 @@ +"""Bearer-token strategies for the worker's coordinator calls. + +A worker authenticates in one of two ways: + +* a *static* token — the shared service token or a directly supplied JWT, fixed + for the life of the process; or +* a *worker key* — a long-lived per-user credential the worker trades for a + short-lived JWT at the userservice, refreshing before that JWT expires. + +Both are exposed through the small ``TokenProvider`` protocol so the HTTP +clients neither know nor care which one is in play. +""" + +from __future__ import annotations + +import json +import time +from typing import Callable, Protocol +from urllib.error import HTTPError, URLError +from urllib.request import Request, build_opener + +from .transport import NoRedirectHandler + + +class TokenExchangeError(RuntimeError): + """The userservice refused or failed to exchange a worker key.""" + + +class TokenProvider(Protocol): + def token(self) -> str | None: + """Return the current bearer token, refreshing it if necessary.""" + + def refresh(self) -> None: + """Force the next token to be re-fetched (e.g. after a 401).""" + + +class StaticTokenProvider: + """Serves a fixed token forever. ``None`` means "send no Authorization".""" + + def __init__(self, token: str | None) -> None: + self._token = token + + def token(self) -> str | None: + return self._token + + def refresh(self) -> None: # noqa: D401 - nothing to refresh + return None + + +class WorkerKeyTokenProvider: + """Exchanges a long-lived worker key for short-lived JWTs and refreshes them. + + The token is cached until roughly ``1 - refresh_leeway`` of its lifetime has + elapsed, so the worker renews ahead of expiry instead of waiting for a 401. + A monotonic clock is injectable to keep tests deterministic. + """ + + def __init__( + self, + userservice_url: str, + worker_key: str, + timeout: float, + *, + refresh_leeway: float = 0.2, + now: Callable[[], float] = time.monotonic, + ) -> None: + self._url = userservice_url.rstrip("/") + self._key = worker_key + self._timeout = timeout + self._leeway = refresh_leeway + self._now = now + self._token: str | None = None + self._refresh_at: float = 0.0 + self._opener = build_opener(NoRedirectHandler()) + + def token(self) -> str: + if self._token is None or self._now() >= self._refresh_at: + self._exchange() + assert self._token is not None # _exchange sets it or raises + return self._token + + def refresh(self) -> None: + self._exchange() + + def _exchange(self) -> None: + request = Request( + f"{self._url}/worker-tokens/exchange", + data=json.dumps({"key": self._key}).encode(), + method="POST", + headers={"Content-Type": "application/json"}, + ) + try: + with self._opener.open(request, timeout=self._timeout) as response: + raw = response.read() + data = json.loads(raw) if raw else {} + except HTTPError as error: + # A revoked or unknown key is a permanent 401; there is nothing the + # worker can do but stop, so surface it rather than retry forever. + raise TokenExchangeError( + f"worker key exchange rejected with status {error.code}" + ) from error + except (URLError, TimeoutError, json.JSONDecodeError) as error: + raise TokenExchangeError("worker key exchange request failed") from error + + token = data.get("token") + if not isinstance(token, str) or not token: + raise TokenExchangeError("worker key exchange response is missing a token") + + expires_in = data.get("expires_in") + ttl = float(expires_in) if isinstance(expires_in, (int, float)) and expires_in > 0 else 0.0 + self._token = token + # Renew once ~(1 - leeway) of the lifetime is gone. An unknown TTL falls + # back to re-exchanging on the next call — correct, just chattier. + self._refresh_at = self._now() + ttl * (1.0 - self._leeway) + + +def provider_from_config( + *, + worker_key: str | None, + userservice_url: str | None, + bearer_token: str | None, + request_timeout: float, +) -> TokenProvider: + """Pick the token strategy: a worker key (exchange mode) wins over a static + bearer token, which in turn wins over no credential at all.""" + if worker_key and userservice_url: + return WorkerKeyTokenProvider(userservice_url, worker_key, request_timeout) + return StaticTokenProvider(bearer_token) diff --git a/scimesh/worker/cli.py b/scimesh/worker/cli.py index 422c178..f8a39ec 100644 --- a/scimesh/worker/cli.py +++ b/scimesh/worker/cli.py @@ -7,6 +7,7 @@ import logging from pathlib import Path from .artifacts import HttpArtifactClient +from .auth import provider_from_config from .config import WorkerConfig from .coordinator import HttpCoordinatorClient from .daemon import WorkerDaemon @@ -22,14 +23,23 @@ def build_parser() -> argparse.ArgumentParser: "SCIMESH_WORKER_NAME, SCIMESH_CPU_COUNT, SCIMESH_MEMORY_MB, " "SCIMESH_POLL_INTERVAL, SCIMESH_REQUEST_TIMEOUT, " "SCIMESH_HEARTBEAT_INTERVAL, SCIMESH_CLEANUP_AFTER_SECONDS, " - "SCIMESH_MAX_TASKS, and SCIMESH_BEARER_TOKEN. " - "SCIMESH_WORKER_ID is a legacy/test override." + "SCIMESH_MAX_TASKS, SCIMESH_BEARER_TOKEN, SCIMESH_WORKER_KEY, and " + "SCIMESH_USERSERVICE_URL. SCIMESH_WORKER_ID is a legacy/test override." ), ) parser.add_argument("--coordinator-url") parser.add_argument("--worker-id") parser.add_argument("--work-dir") parser.add_argument("--worker-name") + parser.add_argument( + "--worker-key", + help="Long-lived worker key from the web UI; the worker exchanges it for " + "short-lived tokens, binding it to your account. Requires --userservice-url.", + ) + parser.add_argument( + "--userservice-url", + help="Base URL of the userservice that issues tokens for --worker-key", + ) parser.add_argument("--cpu-count", type=int) parser.add_argument("--memory-mb", type=int) parser.add_argument("--poll-interval", type=float) @@ -68,11 +78,23 @@ def main(argv: list[str] | None = None) -> int: except (TypeError, ValueError) as error: parser.error(str(error)) logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") - client = HttpCoordinatorClient(config.coordinator_url, config.request_timeout, config.bearer_token) + # One shared token strategy backs both clients: a worker key (exchanged and + # refreshed) or a static bearer token, decided by what the config carries. + tokens = provider_from_config( + worker_key=config.worker_key, + userservice_url=config.userservice_url, + bearer_token=config.bearer_token, + request_timeout=config.request_timeout, + ) + client = HttpCoordinatorClient( + config.coordinator_url, config.request_timeout, token_provider=tokens + ) completed_without_interruption = WorkerDaemon( config, client, - HttpArtifactClient(config.coordinator_url, config.request_timeout, config.bearer_token), + HttpArtifactClient( + config.coordinator_url, config.request_timeout, token_provider=tokens + ), SciMeshRunner(), ).run_forever() return 0 if completed_without_interruption else 130 diff --git a/scimesh/worker/config.py b/scimesh/worker/config.py index 84a1157..b56fba7 100644 --- a/scimesh/worker/config.py +++ b/scimesh/worker/config.py @@ -11,6 +11,14 @@ from typing import Mapping from urllib.parse import urlsplit +def _clean_url(value: object | None) -> str | None: + """Normalise an optional URL: drop a blank one, strip a trailing slash.""" + if value is None: + return None + text = str(value).strip() + return text.rstrip("/") or None + + def _positive_number(value: object, name: str, *, allow_zero: bool = False) -> None: if ( isinstance(value, bool) @@ -35,6 +43,11 @@ class WorkerConfig: request_timeout: float = 30.0 heartbeat_interval: float = 15.0 bearer_token: str | None = None + # A long-lived per-user credential. When set (with userservice_url), the + # worker exchanges it for short-lived JWTs instead of using bearer_token, + # binding the worker to that user's account. + worker_key: str | None = None + userservice_url: str | None = None cleanup_after_seconds: float | None = None max_tasks: int | None = None exit_when_idle: bool = False @@ -54,6 +67,12 @@ class WorkerConfig: raise ValueError("coordinator_url must be an absolute HTTP(S) URL") if not isinstance(self.worker_name, str) or not self.worker_name.strip(): raise ValueError("worker_name must be non-empty") + if self.userservice_url is not None: + us = urlsplit(self.userservice_url) + if us.scheme not in {"http", "https"} or not us.hostname: + raise ValueError("userservice_url must be an absolute HTTP(S) URL") + if self.worker_key is not None and not self.userservice_url: + raise ValueError("worker_key requires userservice_url (SCIMESH_USERSERVICE_URL)") if isinstance(self.cpu_count, bool) or not isinstance(self.cpu_count, int) or self.cpu_count < 1: raise ValueError("cpu_count must be positive") if self.worker_id is not None and not isinstance(self.worker_id, str): @@ -114,6 +133,8 @@ class WorkerConfig: request_timeout=float(value("request_timeout", "SCIMESH_REQUEST_TIMEOUT", "30")), heartbeat_interval=float(value("heartbeat_interval", "SCIMESH_HEARTBEAT_INTERVAL", "15")), bearer_token=value("bearer_token", "SCIMESH_BEARER_TOKEN"), + worker_key=value("worker_key", "SCIMESH_WORKER_KEY"), + userservice_url=_clean_url(value("userservice_url", "SCIMESH_USERSERVICE_URL")), cleanup_after_seconds=float(cleanup) if cleanup else None, max_tasks=int(max_tasks) if max_tasks is not None else None, exit_when_idle=bool(values.get("exit_when_idle", False)), diff --git a/scimesh/worker/coordinator.py b/scimesh/worker/coordinator.py index 7b575b7..4367232 100644 --- a/scimesh/worker/coordinator.py +++ b/scimesh/worker/coordinator.py @@ -7,6 +7,7 @@ from typing import Any, Protocol from urllib.error import HTTPError, URLError from urllib.request import Request, build_opener +from .auth import StaticTokenProvider, TokenProvider from .models import ClaimedTask, RegisteredWorker from .transport import NoRedirectHandler @@ -38,12 +39,25 @@ class CoordinatorClient(Protocol): class HttpCoordinatorClient: - def __init__(self, base_url: str, timeout: float, bearer_token: str | None = None) -> None: + def __init__( + self, + base_url: str, + timeout: float, + bearer_token: str | None = None, + *, + token_provider: TokenProvider | None = None, + ) -> None: self.base_url = base_url.rstrip("/") self.timeout = timeout - self.bearer_token = bearer_token + # A bearer_token argument keeps older call sites working; internally + # everything goes through a provider so refresh is uniform. + self._tokens: TokenProvider = token_provider or StaticTokenProvider(bearer_token) self._opener = build_opener(NoRedirectHandler()) + @property + def bearer_token(self) -> str | None: + return self._tokens.token() + def register( self, name: str, capabilities: tuple[str, ...], cpu_count: int, memory_mb: int | None ) -> RegisteredWorker: @@ -102,6 +116,11 @@ class HttpCoordinatorClient: return lease_expires_at def _request(self, method: str, path: str, payload: dict[str, Any]) -> tuple[int, dict[str, Any]]: + return self._request_once(method, path, payload, allow_refresh=True) + + def _request_once( + self, method: str, path: str, payload: dict[str, Any], *, allow_refresh: bool + ) -> tuple[int, dict[str, Any]]: request = Request( f"{self.base_url}{path}", data=json.dumps(payload).encode(), method=method, headers={"Content-Type": "application/json", **self._auth_header()}, @@ -114,6 +133,11 @@ class HttpCoordinatorClient: except json.JSONDecodeError as error: raise CoordinatorError("coordinator returned invalid JSON") from error except HTTPError as error: + # A 401 usually means the short-lived JWT expired; mint a fresh one + # and retry exactly once so an in-flight worker rides over the gap. + if error.code == 401 and allow_refresh: + self._tokens.refresh() + return self._request_once(method, path, payload, allow_refresh=False) if error.code >= 500: raise CoordinatorTransientError(f"coordinator returned {error.code}") from error return error.code, {} @@ -121,4 +145,5 @@ class HttpCoordinatorClient: raise CoordinatorTransientError("coordinator request failed") from error def _auth_header(self) -> dict[str, str]: - return {"Authorization": f"Bearer {self.bearer_token}"} if self.bearer_token else {} + token = self._tokens.token() + return {"Authorization": f"Bearer {token}"} if token else {} diff --git a/tests/test_worker_auth.py b/tests/test_worker_auth.py new file mode 100644 index 0000000..3d5131d --- /dev/null +++ b/tests/test_worker_auth.py @@ -0,0 +1,205 @@ +"""Tests for worker token strategies and the client's 401 refresh.""" + +from __future__ import annotations + +import json +from pathlib import Path +from urllib.error import HTTPError + +import pytest + +from scimesh.worker.auth import ( + StaticTokenProvider, + TokenExchangeError, + WorkerKeyTokenProvider, + provider_from_config, +) +from scimesh.worker.config import WorkerConfig +from scimesh.worker.coordinator import HttpCoordinatorClient + + +class FakeResponse: + def __init__(self, status: int, body: bytes) -> None: + self.status = status + self._body = body + + def read(self) -> bytes: + return self._body + + def __enter__(self) -> "FakeResponse": + return self + + def __exit__(self, *exc) -> bool: + return False + + +class SeqOpener: + """Returns/raises a scripted sequence of responses, recording each request.""" + + def __init__(self, actions: list) -> None: + self.actions = list(actions) + self.requests: list = [] + + def open(self, request, timeout=None): + self.requests.append(request) + action = self.actions.pop(0) + if isinstance(action, Exception): + raise action + return action + + +def _exchange_response(token: str, expires_in: int) -> FakeResponse: + return FakeResponse(200, json.dumps({"token": token, "expires_in": expires_in}).encode()) + + +def test_static_provider_returns_fixed_token_and_never_refreshes(): + provider = StaticTokenProvider("tok") + assert provider.token() == "tok" + provider.refresh() + assert provider.token() == "tok" + + +def test_static_provider_none_means_no_auth(): + assert StaticTokenProvider(None).token() is None + + +def test_worker_key_provider_exchanges_once_then_caches(): + clock = {"t": 1000.0} + provider = WorkerKeyTokenProvider( + "http://users", "scimesh_wk_live_x", timeout=5, now=lambda: clock["t"] + ) + provider._opener = SeqOpener([_exchange_response("jwt-1", 100)]) + + # First call exchanges; a second call well within the TTL reuses the cache. + assert provider.token() == "jwt-1" + clock["t"] = 1050.0 # 50s later, TTL 100s with 0.2 leeway → refresh at +80s + assert provider.token() == "jwt-1" + assert len(provider._opener.requests) == 1 + + +def test_worker_key_provider_refreshes_after_leeway(): + clock = {"t": 0.0} + provider = WorkerKeyTokenProvider( + "http://users", "k", timeout=5, now=lambda: clock["t"] + ) + provider._opener = SeqOpener([ + _exchange_response("jwt-1", 100), + _exchange_response("jwt-2", 100), + ]) + assert provider.token() == "jwt-1" + clock["t"] = 85.0 # past the 80s refresh point + assert provider.token() == "jwt-2" + assert len(provider._opener.requests) == 2 + + +def test_worker_key_provider_force_refresh(): + provider = WorkerKeyTokenProvider("http://users", "k", timeout=5, now=lambda: 0.0) + provider._opener = SeqOpener([ + _exchange_response("jwt-1", 100), + _exchange_response("jwt-2", 100), + ]) + assert provider.token() == "jwt-1" + provider.refresh() + assert provider.token() == "jwt-2" + + +def test_worker_key_provider_raises_on_rejected_key(): + provider = WorkerKeyTokenProvider("http://users", "bad", timeout=5, now=lambda: 0.0) + provider._opener = SeqOpener([HTTPError("http://users", 401, "unauthorized", {}, None)]) + with pytest.raises(TokenExchangeError): + provider.token() + + +def test_worker_key_provider_raises_when_token_missing(): + provider = WorkerKeyTokenProvider("http://users", "k", timeout=5, now=lambda: 0.0) + provider._opener = SeqOpener([FakeResponse(200, json.dumps({"expires_in": 100}).encode())]) + with pytest.raises(TokenExchangeError): + provider.token() + + +def test_provider_from_config_selects_worker_key_mode(): + provider = provider_from_config( + worker_key="scimesh_wk_live_x", + userservice_url="http://users", + bearer_token="ignored", + request_timeout=5, + ) + assert isinstance(provider, WorkerKeyTokenProvider) + + +def test_provider_from_config_falls_back_to_static(): + provider = provider_from_config( + worker_key=None, userservice_url=None, bearer_token="tok", request_timeout=5 + ) + assert isinstance(provider, StaticTokenProvider) + assert provider.token() == "tok" + + +class RefreshCountingProvider: + def __init__(self) -> None: + self.tokens = ["stale", "fresh"] + self.index = 0 + self.refreshes = 0 + + def token(self) -> str: + return self.tokens[min(self.index, len(self.tokens) - 1)] + + def refresh(self) -> None: + self.refreshes += 1 + self.index += 1 + + +def test_coordinator_client_refreshes_and_retries_once_on_401(): + provider = RefreshCountingProvider() + client = HttpCoordinatorClient("http://coord", timeout=5, token_provider=provider) + client._opener = SeqOpener([ + HTTPError("http://coord/tasks/claim", 401, "unauthorized", {}, None), + FakeResponse(204, b""), + ]) + + status, _ = client._request("POST", "/tasks/claim", {"worker_id": "w"}) + + assert status == 204 + assert provider.refreshes == 1 + # The retry carried the refreshed token. + assert provider.index == 1 + + +def test_coordinator_client_does_not_loop_on_persistent_401(): + provider = RefreshCountingProvider() + client = HttpCoordinatorClient("http://coord", timeout=5, token_provider=provider) + client._opener = SeqOpener([ + HTTPError("http://coord/x", 401, "unauthorized", {}, None), + HTTPError("http://coord/x", 401, "unauthorized", {}, None), + ]) + + status, _ = client._request("POST", "/x", {}) + + # One refresh, one retry, then the second 401 is surfaced rather than retried. + assert status == 401 + assert provider.refreshes == 1 + + +def _base_config(**extra) -> dict: + return { + "coordinator_url": "http://coord", + "worker_id": None, + "work_dir": Path("."), + **extra, + } + + +def test_worker_key_requires_userservice_url(): + with pytest.raises(ValueError, match="userservice_url"): + WorkerConfig(**_base_config(worker_key="scimesh_wk_live_x")) + + +def test_worker_key_with_userservice_url_is_valid(): + cfg = WorkerConfig(**_base_config(worker_key="scimesh_wk_live_x", userservice_url="http://users")) + assert cfg.worker_key == "scimesh_wk_live_x" + assert cfg.userservice_url == "http://users" + + +def test_userservice_url_must_be_absolute(): + with pytest.raises(ValueError, match="userservice_url"): + WorkerConfig(**_base_config(userservice_url="not-a-url")) diff --git a/users/cmd/userservice/main.go b/users/cmd/userservice/main.go index 97f55f4..9942f9f 100644 --- a/users/cmd/userservice/main.go +++ b/users/cmd/userservice/main.go @@ -49,16 +49,21 @@ func run() error { // Adapters implementing the usecase ports. users := postgres.NewUserRepo(pool) + workerKeys := postgres.NewWorkerKeyRepo(pool) hasher := auth.NewHasher(cfg.BcryptCost) clock := infra.NewClock() issuer := auth.NewIssuer(cfg.JWTSecret, cfg.TokenTTL, clock.Now) uc := apihttp.UseCases{ - Register: usecase.NewRegister(users, hasher, clock), - Login: usecase.NewLogin(users, hasher, issuer), - SetVerified: usecase.NewSetVerified(users), - SetRole: usecase.NewSetRole(users), - Users: users, + Register: usecase.NewRegister(users, hasher, clock), + Login: usecase.NewLogin(users, hasher, issuer), + SetVerified: usecase.NewSetVerified(users), + SetRole: usecase.NewSetRole(users), + CreateWorkerKey: usecase.NewCreateWorkerKey(workerKeys, clock), + ListWorkerKeys: usecase.NewListWorkerKeys(workerKeys), + RevokeWorkerKey: usecase.NewRevokeWorkerKey(workerKeys), + ExchangeWorkerKey: usecase.NewExchangeWorkerKey(workerKeys, users, issuer, cfg.TokenTTL), + Users: users, } // Seed the first admin, if configured. Idempotent: a no-op once it exists. diff --git a/users/internal/domain/errors.go b/users/internal/domain/errors.go index 39bd39a..05de0a8 100644 --- a/users/internal/domain/errors.go +++ b/users/internal/domain/errors.go @@ -8,4 +8,6 @@ var ( ErrEmptyEmail = errors.New("email is required") ErrInvalidEmail = errors.New("email is not a valid address") ErrEmptyPasswordHash = errors.New("password hash is required") + + ErrWorkerKeyNameTooLong = errors.New("worker key name is too long") ) diff --git a/users/internal/domain/workerkey.go b/users/internal/domain/workerkey.go new file mode 100644 index 0000000..f0f0024 --- /dev/null +++ b/users/internal/domain/workerkey.go @@ -0,0 +1,84 @@ +package domain + +import ( + "crypto/rand" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "strings" + "time" + + "github.com/google/uuid" +) + +const ( + // workerKeyLabel makes a key self-describing when it turns up in a log or an + // env var, and lets a client sanity-check the shape before exchanging it. + workerKeyLabel = "scimesh_wk_live_" + // workerKeyRandomBytes is the entropy behind the secret. 24 bytes (192 bits) + // is far beyond guessable, which is why the stored hash needs no salt. + workerKeyRandomBytes = 24 + // workerKeyPrefixChars is how much of the random tail we keep, alongside the + // label, as the non-secret identifier shown in the UI. + workerKeyPrefixChars = 8 + // workerKeyNameMax caps the user-supplied label. + workerKeyNameMax = 100 + // workerKeyDefaultName is used when the caller supplies no label. + workerKeyDefaultName = "my machine" +) + +// WorkerKey is a long-lived, per-user credential for running a worker. The +// secret itself is never stored — only TokenHash — so the plaintext returned by +// NewWorkerKey is the one and only chance to show it to the user. +type WorkerKey struct { + ID uuid.UUID + UserID uuid.UUID + Name string + TokenHash string + Prefix string + CreatedAt time.Time + LastUsedAt *time.Time + RevokedAt *time.Time +} + +// NewWorkerKey mints a key for a user and returns both the entity (carrying only +// the hash) and the one-time plaintext to hand back to the caller. The label is +// trimmed and defaulted; an over-long one is rejected. +func NewWorkerKey(userID uuid.UUID, name string, now time.Time) (*WorkerKey, string, error) { + name = strings.TrimSpace(name) + if name == "" { + name = workerKeyDefaultName + } + if len(name) > workerKeyNameMax { + return nil, "", ErrWorkerKeyNameTooLong + } + + b := make([]byte, workerKeyRandomBytes) + if _, err := rand.Read(b); err != nil { + return nil, "", err + } + // URL-safe, unpadded: the key rides in env vars and shell commands, so it + // must contain no '=', '+', or '/' that a shell might mangle. + raw := workerKeyLabel + base64.RawURLEncoding.EncodeToString(b) + + key := &WorkerKey{ + ID: uuid.New(), + UserID: userID, + Name: name, + TokenHash: HashWorkerKey(raw), + Prefix: raw[:len(workerKeyLabel)+workerKeyPrefixChars], + CreatedAt: now, + } + return key, raw, nil +} + +// HashWorkerKey returns the hex SHA-256 of a presented key. Exchange hashes the +// incoming key the same way and looks the row up by it, so the plaintext never +// has to be compared directly. +func HashWorkerKey(raw string) string { + sum := sha256.Sum256([]byte(raw)) + return hex.EncodeToString(sum[:]) +} + +// Revoked reports whether the key has been retired and must no longer exchange. +func (k *WorkerKey) Revoked() bool { return k.RevokedAt != nil } diff --git a/users/internal/domain/workerkey_test.go b/users/internal/domain/workerkey_test.go new file mode 100644 index 0000000..ff4eeee --- /dev/null +++ b/users/internal/domain/workerkey_test.go @@ -0,0 +1,61 @@ +package domain_test + +import ( + "strings" + "testing" + "time" + + "github.com/google/uuid" + + "github.com/emil28092005/SciMesh/users/internal/domain" +) + +func TestNewWorkerKeyShape(t *testing.T) { + owner := uuid.New() + now := time.Date(2026, 7, 26, 0, 0, 0, 0, time.UTC) + + key, raw, err := domain.NewWorkerKey(owner, "home-desktop", now) + if err != nil { + t.Fatalf("NewWorkerKey: %v", err) + } + if !strings.HasPrefix(raw, "scimesh_wk_live_") { + t.Errorf("raw key has no recognisable label: %q", raw) + } + if key.TokenHash != domain.HashWorkerKey(raw) { + t.Error("stored hash does not match the plaintext") + } + if key.TokenHash == raw || strings.Contains(key.TokenHash, raw) { + t.Error("plaintext leaked into the stored hash") + } + if !strings.HasPrefix(raw, key.Prefix) { + t.Errorf("prefix %q is not a leading slice of the key", key.Prefix) + } + if key.UserID != owner || key.CreatedAt != now || key.Revoked() { + t.Errorf("unexpected key metadata: %+v", key) + } +} + +func TestNewWorkerKeyDefaultsBlankName(t *testing.T) { + key, _, err := domain.NewWorkerKey(uuid.New(), " ", time.Now()) + if err != nil { + t.Fatalf("NewWorkerKey: %v", err) + } + if key.Name == "" { + t.Error("blank name was not defaulted") + } +} + +func TestNewWorkerKeyRejectsLongName(t *testing.T) { + _, _, err := domain.NewWorkerKey(uuid.New(), strings.Repeat("x", 101), time.Now()) + if err != domain.ErrWorkerKeyNameTooLong { + t.Errorf("got %v, want ErrWorkerKeyNameTooLong", err) + } +} + +func TestNewWorkerKeyUniquePerCall(t *testing.T) { + a, rawA, _ := domain.NewWorkerKey(uuid.New(), "a", time.Now()) + b, rawB, _ := domain.NewWorkerKey(uuid.New(), "b", time.Now()) + if rawA == rawB || a.TokenHash == b.TokenHash || a.ID == b.ID { + t.Error("two keys collided; generation is not random") + } +} diff --git a/users/internal/storage/postgres/workerkey_repo.go b/users/internal/storage/postgres/workerkey_repo.go new file mode 100644 index 0000000..09526b6 --- /dev/null +++ b/users/internal/storage/postgres/workerkey_repo.go @@ -0,0 +1,123 @@ +package postgres + +import ( + "context" + "errors" + + sq "github.com/Masterminds/squirrel" + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + + "github.com/emil28092005/SciMesh/users/internal/domain" + "github.com/emil28092005/SciMesh/users/internal/usecase" +) + +var workerKeyColumns = []string{ + "id", "user_id", "name", "token_hash", "prefix", "created_at", "last_used_at", "revoked_at", +} + +// WorkerKeyRepo implements usecase.WorkerKeyRepository on PostgreSQL. +type WorkerKeyRepo struct { + pool *pgxpool.Pool +} + +func NewWorkerKeyRepo(pool *pgxpool.Pool) *WorkerKeyRepo { + return &WorkerKeyRepo{pool: pool} +} + +func (r *WorkerKeyRepo) Insert(ctx context.Context, k *domain.WorkerKey) error { + sql, args, err := psql.Insert("worker_keys"). + Columns(workerKeyColumns...). + Values(k.ID, k.UserID, k.Name, k.TokenHash, k.Prefix, k.CreatedAt, k.LastUsedAt, k.RevokedAt). + ToSql() + if err != nil { + return err + } + _, err = conn(ctx, r.pool).Exec(ctx, sql, args...) + return err +} + +func (r *WorkerKeyRepo) ListByUser(ctx context.Context, userID uuid.UUID) ([]*domain.WorkerKey, error) { + sql, args, err := psql.Select(workerKeyColumns...). + From("worker_keys"). + Where(sq.Eq{"user_id": userID, "revoked_at": nil}). + OrderBy("created_at DESC"). + ToSql() + if err != nil { + return nil, err + } + rows, err := conn(ctx, r.pool).Query(ctx, sql, args...) + if err != nil { + return nil, err + } + defer rows.Close() + + keys := []*domain.WorkerKey{} + for rows.Next() { + k, err := scanWorkerKey(rows) + if err != nil { + return nil, err + } + keys = append(keys, k) + } + return keys, rows.Err() +} + +func (r *WorkerKeyRepo) GetActiveByHash(ctx context.Context, tokenHash string) (*domain.WorkerKey, error) { + sql, args, err := psql.Select(workerKeyColumns...). + From("worker_keys"). + Where(sq.Eq{"token_hash": tokenHash, "revoked_at": nil}). + ToSql() + if err != nil { + return nil, err + } + return scanWorkerKey(conn(ctx, r.pool).QueryRow(ctx, sql, args...)) +} + +// Revoke retires a live key the user owns. Scoping the UPDATE to both id and +// user_id means one user can never revoke another's key, and the revoked_at IS +// NULL guard makes a double-revoke a clean 404 rather than a silent success. +func (r *WorkerKeyRepo) Revoke(ctx context.Context, id, userID uuid.UUID) error { + sql, args, err := psql.Update("worker_keys"). + Set("revoked_at", sq.Expr("now()")). + Where(sq.Eq{"id": id, "user_id": userID, "revoked_at": nil}). + ToSql() + if err != nil { + return err + } + tag, err := conn(ctx, r.pool).Exec(ctx, sql, args...) + if err != nil { + return err + } + if tag.RowsAffected() == 0 { + return usecase.ErrWorkerKeyNotFound + } + return nil +} + +func (r *WorkerKeyRepo) TouchLastUsed(ctx context.Context, id uuid.UUID) error { + sql, args, err := psql.Update("worker_keys"). + Set("last_used_at", sq.Expr("now()")). + Where(sq.Eq{"id": id}). + ToSql() + if err != nil { + return err + } + _, err = conn(ctx, r.pool).Exec(ctx, sql, args...) + return err +} + +func scanWorkerKey(row pgx.Row) (*domain.WorkerKey, error) { + var k domain.WorkerKey + if err := row.Scan( + &k.ID, &k.UserID, &k.Name, &k.TokenHash, &k.Prefix, + &k.CreatedAt, &k.LastUsedAt, &k.RevokedAt, + ); err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return nil, usecase.ErrWorkerKeyNotFound + } + return nil, err + } + return &k, nil +} diff --git a/users/internal/transport/http/dto.go b/users/internal/transport/http/dto.go index 808ee3d..65355b1 100644 --- a/users/internal/transport/http/dto.go +++ b/users/internal/transport/http/dto.go @@ -41,3 +41,53 @@ func toUserResponse(u *domain.User) userResponse { CreatedAt: u.CreatedAt.UTC().Format(time.RFC3339), } } + +// createWorkerKeyRequest is the body for minting a worker key. Name is an +// optional human label; the domain defaults it when blank. +type createWorkerKeyRequest struct { + Name string `json:"name"` +} + +// exchangeWorkerKeyRequest trades a worker key for a short-lived JWT. +type exchangeWorkerKeyRequest struct { + Key string `json:"key"` +} + +type exchangeWorkerKeyResponse struct { + Token string `json:"token"` + ExpiresIn int `json:"expires_in"` +} + +// workerKeyResponse is the public view of a key. It never carries the secret — +// only the non-secret prefix used to identify a row. +type workerKeyResponse struct { + ID string `json:"id"` + Name string `json:"name"` + Prefix string `json:"prefix"` + CreatedAt string `json:"created_at"` + LastUsedAt string `json:"last_used_at,omitempty"` +} + +// createdWorkerKeyResponse extends the public view with the one-time plaintext, +// returned only from the create call and never again. +type createdWorkerKeyResponse struct { + workerKeyResponse + Key string `json:"key"` +} + +type workerKeysResponse struct { + WorkerKeys []workerKeyResponse `json:"worker_keys"` +} + +func toWorkerKeyResponse(k *domain.WorkerKey) workerKeyResponse { + resp := workerKeyResponse{ + ID: k.ID.String(), + Name: k.Name, + Prefix: k.Prefix, + CreatedAt: k.CreatedAt.UTC().Format(time.RFC3339), + } + if k.LastUsedAt != nil { + resp.LastUsedAt = k.LastUsedAt.UTC().Format(time.RFC3339) + } + return resp +} diff --git a/users/internal/transport/http/errors.go b/users/internal/transport/http/errors.go index 8ccbe7e..ac03627 100644 --- a/users/internal/transport/http/errors.go +++ b/users/internal/transport/http/errors.go @@ -48,6 +48,12 @@ func statusForError(err error) (int, string) { return http.StatusUnauthorized, "invalid email or password" case errors.Is(err, usecase.ErrUserNotFound): return http.StatusNotFound, "user not found" + case errors.Is(err, usecase.ErrWorkerKeyNotFound): + return http.StatusNotFound, "worker key not found" + case errors.Is(err, usecase.ErrInvalidWorkerKey): + return http.StatusUnauthorized, "invalid worker key" + case errors.Is(err, domain.ErrWorkerKeyNameTooLong): + return http.StatusBadRequest, "worker key name is too long" case errors.Is(err, usecase.ErrPasswordTooShort): return http.StatusBadRequest, "password must be at least 8 characters" case errors.Is(err, usecase.ErrPasswordTooLong): diff --git a/users/internal/transport/http/handlers.go b/users/internal/transport/http/handlers.go index beaca96..b766601 100644 --- a/users/internal/transport/http/handlers.go +++ b/users/internal/transport/http/handlers.go @@ -13,12 +13,16 @@ import ( // Handlers holds the use cases each endpoint drives. type Handlers struct { - register *usecase.Register - login *usecase.Login - setVerified *usecase.SetVerified - setRole *usecase.SetRole - users usecase.UserRepository - log *slog.Logger + register *usecase.Register + login *usecase.Login + setVerified *usecase.SetVerified + setRole *usecase.SetRole + createWorkerKey *usecase.CreateWorkerKey + listWorkerKeys *usecase.ListWorkerKeys + revokeWorkerKey *usecase.RevokeWorkerKey + exchangeWorkerKey *usecase.ExchangeWorkerKey + users usecase.UserRepository + log *slog.Logger } // handleHealth is an unauthenticated liveness probe for the container and load @@ -113,6 +117,90 @@ func (h *Handlers) handleSetRole(role domain.Role) http.HandlerFunc { } } +// handleCreateWorkerKey mints a long-lived worker key for the authenticated +// caller and returns it once, plaintext included. The user copies it into their +// worker's SCIMESH_WORKER_KEY; it is never retrievable again. +func (h *Handlers) handleCreateWorkerKey(w http.ResponseWriter, r *http.Request) { + id, ok := userIDFrom(r.Context()) + if !ok { + unauthorized(w, r) + return + } + var req createWorkerKeyRequest + if !decodeJSON(w, r, &req) { + return + } + key, raw, err := h.createWorkerKey.Execute(r.Context(), id, req.Name) + if err != nil { + writeError(w, r, h.log, err) + return + } + writeJSON(w, http.StatusCreated, createdWorkerKeyResponse{ + workerKeyResponse: toWorkerKeyResponse(key), + Key: raw, + }) +} + +// handleListWorkerKeys returns the caller's live keys (no secrets) for display +// and revocation. +func (h *Handlers) handleListWorkerKeys(w http.ResponseWriter, r *http.Request) { + id, ok := userIDFrom(r.Context()) + if !ok { + unauthorized(w, r) + return + } + keys, err := h.listWorkerKeys.Execute(r.Context(), id) + if err != nil { + writeError(w, r, h.log, err) + return + } + out := make([]workerKeyResponse, 0, len(keys)) + for _, k := range keys { + out = append(out, toWorkerKeyResponse(k)) + } + writeJSON(w, http.StatusOK, workerKeysResponse{WorkerKeys: out}) +} + +// handleRevokeWorkerKey retires one of the caller's keys. The repository scopes +// the delete to the owner, so a mismatched id is a clean 404, not another user's +// key. +func (h *Handlers) handleRevokeWorkerKey(w http.ResponseWriter, r *http.Request) { + userID, ok := userIDFrom(r.Context()) + if !ok { + unauthorized(w, r) + return + } + keyID, err := uuid.Parse(r.PathValue("id")) + if err != nil { + writeJSON(w, http.StatusBadRequest, errorResponse{ + Error: "invalid worker key id", + RequestID: requestIDFrom(r.Context()), + }) + return + } + if err := h.revokeWorkerKey.Execute(r.Context(), userID, keyID); err != nil { + writeError(w, r, h.log, err) + return + } + w.WriteHeader(http.StatusNoContent) +} + +// handleExchangeWorkerKey trades a worker key for a short-lived JWT. It is +// unauthenticated: the key itself is the credential. A worker calls this on +// startup and again to refresh before the JWT expires. +func (h *Handlers) handleExchangeWorkerKey(w http.ResponseWriter, r *http.Request) { + var req exchangeWorkerKeyRequest + if !decodeJSON(w, r, &req) { + return + } + token, expiresIn, err := h.exchangeWorkerKey.Execute(r.Context(), req.Key) + if err != nil { + writeError(w, r, h.log, err) + return + } + writeJSON(w, http.StatusOK, exchangeWorkerKeyResponse{Token: token, ExpiresIn: expiresIn}) +} + // decodeJSON reads a size-capped JSON body into dst, rejecting unknown fields. // It writes a 400 and returns false on any problem, so callers can `if // !decodeJSON(...) { return }`. diff --git a/users/internal/transport/http/server.go b/users/internal/transport/http/server.go index 3a57af7..fe28394 100644 --- a/users/internal/transport/http/server.go +++ b/users/internal/transport/http/server.go @@ -14,23 +14,31 @@ import ( // UseCases bundles the application services the handlers drive. type UseCases struct { - Register *usecase.Register - Login *usecase.Login - SetVerified *usecase.SetVerified - SetRole *usecase.SetRole - Users usecase.UserRepository + Register *usecase.Register + Login *usecase.Login + SetVerified *usecase.SetVerified + SetRole *usecase.SetRole + CreateWorkerKey *usecase.CreateWorkerKey + ListWorkerKeys *usecase.ListWorkerKeys + RevokeWorkerKey *usecase.RevokeWorkerKey + ExchangeWorkerKey *usecase.ExchangeWorkerKey + Users usecase.UserRepository } // NewServer wires the routes and the middleware stack and returns the handler. // The issuer verifies tokens for the JWT-protected routes. func NewServer(log *slog.Logger, uc UseCases, issuer auth.Issuer) http.Handler { h := &Handlers{ - register: uc.Register, - login: uc.Login, - setVerified: uc.SetVerified, - setRole: uc.SetRole, - users: uc.Users, - log: log, + register: uc.Register, + login: uc.Login, + setVerified: uc.SetVerified, + setRole: uc.SetRole, + createWorkerKey: uc.CreateWorkerKey, + listWorkerKeys: uc.ListWorkerKeys, + revokeWorkerKey: uc.RevokeWorkerKey, + exchangeWorkerKey: uc.ExchangeWorkerKey, + users: uc.Users, + log: log, } mux := http.NewServeMux() @@ -41,6 +49,14 @@ func NewServer(log *slog.Logger, uc UseCases, issuer auth.Issuer) http.Handler { // /me proves a token round-trips; it sits behind JWT auth. mux.Handle("GET /me", chain(http.HandlerFunc(h.handleMe), withJWT(issuer))) + // Worker keys: a user mints a long-lived key (JWT-protected), and a worker + // trades it for a short-lived JWT on the public exchange endpoint — the key + // itself is the credential there, so no prior token is required. + mux.HandleFunc("POST /worker-tokens/exchange", h.handleExchangeWorkerKey) + mux.Handle("POST /worker-keys", chain(http.HandlerFunc(h.handleCreateWorkerKey), withJWT(issuer))) + mux.Handle("GET /worker-keys", chain(http.HandlerFunc(h.handleListWorkerKeys), withJWT(issuer))) + mux.Handle("DELETE /worker-keys/{id}", chain(http.HandlerFunc(h.handleRevokeWorkerKey), withJWT(issuer))) + // Admin-only: grant or revoke the trusted-contributor badge. withAdmin sits // inside withJWT so the role is available from the verified token. mux.Handle("POST /users/{id}/verify", diff --git a/users/internal/usecase/errors.go b/users/internal/usecase/errors.go index 9a0b3e2..bc6978d 100644 --- a/users/internal/usecase/errors.go +++ b/users/internal/usecase/errors.go @@ -7,6 +7,14 @@ var ( ErrEmailExists = errors.New("email already registered") ErrUserNotFound = errors.New("user not found") + // ErrWorkerKeyNotFound is returned by WorkerKeyRepository when no live key + // matches (by id for revoke, by hash for exchange). + ErrWorkerKeyNotFound = errors.New("worker key not found") + // ErrInvalidWorkerKey is surfaced to the transport layer for a key that does + // not exchange (unknown, revoked, or owner gone). Deliberately opaque so a + // caller cannot distinguish the cases while probing. + ErrInvalidWorkerKey = errors.New("invalid worker key") + // Use-case errors surfaced to the transport layer. // // ErrInvalidCredentials is deliberately returned for both an unknown email diff --git a/users/internal/usecase/ports.go b/users/internal/usecase/ports.go index d94bfd1..09c1581 100644 --- a/users/internal/usecase/ports.go +++ b/users/internal/usecase/ports.go @@ -30,6 +30,25 @@ type UserRepository interface { SetRole(ctx context.Context, id uuid.UUID, role domain.Role) error } +// WorkerKeyRepository persists and looks up the long-lived worker keys a user +// creates to run a worker bound to their account. Implementations return the +// sentinel errors in errors.go so the use cases stay free of SQL types. +type WorkerKeyRepository interface { + // Insert stores a freshly minted key. + Insert(ctx context.Context, k *domain.WorkerKey) error + // ListByUser returns a user's live (non-revoked) keys, newest first. + ListByUser(ctx context.Context, userID uuid.UUID) ([]*domain.WorkerKey, error) + // GetActiveByHash returns the non-revoked key with the given hash, or + // ErrWorkerKeyNotFound. + GetActiveByHash(ctx context.Context, tokenHash string) (*domain.WorkerKey, error) + // Revoke retires a key the user owns, returning ErrWorkerKeyNotFound when no + // live key with that id belongs to the user. + Revoke(ctx context.Context, id, userID uuid.UUID) error + // TouchLastUsed records a successful exchange. Best-effort: a failure here + // must not fail the exchange itself. + TouchLastUsed(ctx context.Context, id uuid.UUID) error +} + // PasswordHasher hashes and verifies passwords. The bcrypt adapter satisfies it. type PasswordHasher interface { Hash(password string) (string, error) diff --git a/users/internal/usecase/workerkey.go b/users/internal/usecase/workerkey.go new file mode 100644 index 0000000..6f28bb9 --- /dev/null +++ b/users/internal/usecase/workerkey.go @@ -0,0 +1,112 @@ +package usecase + +import ( + "context" + "errors" + "time" + + "github.com/google/uuid" + + "github.com/emil28092005/SciMesh/users/internal/domain" +) + +// CreateWorkerKey mints a long-lived worker key for a user and returns the +// one-time plaintext to show once. +type CreateWorkerKey struct { + keys WorkerKeyRepository + clock Clock +} + +func NewCreateWorkerKey(keys WorkerKeyRepository, clock Clock) *CreateWorkerKey { + return &CreateWorkerKey{keys: keys, clock: clock} +} + +// Execute returns the stored key (hash only) and the plaintext secret. The +// secret is never persisted, so this is the sole moment it can be surfaced. +func (uc *CreateWorkerKey) Execute(ctx context.Context, userID uuid.UUID, name string) (*domain.WorkerKey, string, error) { + key, raw, err := domain.NewWorkerKey(userID, name, uc.clock.Now()) + if err != nil { + return nil, "", err + } + if err := uc.keys.Insert(ctx, key); err != nil { + return nil, "", err + } + return key, raw, nil +} + +// ListWorkerKeys returns a user's live keys for display and management. +type ListWorkerKeys struct { + keys WorkerKeyRepository +} + +func NewListWorkerKeys(keys WorkerKeyRepository) *ListWorkerKeys { + return &ListWorkerKeys{keys: keys} +} + +func (uc *ListWorkerKeys) Execute(ctx context.Context, userID uuid.UUID) ([]*domain.WorkerKey, error) { + return uc.keys.ListByUser(ctx, userID) +} + +// RevokeWorkerKey retires one of the caller's keys. +type RevokeWorkerKey struct { + keys WorkerKeyRepository +} + +func NewRevokeWorkerKey(keys WorkerKeyRepository) *RevokeWorkerKey { + return &RevokeWorkerKey{keys: keys} +} + +func (uc *RevokeWorkerKey) Execute(ctx context.Context, userID, id uuid.UUID) error { + return uc.keys.Revoke(ctx, id, userID) +} + +// ExchangeWorkerKey trades a valid worker key for a short-lived JWT. The JWT +// carries the owner's current role and verified flag, so a worker that refreshes +// after an admin verifies the owner picks up the upgraded trust on its next +// registration. +type ExchangeWorkerKey struct { + keys WorkerKeyRepository + users UserRepository + tokens TokenIssuer + ttl time.Duration +} + +func NewExchangeWorkerKey(keys WorkerKeyRepository, users UserRepository, tokens TokenIssuer, ttl time.Duration) *ExchangeWorkerKey { + return &ExchangeWorkerKey{keys: keys, users: users, tokens: tokens, ttl: ttl} +} + +// Execute returns a signed token and its lifetime in seconds. Every failure to +// resolve the key to a usable owner collapses to ErrInvalidWorkerKey so a caller +// cannot tell an unknown key from a revoked one or a deleted owner. +func (uc *ExchangeWorkerKey) Execute(ctx context.Context, rawKey string) (string, int, error) { + if rawKey == "" { + return "", 0, ErrInvalidWorkerKey + } + + key, err := uc.keys.GetActiveByHash(ctx, domain.HashWorkerKey(rawKey)) + if err != nil { + if errors.Is(err, ErrWorkerKeyNotFound) { + return "", 0, ErrInvalidWorkerKey + } + return "", 0, err + } + + u, err := uc.users.GetByID(ctx, key.UserID) + if err != nil { + if errors.Is(err, ErrUserNotFound) { + return "", 0, ErrInvalidWorkerKey + } + return "", 0, err + } + + token, err := uc.tokens.Issue(u) + if err != nil { + return "", 0, err + } + + // Best-effort: a failed timestamp update must not sink an otherwise valid + // exchange the worker depends on to keep running. + _ = uc.keys.TouchLastUsed(ctx, key.ID) + + return token, int(uc.ttl.Seconds()), nil +} diff --git a/users/internal/usecase/workerkey_test.go b/users/internal/usecase/workerkey_test.go new file mode 100644 index 0000000..a06eb9a --- /dev/null +++ b/users/internal/usecase/workerkey_test.go @@ -0,0 +1,186 @@ +package usecase_test + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/google/uuid" + + "github.com/emil28092005/SciMesh/users/internal/auth" + "github.com/emil28092005/SciMesh/users/internal/domain" + "github.com/emil28092005/SciMesh/users/internal/memstore" + "github.com/emil28092005/SciMesh/users/internal/usecase" +) + +// fakeKeyRepo is an in-memory WorkerKeyRepository for the use-case tests. +type fakeKeyRepo struct { + byHash map[string]*domain.WorkerKey + byID map[uuid.UUID]*domain.WorkerKey + touched []uuid.UUID +} + +func newFakeKeyRepo() *fakeKeyRepo { + return &fakeKeyRepo{byHash: map[string]*domain.WorkerKey{}, byID: map[uuid.UUID]*domain.WorkerKey{}} +} + +func (r *fakeKeyRepo) Insert(_ context.Context, k *domain.WorkerKey) error { + r.byHash[k.TokenHash] = k + r.byID[k.ID] = k + return nil +} + +func (r *fakeKeyRepo) ListByUser(_ context.Context, userID uuid.UUID) ([]*domain.WorkerKey, error) { + var out []*domain.WorkerKey + for _, k := range r.byID { + if k.UserID == userID && !k.Revoked() { + out = append(out, k) + } + } + return out, nil +} + +func (r *fakeKeyRepo) GetActiveByHash(_ context.Context, hash string) (*domain.WorkerKey, error) { + k, ok := r.byHash[hash] + if !ok || k.Revoked() { + return nil, usecase.ErrWorkerKeyNotFound + } + return k, nil +} + +func (r *fakeKeyRepo) Revoke(_ context.Context, id, userID uuid.UUID) error { + k, ok := r.byID[id] + if !ok || k.UserID != userID || k.Revoked() { + return usecase.ErrWorkerKeyNotFound + } + now := time.Now() + k.RevokedAt = &now + return nil +} + +func (r *fakeKeyRepo) TouchLastUsed(_ context.Context, id uuid.UUID) error { + r.touched = append(r.touched, id) + return nil +} + +func newKeyFixtures(t *testing.T) (*usecase.CreateWorkerKey, *usecase.ExchangeWorkerKey, *usecase.RevokeWorkerKey, *usecase.ListWorkerKeys, *fakeKeyRepo, *domain.User) { + t.Helper() + users := memstore.NewUserRepo() + hasher := auth.NewHasher(4) + clk := memstore.Clock{T: time.Date(2026, 7, 26, 0, 0, 0, 0, time.UTC)} + issuer := auth.NewIssuer(secret, time.Hour, nil) + keys := newFakeKeyRepo() + + u, err := usecase.NewRegister(users, hasher, clk).Execute(context.Background(), "worker@example.com", "password123") + if err != nil { + t.Fatalf("seed user: %v", err) + } + + return usecase.NewCreateWorkerKey(keys, clk), + usecase.NewExchangeWorkerKey(keys, users, issuer, time.Hour), + usecase.NewRevokeWorkerKey(keys), + usecase.NewListWorkerKeys(keys), + keys, u +} + +func TestCreateAndExchangeWorkerKey(t *testing.T) { + create, exchange, _, _, keys, u := newKeyFixtures(t) + ctx := context.Background() + + key, raw, err := create.Execute(ctx, u.ID, "home-desktop") + if err != nil { + t.Fatalf("create: %v", err) + } + if key.Name != "home-desktop" || raw == "" { + t.Fatalf("unexpected key %+v raw=%q", key, raw) + } + + token, expiresIn, err := exchange.Execute(ctx, raw) + if err != nil { + t.Fatalf("exchange: %v", err) + } + if expiresIn != int((time.Hour).Seconds()) { + t.Errorf("expires_in = %d, want 3600", expiresIn) + } + + claims, err := auth.NewIssuer(secret, time.Hour, nil).Verify(token) + if err != nil { + t.Fatalf("issued token does not verify: %v", err) + } + if claims.Subject != u.ID.String() { + t.Errorf("token sub = %q, want owner %q", claims.Subject, u.ID) + } + if len(keys.touched) != 1 || keys.touched[0] != key.ID { + t.Errorf("exchange did not record last-used, touched=%v", keys.touched) + } +} + +func TestExchangeUnknownKeyIsInvalid(t *testing.T) { + _, exchange, _, _, _, _ := newKeyFixtures(t) + if _, _, err := exchange.Execute(context.Background(), "scimesh_wk_live_nope"); !errors.Is(err, usecase.ErrInvalidWorkerKey) { + t.Errorf("got %v, want ErrInvalidWorkerKey", err) + } +} + +func TestExchangeEmptyKeyIsInvalid(t *testing.T) { + _, exchange, _, _, _, _ := newKeyFixtures(t) + if _, _, err := exchange.Execute(context.Background(), ""); !errors.Is(err, usecase.ErrInvalidWorkerKey) { + t.Errorf("got %v, want ErrInvalidWorkerKey", err) + } +} + +func TestExchangeRevokedKeyIsInvalid(t *testing.T) { + create, exchange, revoke, _, _, u := newKeyFixtures(t) + ctx := context.Background() + + key, raw, err := create.Execute(ctx, u.ID, "laptop") + if err != nil { + t.Fatal(err) + } + if err := revoke.Execute(ctx, u.ID, key.ID); err != nil { + t.Fatalf("revoke: %v", err) + } + if _, _, err := exchange.Execute(ctx, raw); !errors.Is(err, usecase.ErrInvalidWorkerKey) { + t.Errorf("revoked key still exchanges: %v", err) + } +} + +func TestRevokeIsScopedToOwner(t *testing.T) { + create, _, revoke, _, _, u := newKeyFixtures(t) + ctx := context.Background() + + key, _, err := create.Execute(ctx, u.ID, "laptop") + if err != nil { + t.Fatal(err) + } + // A different user must not be able to revoke this key. + if err := revoke.Execute(ctx, uuid.New(), key.ID); !errors.Is(err, usecase.ErrWorkerKeyNotFound) { + t.Errorf("cross-owner revoke returned %v, want ErrWorkerKeyNotFound", err) + } +} + +func TestListReturnsOnlyLiveKeys(t *testing.T) { + create, _, revoke, list, _, u := newKeyFixtures(t) + ctx := context.Background() + + live, _, err := create.Execute(ctx, u.ID, "keep") + if err != nil { + t.Fatal(err) + } + dead, _, err := create.Execute(ctx, u.ID, "drop") + if err != nil { + t.Fatal(err) + } + if err := revoke.Execute(ctx, u.ID, dead.ID); err != nil { + t.Fatal(err) + } + + got, err := list.Execute(ctx, u.ID) + if err != nil { + t.Fatalf("list: %v", err) + } + if len(got) != 1 || got[0].ID != live.ID { + t.Errorf("list = %d keys, want only the live one", len(got)) + } +} diff --git a/users/migrations/0003_worker_keys.down.sql b/users/migrations/0003_worker_keys.down.sql new file mode 100644 index 0000000..af1e4b0 --- /dev/null +++ b/users/migrations/0003_worker_keys.down.sql @@ -0,0 +1,5 @@ +BEGIN; + +DROP TABLE IF EXISTS worker_keys; + +COMMIT; diff --git a/users/migrations/0003_worker_keys.up.sql b/users/migrations/0003_worker_keys.up.sql new file mode 100644 index 0000000..5a6b3f7 --- /dev/null +++ b/users/migrations/0003_worker_keys.up.sql @@ -0,0 +1,31 @@ +BEGIN; + +-- A worker key is a long-lived credential a user creates to run a worker on +-- their own machine. Unlike the 24h login JWT, it does not expire on its own: +-- the worker presents it to /worker-tokens/exchange to mint a short-lived JWT +-- and refreshes as needed. Only a SHA-256 hash is stored, never the key itself, +-- so a database leak cannot be replayed as a credential. +CREATE TABLE worker_keys ( + id uuid PRIMARY KEY, + user_id uuid NOT NULL REFERENCES users(id) ON DELETE CASCADE, + -- Human label so a user can tell their machines apart when revoking. + name text NOT NULL, + -- Hex SHA-256 of the presented key. The key is high-entropy, so a fast hash + -- is enough — no per-key salt or bcrypt cost is needed here. + token_hash text NOT NULL, + -- The leading, non-secret slice of the key, shown in the UI to identify a + -- row without ever revealing the secret again. + prefix text NOT NULL, + created_at timestamptz NOT NULL DEFAULT now(), + -- Last successful exchange; NULL until first use. + last_used_at timestamptz, + -- Set when the user revokes the key; a revoked key never exchanges again. + revoked_at timestamptz, + + CONSTRAINT uq_worker_keys_token_hash UNIQUE (token_hash) +); + +-- Listing and revoking are always scoped to one owner's live keys. +CREATE INDEX ix_worker_keys_user_active ON worker_keys (user_id) WHERE revoked_at IS NULL; + +COMMIT;