diff --git a/workspace-server/internal/handlers/org_tokens.go b/workspace-server/internal/handlers/org_tokens.go index fb79271d..3a9b36d6 100644 --- a/workspace-server/internal/handlers/org_tokens.go +++ b/workspace-server/internal/handlers/org_tokens.go @@ -112,27 +112,37 @@ func (h *OrgTokenHandler) Revoke(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"revoked": id}) } +// Provenance labels used in the org_api_tokens.created_by column +// and in mint/revoke audit logs. Kept as constants so the labels +// are greppable across services (log pipelines, audit queries). +const ( + actorOrgTokenPrefix = "org-token:" // appended: 8-char plaintext prefix from the UI + actorSession = "session" // WorkOS-session-verified call + actorAdminToken = "admin-token" // bootstrap ADMIN_TOKEN env +) + // orgTokenActor derives a short provenance string for audit. // // - If the request was authed via another org token, return -// "org-token:" so revoke audits show which token minted -// the new one. +// "org-token:" where prefix is the 8-char plaintext +// prefix shown in the UI — correlates audit rows directly with +// the revoke button a user sees. // - If authed via session cookie (AdminAuth's session tier), the -// middleware doesn't set anything on c for us — return "session" -// as a generic label. When we grow a session-user-id capture -// upgrade this to return the real WorkOS user_id. +// middleware doesn't stash a WorkOS user_id today — return +// "session" as a generic label. Follow-up (see +// docs/architecture/org-api-keys-followups.md #6) captures the +// user_id through the session tier for full attribution. // - Else (ADMIN_TOKEN / bootstrap), return "admin-token". func orgTokenActor(c *gin.Context) string { - if v, ok := c.Get("org_token_id"); ok { - if s, ok := v.(string); ok && len(s) >= 8 { - return "org-token:" + s[:8] + if v, ok := c.Get("org_token_prefix"); ok { + if s, ok := v.(string); ok && s != "" { + return actorOrgTokenPrefix + s } } // Session-tier auth doesn't stash an identity in the gin context - // today. Until it does, treat session requests as "session". A - // follow-up issue captures WorkOS user_id here. + // today. Until it does, treat session requests as "session". if c.GetHeader("Cookie") != "" { - return "session" + return actorSession } - return "admin-token" + return actorAdminToken } diff --git a/workspace-server/internal/handlers/org_tokens_test.go b/workspace-server/internal/handlers/org_tokens_test.go new file mode 100644 index 00000000..633df653 --- /dev/null +++ b/workspace-server/internal/handlers/org_tokens_test.go @@ -0,0 +1,287 @@ +package handlers + +import ( + "bytes" + "context" + "database/sql" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/Molecule-AI/molecule-monorepo/platform/internal/db" + "github.com/gin-gonic/gin" +) + +// setupOrgTokenTest wires the package-global db.DB to a sqlmock for +// the duration of a test, returning the handler + mock + cleanup. +// Gin runs in release mode to suppress debug noise. +func setupOrgTokenTest(t *testing.T) (*OrgTokenHandler, sqlmock.Sqlmock, func()) { + t.Helper() + gin.SetMode(gin.ReleaseMode) + mock, mockDB, cleanup := mockSQLDB(t) + prev := db.DB + db.DB = mockDB + return NewOrgTokenHandler(), mock, func() { + db.DB = prev + cleanup() + } +} + +// mockSQLDB returns an sqlmock + *sql.DB pair. Caller restores +// package state via the cleanup func. +func mockSQLDB(t *testing.T) (sqlmock.Sqlmock, *sql.DB, func()) { + t.Helper() + d, m, err := sqlmock.New() + if err != nil { + t.Fatalf("sqlmock: %v", err) + } + return m, d, func() { _ = d.Close() } +} + +// buildCtx returns a gin.Context + recorder wired for the given +// method+path+body. Test code can set headers / context values +// before calling the handler. +func buildCtx(method, path, body string) (*gin.Context, *httptest.ResponseRecorder) { + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + var r *http.Request + if body != "" { + r = httptest.NewRequest(method, path, bytes.NewBufferString(body)) + r.Header.Set("Content-Type", "application/json") + } else { + r = httptest.NewRequest(method, path, nil) + } + c.Request = r.WithContext(context.Background()) + return c, w +} + +// ---- List ----------------------------------------------------------------- + +func TestOrgTokenHandler_List_HappyPath(t *testing.T) { + h, mock, cleanup := setupOrgTokenTest(t) + defer cleanup() + + now := time.Now().UTC() + mock.ExpectQuery(`SELECT id, prefix.*FROM org_api_tokens`). + WillReturnRows(sqlmock.NewRows([]string{"id", "prefix", "name", "created_by", "created_at", "last_used_at"}). + AddRow("tok-1", "abcd1234", "zapier", "session", now, nil)) + + c, w := buildCtx("GET", "/org/tokens", "") + h.List(c) + + if w.Code != http.StatusOK { + t.Fatalf("want 200, got %d: %s", w.Code, w.Body.String()) + } + var body struct { + Count int `json:"count"` + Tokens []struct { + ID string `json:"id"` + Prefix string `json:"prefix"` + Name string `json:"name"` + } `json:"tokens"` + } + if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil { + t.Fatalf("parse: %v", err) + } + if body.Count != 1 || len(body.Tokens) != 1 { + t.Fatalf("expected 1 token, got %+v", body) + } + if body.Tokens[0].Prefix != "abcd1234" { + t.Errorf("prefix not propagated: %q", body.Tokens[0].Prefix) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet: %v", err) + } +} + +func TestOrgTokenHandler_List_DBError500(t *testing.T) { + h, mock, cleanup := setupOrgTokenTest(t) + defer cleanup() + mock.ExpectQuery(`SELECT id, prefix.*FROM org_api_tokens`). + WillReturnError(sql.ErrConnDone) + + c, w := buildCtx("GET", "/org/tokens", "") + h.List(c) + + if w.Code != http.StatusInternalServerError { + t.Errorf("db error → 500 expected; got %d", w.Code) + } +} + +// ---- Create --------------------------------------------------------------- + +func TestOrgTokenHandler_Create_ActorFromAdminToken(t *testing.T) { + h, mock, cleanup := setupOrgTokenTest(t) + defer cleanup() + + // No Cookie header, no org_token_prefix → actor should be + // "admin-token" (bootstrap path). + mock.ExpectQuery(`INSERT INTO org_api_tokens`). + WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg(), "my-ci", actorAdminToken). + WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("tok-1")) + + c, w := buildCtx("POST", "/org/tokens", `{"name":"my-ci"}`) + h.Create(c) + + if w.Code != http.StatusOK { + t.Fatalf("want 200, got %d: %s", w.Code, w.Body.String()) + } + var body struct { + ID string `json:"id"` + Prefix string `json:"prefix"` + Name string `json:"name"` + Token string `json:"auth_token"` + Warning string `json:"warning"` + } + if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil { + t.Fatalf("parse: %v", err) + } + if body.Token == "" { + t.Errorf("plaintext auth_token missing from response") + } + if body.Prefix != body.Token[:8] { + t.Errorf("prefix %q should equal first 8 chars of token %q", body.Prefix, body.Token[:8]) + } + if body.Name != "my-ci" { + t.Errorf("name round-trip mismatch: %q", body.Name) + } + if body.Warning == "" { + t.Errorf("warning text missing") + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet: %v", err) + } +} + +func TestOrgTokenHandler_Create_ActorFromOrgTokenPrefix(t *testing.T) { + h, mock, cleanup := setupOrgTokenTest(t) + defer cleanup() + + // When an existing org token authenticated the mint, audit + // records "org-token:" using the 8-char plaintext + // prefix from the presenter's token. + mock.ExpectQuery(`INSERT INTO org_api_tokens`). + WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg(), nil, actorOrgTokenPrefix+"parent12"). + WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("tok-new")) + + c, w := buildCtx("POST", "/org/tokens", `{}`) + c.Set("org_token_prefix", "parent12") + h.Create(c) + + if w.Code != http.StatusOK { + t.Fatalf("want 200, got %d: %s", w.Code, w.Body.String()) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet: %v", err) + } +} + +func TestOrgTokenHandler_Create_ActorFromSession(t *testing.T) { + h, mock, cleanup := setupOrgTokenTest(t) + defer cleanup() + + // Cookie present but no org_token_prefix — that's the canvas + // session path. Today recorded as bare "session". When the + // follow-up captures WorkOS user_id, this test changes to + // "session:user_01XXX". + mock.ExpectQuery(`INSERT INTO org_api_tokens`). + WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg(), "from-browser", actorSession). + WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("tok-browser")) + + c, w := buildCtx("POST", "/org/tokens", `{"name":"from-browser"}`) + c.Request.Header.Set("Cookie", "mcp_session=abc") + h.Create(c) + + if w.Code != http.StatusOK { + t.Fatalf("want 200, got %d: %s", w.Code, w.Body.String()) + } +} + +func TestOrgTokenHandler_Create_NameTooLong400(t *testing.T) { + h, _, cleanup := setupOrgTokenTest(t) + defer cleanup() + longName := string(make([]byte, 101)) + for i := range longName { + _ = i + } + // Build a 101-char name (any ASCII works; we're hitting the + // length bound). + buf := make([]byte, 101) + for i := range buf { + buf[i] = 'a' + } + c, w := buildCtx("POST", "/org/tokens", `{"name":"`+string(buf)+`"}`) + h.Create(c) + + if w.Code != http.StatusBadRequest { + t.Errorf("oversize name: want 400, got %d", w.Code) + } +} + +func TestOrgTokenHandler_Create_EmptyBodyOK(t *testing.T) { + h, mock, cleanup := setupOrgTokenTest(t) + defer cleanup() + // Empty POST must still mint a token — name is optional. + mock.ExpectQuery(`INSERT INTO org_api_tokens`). + WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg(), nil, actorAdminToken). + WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("tok-min")) + + c, w := buildCtx("POST", "/org/tokens", "") + h.Create(c) + + if w.Code != http.StatusOK { + t.Errorf("empty body should mint OK; got %d: %s", w.Code, w.Body.String()) + } +} + +// ---- Revoke --------------------------------------------------------------- + +func TestOrgTokenHandler_Revoke_HappyPath200(t *testing.T) { + h, mock, cleanup := setupOrgTokenTest(t) + defer cleanup() + + mock.ExpectExec(`UPDATE org_api_tokens`). + WithArgs("tok-1"). + WillReturnResult(sqlmock.NewResult(0, 1)) + + c, w := buildCtx("DELETE", "/org/tokens/tok-1", "") + c.Params = gin.Params{{Key: "id", Value: "tok-1"}} + h.Revoke(c) + + if w.Code != http.StatusOK { + t.Errorf("want 200, got %d: %s", w.Code, w.Body.String()) + } +} + +func TestOrgTokenHandler_Revoke_Missing404(t *testing.T) { + // Idempotency: revoking a non-existent or already-revoked id + // returns 404 — callers can tell "worked" from "already done". + h, mock, cleanup := setupOrgTokenTest(t) + defer cleanup() + mock.ExpectExec(`UPDATE org_api_tokens`). + WithArgs("ghost"). + WillReturnResult(sqlmock.NewResult(0, 0)) + + c, w := buildCtx("DELETE", "/org/tokens/ghost", "") + c.Params = gin.Params{{Key: "id", Value: "ghost"}} + h.Revoke(c) + + if w.Code != http.StatusNotFound { + t.Errorf("want 404, got %d", w.Code) + } +} + +func TestOrgTokenHandler_Revoke_MissingID400(t *testing.T) { + h, _, cleanup := setupOrgTokenTest(t) + defer cleanup() + c, w := buildCtx("DELETE", "/org/tokens/", "") + // No id param set. + h.Revoke(c) + + if w.Code != http.StatusBadRequest { + t.Errorf("want 400, got %d", w.Code) + } +} diff --git a/workspace-server/internal/middleware/wsauth_middleware.go b/workspace-server/internal/middleware/wsauth_middleware.go index 2fd8ae16..e516a5fc 100644 --- a/workspace-server/internal/middleware/wsauth_middleware.go +++ b/workspace-server/internal/middleware/wsauth_middleware.go @@ -60,8 +60,9 @@ func WorkspaceAuth(database *sql.DB) gin.HandlerFunc { // power surface as ADMIN_TOKEN but named, revocable, audited. // Check before per-workspace token so an org-key presenter // doesn't hit the narrower ValidateToken failure path. - if id, err := orgtoken.Validate(ctx, database, tok); err == nil { + if id, prefix, err := orgtoken.Validate(ctx, database, tok); err == nil { c.Set("org_token_id", id) + c.Set("org_token_prefix", prefix) c.Next() return } else if !errors.Is(err, orgtoken.ErrInvalidToken) { @@ -180,8 +181,9 @@ func AdminAuth(database *sql.DB) gin.HandlerFunc { // index with revoked_at IS NULL) + an async last_used_at // bump. Cost per request: one SELECT + one UPDATE, both // hitting the same narrow partial index. - if id, err := orgtoken.Validate(ctx, database, tok); err == nil { + if id, prefix, err := orgtoken.Validate(ctx, database, tok); err == nil { c.Set("org_token_id", id) + c.Set("org_token_prefix", prefix) c.Next() return } else if !errors.Is(err, orgtoken.ErrInvalidToken) { diff --git a/workspace-server/internal/orgtoken/tokens.go b/workspace-server/internal/orgtoken/tokens.go index e340247c..45be465b 100644 --- a/workspace-server/internal/orgtoken/tokens.go +++ b/workspace-server/internal/orgtoken/tokens.go @@ -36,6 +36,12 @@ const ( // crackable on its own (6 bits × 8 = 48 bits of prefix space — // good enough to disambiguate, nowhere near guessable). tokenPrefixLen = 8 + + // listMax caps the number of rows List returns. Realistic org- + // admin UIs show on the order of 10. 500 is enough headroom that + // no legitimate flow hits it, low enough that a mint-storm can't + // force a big allocation on every list render. + listMax = 500 ) // ErrInvalidToken is returned when a presented bearer doesn't match @@ -86,33 +92,42 @@ func Issue(ctx context.Context, db *sql.DB, name, createdBy string) (plaintext, // Validate looks up a presented bearer, returns ErrInvalidToken on // any mismatch (bad bytes, revoked, deleted). On success, updates // last_used_at best-effort (the hot path — failure to update doesn't -// fail the request) and returns the token id for audit logging. -func Validate(ctx context.Context, db *sql.DB, plaintext string) (string, error) { +// fail the request) and returns the token id + display prefix for +// audit logging. +// +// Returning the prefix alongside the id lets callers produce audit +// strings that match what users see in the UI (the plaintext prefix, +// not the UUID). Keeps the "who did what" trail visually +// correlatable to the revoke button in the token list. +func Validate(ctx context.Context, db *sql.DB, plaintext string) (id, prefix string, err error) { if plaintext == "" { - return "", ErrInvalidToken + return "", "", ErrInvalidToken } hash := sha256.Sum256([]byte(plaintext)) - var id string - err := db.QueryRowContext(ctx, ` - SELECT id FROM org_api_tokens + queryErr := db.QueryRowContext(ctx, ` + SELECT id, prefix FROM org_api_tokens WHERE token_hash = $1 AND revoked_at IS NULL - `, hash[:]).Scan(&id) - if err != nil { + `, hash[:]).Scan(&id, &prefix) + if queryErr != nil { // Collapse all failure shapes into ErrInvalidToken so the // caller can't accidentally leak "row exists but revoked" vs // "row never existed" via response shape. - return "", ErrInvalidToken + return "", "", ErrInvalidToken } // Best-effort last_used_at bump. Failure here is acceptable — the // request is already authenticated; we don't want a transient DB // blip to flip a 200 into a 500. _, _ = db.ExecContext(ctx, `UPDATE org_api_tokens SET last_used_at = now() WHERE id = $1`, id) - return id, nil + return id, prefix, nil } // List returns live (non-revoked) tokens newest-first. Safe to // expose to the admin UI — no hash, no plaintext, only prefix. +// +// Capped at listMax rows. A UI page with more than that is a +// symptom of abuse or a bug — the hard cap prevents one runaway +// minting loop from O(N) pageloads in the admin UI. func List(ctx context.Context, db *sql.DB) ([]Token, error) { rows, err := db.QueryContext(ctx, ` SELECT id, prefix, COALESCE(name,''), COALESCE(created_by,''), @@ -120,7 +135,8 @@ func List(ctx context.Context, db *sql.DB) ([]Token, error) { FROM org_api_tokens WHERE revoked_at IS NULL ORDER BY created_at DESC - `) + LIMIT $1 + `, listMax) if err != nil { return nil, fmt.Errorf("orgtoken: list: %w", err) } diff --git a/workspace-server/internal/orgtoken/tokens_test.go b/workspace-server/internal/orgtoken/tokens_test.go index e5b4a918..dea3f729 100644 --- a/workspace-server/internal/orgtoken/tokens_test.go +++ b/workspace-server/internal/orgtoken/tokens_test.go @@ -72,26 +72,29 @@ func TestValidate_HappyPath(t *testing.T) { plaintext := "known-plaintext-for-test" hash := sha256.Sum256([]byte(plaintext)) - mock.ExpectQuery(`SELECT id FROM org_api_tokens`). + mock.ExpectQuery(`SELECT id, prefix FROM org_api_tokens`). WithArgs(hash[:]). - WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow("tok-live")) + WillReturnRows(sqlmock.NewRows([]string{"id", "prefix"}).AddRow("tok-live", "abcd1234")) mock.ExpectExec(`UPDATE org_api_tokens SET last_used_at`). WithArgs("tok-live"). WillReturnResult(sqlmock.NewResult(0, 1)) - id, err := Validate(context.Background(), db, plaintext) + id, prefix, err := Validate(context.Background(), db, plaintext) if err != nil { t.Fatalf("Validate: %v", err) } if id != "tok-live" { t.Errorf("id = %q, want tok-live", id) } + if prefix != "abcd1234" { + t.Errorf("prefix = %q, want abcd1234", prefix) + } } func TestValidate_EmptyPlaintextRejected(t *testing.T) { db, _, _ := sqlmock.New() defer db.Close() - if _, err := Validate(context.Background(), db, ""); !errors.Is(err, ErrInvalidToken) { + if _, _, err := Validate(context.Background(), db, ""); !errors.Is(err, ErrInvalidToken) { t.Errorf("empty plaintext should be ErrInvalidToken, got %v", err) } } @@ -103,11 +106,11 @@ func TestValidate_UnknownHashErrInvalid(t *testing.T) { } defer db.Close() - mock.ExpectQuery(`SELECT id FROM org_api_tokens`). + mock.ExpectQuery(`SELECT id, prefix FROM org_api_tokens`). WithArgs(sqlmock.AnyArg()). WillReturnError(sql.ErrNoRows) - if _, err := Validate(context.Background(), db, "ghost"); !errors.Is(err, ErrInvalidToken) { + if _, _, err := Validate(context.Background(), db, "ghost"); !errors.Is(err, ErrInvalidToken) { t.Errorf("unknown hash should be ErrInvalidToken, got %v", err) } } @@ -120,11 +123,11 @@ func TestValidate_RevokedTokenNotAccepted(t *testing.T) { defer db.Close() // Query has `AND revoked_at IS NULL` — sqlmock will return // ErrNoRows because the revoked row is filtered out. - mock.ExpectQuery(`SELECT id FROM org_api_tokens`). + mock.ExpectQuery(`SELECT id, prefix FROM org_api_tokens`). WithArgs(sqlmock.AnyArg()). WillReturnError(sql.ErrNoRows) - if _, err := Validate(context.Background(), db, "revoked-plaintext"); !errors.Is(err, ErrInvalidToken) { + if _, _, err := Validate(context.Background(), db, "revoked-plaintext"); !errors.Is(err, ErrInvalidToken) { t.Errorf("revoked token should be ErrInvalidToken, got %v", err) } } @@ -139,6 +142,7 @@ func TestList_NewestFirst(t *testing.T) { now := time.Now() earlier := now.Add(-1 * time.Hour) mock.ExpectQuery(`SELECT id, prefix.*FROM org_api_tokens.*ORDER BY created_at DESC`). + WithArgs(listMax). WillReturnRows(sqlmock.NewRows([]string{"id", "prefix", "name", "created_by", "created_at", "last_used_at"}). AddRow("t2", "abcd1234", "zapier", "user_01", now, now). AddRow("t1", "efgh5678", "", "", earlier, nil)) diff --git a/workspace-server/internal/router/router.go b/workspace-server/internal/router/router.go index b547a2e3..ee0767ad 100644 --- a/workspace-server/internal/router/router.go +++ b/workspace-server/internal/router/router.go @@ -476,11 +476,20 @@ func Setup(hub *ws.Hub, broadcaster *events.Broadcaster, prov *provisioner.Provi // session cookie, OR an existing org token to mint more. That's // bootstrap-friendly (first token from ADMIN_TOKEN or canvas // session) and self-sustaining afterwards (tokens mint tokens). + // + // The mint endpoint gets an extra per-IP rate limiter — a + // compromised session or leaked bearer could otherwise mint + // thousands of tokens per second, making forensic cleanup + // painful. 10 mints per hour per IP is ample for real usage; + // legitimate bursts fit in the ceiling and abuse bounces off. + // List + Delete don't need the extra limit (they can't be used + // to generate new secret material). { orgTokenHandler := handlers.NewOrgTokenHandler() orgTokenAdmin := r.Group("", middleware.AdminAuth(db.DB)) orgTokenAdmin.GET("/org/tokens", orgTokenHandler.List) - orgTokenAdmin.POST("/org/tokens", orgTokenHandler.Create) + orgTokenMintLimiter := middleware.NewRateLimiter(10, time.Hour, context.Background()) + orgTokenAdmin.POST("/org/tokens", orgTokenMintLimiter.Middleware(), orgTokenHandler.Create) orgTokenAdmin.DELETE("/org/tokens/:id", orgTokenHandler.Revoke) }