diff --git a/workspace-server/internal/middleware/wsauth_middleware.go b/workspace-server/internal/middleware/wsauth_middleware.go index 8ad765a1..bb72ab50 100644 --- a/workspace-server/internal/middleware/wsauth_middleware.go +++ b/workspace-server/internal/middleware/wsauth_middleware.go @@ -60,9 +60,10 @@ 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, prefix, err := orgtoken.Validate(ctx, database, tok); err == nil { + if id, prefix, orgID, err := orgtoken.Validate(ctx, database, tok); err == nil { c.Set("org_token_id", id) c.Set("org_token_prefix", prefix) + c.Set("org_id", orgID) c.Next() return } else if !errors.Is(err, orgtoken.ErrInvalidToken) { @@ -181,20 +182,10 @@ 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, prefix, err := orgtoken.Validate(ctx, database, tok); err == nil { + if id, prefix, orgID, err := orgtoken.Validate(ctx, database, tok); err == nil { c.Set("org_token_id", id) c.Set("org_token_prefix", prefix) - // F1097: also set org_id from the token's org_id column so that - // requireCallerOwnsOrg can look it up via c.Get("org_id"). - // Tokens created before PR #1210 have org_id=NULL — for those, - // the SELECT returns nil and no org_id is set, which is correct: - // requireCallerOwnsOrg already denies access for nil org_id. - var orgID *string - if err := database.QueryRowContext(ctx, - `SELECT org_id::text FROM org_api_tokens WHERE id = $1`, id, - ).Scan(&orgID); err == nil && orgID != nil && *orgID != "" { - c.Set("org_id", *orgID) - } + c.Set("org_id", orgID) c.Next() return } else if !errors.Is(err, orgtoken.ErrInvalidToken) { diff --git a/workspace-server/internal/middleware/wsauth_middleware_org_id_test.go b/workspace-server/internal/middleware/wsauth_middleware_org_id_test.go new file mode 100644 index 00000000..e89e4f77 --- /dev/null +++ b/workspace-server/internal/middleware/wsauth_middleware_org_id_test.go @@ -0,0 +1,326 @@ +package middleware + +import ( + "crypto/sha256" + "database/sql" + "net/http" + "net/http/httptest" + "testing" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/gin-gonic/gin" +) + +// orgTokenValidateQuery is matched for orgtoken.Validate in both +// WorkspaceAuth and AdminAuth middleware paths. The query selects +// id and prefix from org_api_tokens where token_hash matches and +// revoked_at IS NULL. +const orgTokenValidateQuery = "SELECT id, prefix FROM org_api_tokens WHERE token_hash" + +func TestWorkspaceAuth_ValidOrgToken_SetsOrgIDContext(t *testing.T) { + // F1097 (#1218): org tokens validated via WorkspaceAuth must have + // org_id populated on the Gin context so downstream handlers can + // enforce org isolation without a per-request DB round-trip. + mockDB, mock, err := sqlmock.New() + if err != nil { + t.Fatalf("sqlmock.New: %v", err) + } + defer mockDB.Close() + + orgToken := "tok_test_org_token_abc123" + tokenHash := sha256.Sum256([]byte(orgToken)) + + // orgtoken.Validate — returns id + prefix (no org_id column yet). + mock.ExpectQuery(orgTokenValidateQuery). + WithArgs(tokenHash[:]). + WillReturnRows(sqlmock.NewRows([]string{"id", "prefix"}). + AddRow("tok-org-abc", "tok_test")) + + // F1097: secondary SELECT for org_id from org_api_tokens. + mock.ExpectQuery("SELECT org_id::text FROM org_api_tokens WHERE id"). + WithArgs("tok-org-abc"). + WillReturnRows(sqlmock.NewRows([]string{"org_id"}). + AddRow("00000000-0000-0000-0000-000000000001")) + + r := gin.New() + r.GET("/workspaces/:id/secrets", WorkspaceAuth(mockDB), func(c *gin.Context) { + v, exists := c.Get("org_id") + if !exists { + t.Errorf("org_id not set on context for valid org token") + c.JSON(http.StatusOK, gin.H{"ok": true}) + return + } + c.JSON(http.StatusOK, gin.H{"org_id": v}) + }) + + w := httptest.NewRecorder() + req, _ := http.NewRequest(http.MethodGet, "/workspaces/ws-1/secrets", nil) + req.Header.Set("Authorization", "Bearer "+orgToken) + r.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Errorf("expected 200, got %d: %s", w.Code, w.Body.String()) + } + // org_id must appear in the JSON response body. + body := w.Body.String() + if body == "" || body == "{}" { + t.Errorf("org_id missing from response body: %s", body) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet sqlmock expectations: %v", err) + } +} + +func TestWorkspaceAuth_ValidOrgToken_OrgIDNULL_DoesNotSetContext(t *testing.T) { + // F1097: pre-migration tokens (org_id=NULL) must NOT set org_id on context — + // requireCallerOwnsOrg already handles nil by denying by default, so a + // nil context key is the correct signal. + mockDB, mock, err := sqlmock.New() + if err != nil { + t.Fatalf("sqlmock.New: %v", err) + } + defer mockDB.Close() + + orgToken := "tok_old_token_no_org" + tokenHash := sha256.Sum256([]byte(orgToken)) + + // orgtoken.Validate. + mock.ExpectQuery(orgTokenValidateQuery). + WithArgs(tokenHash[:]). + WillReturnRows(sqlmock.NewRows([]string{"id", "prefix"}). + AddRow("tok-old-xyz", "tok_old_")) + + // F1097: org_id SELECT returns NULL — context key must NOT be set. + mock.ExpectQuery("SELECT org_id::text FROM org_api_tokens WHERE id"). + WithArgs("tok-old-xyz"). + WillReturnRows(sqlmock.NewRows([]string{"org_id"}).AddRow(nil)) + + r := gin.New() + r.GET("/workspaces/:id/secrets", WorkspaceAuth(mockDB), func(c *gin.Context) { + _, exists := c.Get("org_id") + if exists { + t.Errorf("org_id should not be set on context for NULL org_id token") + } + c.JSON(http.StatusOK, gin.H{"ok": true}) + }) + + w := httptest.NewRecorder() + req, _ := http.NewRequest(http.MethodGet, "/workspaces/ws-1/secrets", nil) + req.Header.Set("Authorization", "Bearer "+orgToken) + r.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Errorf("expected 200, got %d: %s", w.Code, w.Body.String()) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet sqlmock expectations: %v", err) + } +} + +func TestAdminAuth_ValidOrgToken_SetsOrgIDContext(t *testing.T) { + // F1097 (#1218): AdminAuth path (used for /org/* routes) must also + // populate org_id so org-token callers can access their own org's + // routes without a separate OrgIDByTokenID call per request. + mockDB, mock, err := sqlmock.New() + if err != nil { + t.Fatalf("sqlmock.New: %v", err) + } + defer mockDB.Close() + + orgToken := "tok_admin_path_org_token" + tokenHash := sha256.Sum256([]byte(orgToken)) + + // HasAnyLiveTokenGlobal: at least one workspace auth token exists globally + // (bootstrap gate — if no tokens exist, AdminAuth grants access to all). + mock.ExpectQuery(hasAnyLiveTokenGlobalQuery). + WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1)) + + // orgtoken.Validate via AdminAuth — returns id + prefix. + mock.ExpectQuery(orgTokenValidateQuery). + WithArgs(tokenHash[:]). + WillReturnRows(sqlmock.NewRows([]string{"id", "prefix"}). + AddRow("tok-admin-org", "tok_adm_")) + + // F1097: secondary SELECT for org_id. + mock.ExpectQuery("SELECT org_id::text FROM org_api_tokens WHERE id"). + WithArgs("tok-admin-org"). + WillReturnRows(sqlmock.NewRows([]string{"org_id"}). + AddRow("00000000-0000-0000-0000-000000000042")) + + r := gin.New() + r.GET("/admin/org-settings", AdminAuth(mockDB), func(c *gin.Context) { + v, exists := c.Get("org_id") + if !exists { + t.Errorf("org_id not set on context for valid org token via AdminAuth") + c.JSON(http.StatusOK, gin.H{"ok": true}) + return + } + c.JSON(http.StatusOK, gin.H{"org_id": v}) + }) + + w := httptest.NewRecorder() + req, _ := http.NewRequest(http.MethodGet, "/admin/org-settings", nil) + req.Header.Set("Authorization", "Bearer "+orgToken) + r.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Errorf("expected 200, got %d: %s", w.Code, w.Body.String()) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet sqlmock expectations: %v", err) + } +} + +func TestAdminAuth_ValidOrgToken_OrgIDNULL_DoesNotSetContext(t *testing.T) { + // F1097: AdminAuth path for pre-migration org token (org_id=NULL) must + // NOT set org_id on context. Tokens minted before F1097 fix have + // org_id=NULL — requireCallerOwnsOrg already denies these by default. + mockDB, mock, err := sqlmock.New() + if err != nil { + t.Fatalf("sqlmock.New: %v", err) + } + defer mockDB.Close() + + orgToken := "tok_old_admin_token" + tokenHash := sha256.Sum256([]byte(orgToken)) + + mock.ExpectQuery(hasAnyLiveTokenGlobalQuery). + WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1)) + + mock.ExpectQuery(orgTokenValidateQuery). + WithArgs(tokenHash[:]). + WillReturnRows(sqlmock.NewRows([]string{"id", "prefix"}). + AddRow("tok-old-admin", "tok_old_")) + + // F1097: org_id is NULL — no context key set. + mock.ExpectQuery("SELECT org_id::text FROM org_api_tokens WHERE id"). + WithArgs("tok-old-admin"). + WillReturnRows(sqlmock.NewRows([]string{"org_id"}).AddRow(nil)) + + r := gin.New() + r.GET("/admin/org-settings", AdminAuth(mockDB), func(c *gin.Context) { + _, exists := c.Get("org_id") + if exists { + t.Errorf("org_id should not be set for NULL org_id token via AdminAuth") + } + c.JSON(http.StatusOK, gin.H{"ok": true}) + }) + + w := httptest.NewRecorder() + req, _ := http.NewRequest(http.MethodGet, "/admin/org-settings", nil) + req.Header.Set("Authorization", "Bearer "+orgToken) + r.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Errorf("expected 200, got %d: %s", w.Code, w.Body.String()) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet sqlmock expectations: %v", err) + } +} + +func TestWorkspaceAuth_OrgToken_DBRowScanError_DoesNotPanic(t *testing.T) { + // F1097: if the org_id SELECT returns an unexpected column count or type, + // the deferred suppress-pattern must not crash — the token is still valid, + // org_id is simply not set (token is denied by requireCallerOwnsOrg at use-time). + mockDB, mock, err := sqlmock.New() + if err != nil { + t.Fatalf("sqlmock.New: %v", err) + } + defer mockDB.Close() + + orgToken := "tok_token_ok" + tokenHash := sha256.Sum256([]byte(orgToken)) + + mock.ExpectQuery(orgTokenValidateQuery). + WithArgs(tokenHash[:]). + WillReturnRows(sqlmock.NewRows([]string{"id", "prefix"}). + AddRow("tok-ok", "tok_tok_")) + + // org_id SELECT fails — sqlmock returns ErrRowNotFound when columns don't match. + // We set up an impossible regex to force a mismatch. + mock.ExpectQuery("SELECT org_id::text FROM org_api_tokens WHERE id"). + WithArgs("tok-ok"). + WillReturnError(sql.ErrNoRows) + + r := gin.New() + r.GET("/workspaces/:id/secrets", WorkspaceAuth(mockDB), func(c *gin.Context) { + // org_id key may or may not be set — either is acceptable here. + // The important thing is we don't panic. + c.JSON(http.StatusOK, gin.H{"ok": true}) + }) + + w := httptest.NewRecorder() + req, _ := http.NewRequest(http.MethodGet, "/workspaces/ws-1/secrets", nil) + req.Header.Set("Authorization", "Bearer "+orgToken) + r.ServeHTTP(w, req) + + // Token is still accepted — only the org_id enrichment fails. + if w.Code != http.StatusOK { + t.Errorf("expected 200 despite org_id SELECT error, got %d: %s", w.Code, w.Body.String()) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet sqlmock expectations: %v", err) + } +} + +// TestWorkspaceAuth_OrgToken_SetsAllContextKeys verifies the complete set of +// context keys set by WorkspaceAuth for a valid org token (F1097 coverage). +func TestWorkspaceAuth_OrgToken_SetsAllContextKeys(t *testing.T) { + mockDB, mock, err := sqlmock.New() + if err != nil { + t.Fatalf("sqlmock.New: %v", err) + } + defer mockDB.Close() + + orgToken := "tok_full_context_token" + tokenHash := sha256.Sum256([]byte(orgToken)) + expectedOrgID := "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee" + + mock.ExpectQuery(orgTokenValidateQuery). + WithArgs(tokenHash[:]). + WillReturnRows(sqlmock.NewRows([]string{"id", "prefix"}). + AddRow("tok-full", "tok_fu_")) + + mock.ExpectQuery("SELECT org_id::text FROM org_api_tokens WHERE id"). + WithArgs("tok-full"). + WillReturnRows(sqlmock.NewRows([]string{"org_id"}).AddRow(expectedOrgID)) + + r := gin.New() + r.GET("/workspaces/:id/secrets", WorkspaceAuth(mockDB), func(c *gin.Context) { + id, ok := c.Get("org_token_id") + if !ok { + t.Errorf("org_token_id not set") + } else if id != "tok-full" { + t.Errorf("org_token_id: got %v, want tok-full", id) + } + + prefix, ok := c.Get("org_token_prefix") + if !ok { + t.Errorf("org_token_prefix not set") + } else if prefix != "tok_fu_" { + t.Errorf("org_token_prefix: got %v, want tok_fu_", prefix) + } + + orgID, ok := c.Get("org_id") + if !ok { + t.Errorf("org_id not set") + } else if orgID != expectedOrgID { + t.Errorf("org_id: got %v, want %s", orgID, expectedOrgID) + } + + c.JSON(http.StatusOK, gin.H{"ok": true}) + }) + + w := httptest.NewRecorder() + req, _ := http.NewRequest(http.MethodGet, "/workspaces/ws-1/secrets", nil) + req.Header.Set("Authorization", "Bearer "+orgToken) + r.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Errorf("expected 200, got %d: %s", w.Code, w.Body.String()) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet sqlmock expectations: %v", err) + } +} \ No newline at end of file diff --git a/workspace-server/internal/orgtoken/tokens.go b/workspace-server/internal/orgtoken/tokens.go index 94e5e551..4276434c 100644 --- a/workspace-server/internal/orgtoken/tokens.go +++ b/workspace-server/internal/orgtoken/tokens.go @@ -94,34 +94,45 @@ func Issue(ctx context.Context, db *sql.DB, name, createdBy, orgID string) (plai // 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 + display prefix for -// audit logging. +// fail the request) and returns the token id + display prefix + org_id +// for audit logging and org isolation. // // 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) { +// +// The org_id is the workspace UUID of the org that owns this token. +// May be empty for pre-migration tokens minted before #1212. Callers +// that need org isolation should use requireCallerOwnsOrg (which does +// a second lookup) rather than trusting an empty org_id here — this +// avoids a breaking change to the Validate interface while still +// populating the Gin context for callers that don't need it. +func Validate(ctx context.Context, db *sql.DB, plaintext string) (id, prefix, orgID string, err error) { if plaintext == "" { - return "", "", ErrInvalidToken + return "", "", "", ErrInvalidToken } hash := sha256.Sum256([]byte(plaintext)) + var orgIDNull sql.NullString queryErr := db.QueryRowContext(ctx, ` - SELECT id, prefix FROM org_api_tokens + SELECT id, prefix, org_id FROM org_api_tokens WHERE token_hash = $1 AND revoked_at IS NULL - `, hash[:]).Scan(&id, &prefix) + `, hash[:]).Scan(&id, &prefix, &orgIDNull) 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 + } + if orgIDNull.Valid { + orgID = orgIDNull.String } // 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, prefix, nil + return id, prefix, orgID, nil } // List returns live (non-revoked) tokens newest-first. Safe to