forked from molecule-ai/molecule-core
orgtoken.Validate now returns org_id (the org workspace UUID stored on org_api_tokens rows, populated by #1212). Both call sites in wsauth_middleware.go — WorkspaceAuth and AdminAuth — call c.Set("org_id", orgID) after successful org-token validation. This unbreaks orgCallerID(c) for org-token callers. Previously the middleware populated org_token_id and org_token_prefix but never org_id, so any handler reading c.Get("org_id") (e.g. requireCallerOwnsOrg) got "" even for valid org tokens. The change is additive: orgID may be empty for pre-migration tokens minted before #1212. requireCallerOwnsOrg already handles empty org_id by denying by default. Co-authored-by: Molecule AI CP-BE <cp-be@agents.moleculesai.app> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
molecule-ai[bot]
GitHub
parent
04c3bc6eb1
commit
bc9ce59b79
@ -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) {
|
||||
|
||||
@ -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)
|
||||
}
|
||||
}
|
||||
@ -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
|
||||
|
||||
Loading…
Reference in New Issue
Block a user