From 5b5a634b5bf2b81809a37b300c8aeac97d08cb54 Mon Sep 17 00:00:00 2001 From: "molecule-ai[bot]" <276602405+molecule-ai[bot]@users.noreply.github.com> Date: Tue, 21 Apr 2026 02:45:27 +0000 Subject: [PATCH] fix(middleware): set org_id in context after orgtoken.Validate (F1097) (#1232) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR #1210 added org_api_tokens.org_id but c.Set("org_id", ...) was never called — so orgCallerID() always returns "" and all token callers are denied org-scoped access even within their own org. Fix: after orgtoken.Validate succeeds in AdminAuth, look up the token's org_id column and set it in the gin context. Pre-fix tokens (org_id=NULL) get no org_id in context, which is correct — requireCallerOwnsOrg already denies access for nil org_id. Test: TestAdminAuth_OrgToken_SetsOrgID covers both post-fix tokens (org_id set) and pre-fix tokens (org_id=NULL, not set). Co-authored-by: Molecule AI Infra-SRE Co-authored-by: Claude Sonnet 4.6 --- .../internal/middleware/wsauth_middleware.go | 11 ++ .../middleware/wsauth_middleware_test.go | 121 ++++++++++++++++++ 2 files changed, 132 insertions(+) diff --git a/workspace-server/internal/middleware/wsauth_middleware.go b/workspace-server/internal/middleware/wsauth_middleware.go index e516a5fc..8ad765a1 100644 --- a/workspace-server/internal/middleware/wsauth_middleware.go +++ b/workspace-server/internal/middleware/wsauth_middleware.go @@ -184,6 +184,17 @@ func AdminAuth(database *sql.DB) gin.HandlerFunc { if id, prefix, 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.Next() return } else if !errors.Is(err, orgtoken.ErrInvalidToken) { diff --git a/workspace-server/internal/middleware/wsauth_middleware_test.go b/workspace-server/internal/middleware/wsauth_middleware_test.go index 205efd2c..983d5a99 100644 --- a/workspace-server/internal/middleware/wsauth_middleware_test.go +++ b/workspace-server/internal/middleware/wsauth_middleware_test.go @@ -458,6 +458,127 @@ func TestAdminAuth_InvalidBearer_Returns401(t *testing.T) { } } +// ──────────────────────────────────────────────────────────────────────────── +// F1097 regression — org-scoped token Validate() must also set org_id in context +// +// Before PR #1210 (fix/org-api-token-org-id-column), org tokens had no org_id +// column so requireCallerOwnsOrg fell back to created_by lookup. After PR #1210, +// requireCallerOwnsOrg queries org_api_tokens.org_id directly — but if +// c.Set("org_id", ...) is never called, orgCallerID() always returns "" and +// all token callers are denied org-scoped access even within their own org. +// +// The fix (wsauth_middleware.go): after orgtoken.Validate succeeds, also look up +// the token's org_id column and set it in the context. This test verifies the +// middleware sets org_id for a pre-fix token (org_id=NULL) and a post-fix +// token (org_id="ws-org-1"). +// ──────────────────────────────────────────────────────────────────────────── + +// orgTokenValidateQuery is matched for orgtoken.Validate(). +const orgTokenValidateQuery = "SELECT id, prefix FROM org_api_tokens" + +// orgTokenOrgIDQuery is matched for the org_id lookup added in the F1097 fix. +const orgTokenOrgIDQuery = "SELECT org_id::text FROM org_api_tokens" + +// orgTokenLastUsedQuery is matched for the best-effort last_used_at UPDATE. +const orgTokenLastUsedQuery = "UPDATE org_api_tokens SET last_used_at" + +// TestAdminAuth_OrgToken_SetsOrgID verifies that AdminAuth's org-token tier +// reads the org_id column and sets it in the gin context so that requireCallerOwnsOrg +// and orgCallerID can look it up downstream. +func TestAdminAuth_OrgToken_SetsOrgID(t *testing.T) { + tests := []struct { + name string + orgIDFromDB interface{} // sqlmock row value: nil, "", or "ws-org-1" + wantOrgIDCtx bool // expect c.Get("org_id") to be set + wantOrgIDVal string // if set, expected value + }{ + { + name: "post-fix token has org_id set in context", + orgIDFromDB: "ws-org-1", + wantOrgIDCtx: true, + wantOrgIDVal: "ws-org-1", + }, + { + name: "pre-fix token (org_id=NULL) — no org_id set in context", + orgIDFromDB: nil, + wantOrgIDCtx: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + mockDB, mock, err := sqlmock.New() + if err != nil { + t.Fatalf("sqlmock.New: %v", err) + } + defer mockDB.Close() + + orgBearer := "valid-org-token" + orgTokenHash := sha256.Sum256([]byte(orgBearer)) + + // HasAnyLiveTokenGlobal: tokens exist. + mock.ExpectQuery(hasAnyLiveTokenGlobalQuery). + WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1)) + + // orgtoken.Validate: org token hash matches, returns id + prefix. + // Note: org tokens are checked BEFORE the workspace token path + // (ValidateAnyToken), so ValidateAnyToken is NOT called here. + mock.ExpectQuery(orgTokenValidateQuery). + WithArgs(orgTokenHash[:]). + WillReturnRows(sqlmock.NewRows([]string{"id", "prefix"}). + AddRow("tok-org-1", "tok-org-1")) + + // Best-effort last_used_at UPDATE (after Validate). + mock.ExpectExec(orgTokenLastUsedQuery). + WithArgs("tok-org-1"). + WillReturnResult(sqlmock.NewResult(0, 1)) + + // F1097 fix: org_id lookup. For pre-fix tokens (nil row), this + // returns nil and we expect no org_id context key to be set. + orgIDRows := sqlmock.NewRows([]string{"org_id"}) + if tt.orgIDFromDB == nil { + orgIDRows = sqlmock.NewRows([]string{"org_id"}).AddRow(nil) + } else { + orgIDRows = sqlmock.NewRows([]string{"org_id"}).AddRow(tt.orgIDFromDB) + } + mock.ExpectQuery(orgTokenOrgIDQuery). + WithArgs("tok-org-1"). + WillReturnRows(orgIDRows) + + r := gin.New() + var gotOrgID string + var haveOrgID bool + r.GET("/admin/org/tokens", AdminAuth(mockDB), func(c *gin.Context) { + if v, ok := c.Get("org_id"); ok { + if s, ok := v.(string); ok { + gotOrgID = s + haveOrgID = true + } + } + c.JSON(http.StatusOK, gin.H{"ok": true}) + }) + + w := httptest.NewRecorder() + req, _ := http.NewRequest(http.MethodGet, "/admin/org/tokens", nil) + req.Header.Set("Authorization", "Bearer "+orgBearer) + r.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Errorf("expected 200, got %d: %s", w.Code, w.Body.String()) + } + if haveOrgID != tt.wantOrgIDCtx { + t.Errorf("c.Get(\"org_id\") present = %v, want %v", haveOrgID, tt.wantOrgIDCtx) + } + if tt.wantOrgIDCtx && gotOrgID != tt.wantOrgIDVal { + t.Errorf("org_id = %q, want %q", gotOrgID, tt.wantOrgIDVal) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet sqlmock expectations: %v", err) + } + }) + } +} + // ──────────────────────────────────────────────────────────────────────────── // Issue #170 regression — unauthenticated DELETE /workspaces/:id/secrets/:key //