From 07457ad556f76af0916a5df3532962709c8012f4 Mon Sep 17 00:00:00 2001 From: core-devops Date: Wed, 20 May 2026 22:02:54 -0700 Subject: [PATCH] fix(core): add admin workspace token mint route --- tests/e2e/test_peer_visibility_mcp_staging.sh | 9 ++ .../handlers/admin_workspace_tokens.go | 72 +++++++++++++ .../handlers/admin_workspace_tokens_test.go | 102 ++++++++++++++++++ workspace-server/internal/router/router.go | 2 + 4 files changed, 185 insertions(+) create mode 100644 workspace-server/internal/handlers/admin_workspace_tokens.go create mode 100644 workspace-server/internal/handlers/admin_workspace_tokens_test.go diff --git a/tests/e2e/test_peer_visibility_mcp_staging.sh b/tests/e2e/test_peer_visibility_mcp_staging.sh index fdea4b959..79d8062c3 100755 --- a/tests/e2e/test_peer_visibility_mcp_staging.sh +++ b/tests/e2e/test_peer_visibility_mcp_staging.sh @@ -227,6 +227,15 @@ except Exception: print(''); sys.exit(0) print(d.get('auth_token') or d.get('connection', {}).get('auth_token') or '') " 2>/dev/null) [ -n "$WID" ] || fail "$rt workspace create failed: $(echo "$R" | head -c 300)" + if [ -z "$WTOK" ]; then + TTOK_RESP=$(tenant_call POST "/admin/workspaces/$WID/tokens" 2>/dev/null || true) + WTOK=$(echo "$TTOK_RESP" | python3 -c " +import sys, json +try: d = json.load(sys.stdin) +except Exception: print(''); sys.exit(0) +print(d.get('auth_token') or '') +" 2>/dev/null) + fi if [ -z "$WTOK" ]; then TTOK_RESP=$(tenant_call GET "/admin/workspaces/$WID/test-token" 2>/dev/null || true) WTOK=$(echo "$TTOK_RESP" | python3 -c " diff --git a/workspace-server/internal/handlers/admin_workspace_tokens.go b/workspace-server/internal/handlers/admin_workspace_tokens.go new file mode 100644 index 000000000..fb3e9d82f --- /dev/null +++ b/workspace-server/internal/handlers/admin_workspace_tokens.go @@ -0,0 +1,72 @@ +package handlers + +import ( + "database/sql" + "fmt" + "log" + "net/http" + + "github.com/Molecule-AI/molecule-monorepo/platform/internal/db" + "github.com/Molecule-AI/molecule-monorepo/platform/internal/wsauth" + "github.com/gin-gonic/gin" +) + +// AdminWorkspaceTokenHandler lets tenant admins mint the first workspace +// bearer for managed SaaS workspaces whose runtime receives its token later +// through registry registration. +type AdminWorkspaceTokenHandler struct{} + +func NewAdminWorkspaceTokenHandler() *AdminWorkspaceTokenHandler { + return &AdminWorkspaceTokenHandler{} +} + +// Create handles POST /admin/workspaces/:id/tokens. The route must be mounted +// behind AdminAuth; the plaintext token is returned exactly once. +func (h *AdminWorkspaceTokenHandler) Create(c *gin.Context) { + workspaceID := c.Param("id") + if !validWorkspaceID(workspaceID) { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid workspace id"}) + return + } + + var existing string + err := db.DB.QueryRowContext(c.Request.Context(), + `SELECT id FROM workspaces WHERE id = $1 AND status <> 'removed'`, + workspaceID).Scan(&existing) + if err != nil { + if err == sql.ErrNoRows { + c.JSON(http.StatusNotFound, gin.H{"error": "workspace not found"}) + return + } + log.Printf("admin workspace tokens: workspace lookup failed for %s: %v", workspaceID, err) + c.JSON(http.StatusInternalServerError, gin.H{"error": "workspace lookup failed"}) + return + } + + var count int + if err := db.DB.QueryRowContext(c.Request.Context(), + `SELECT COUNT(*) FROM workspace_auth_tokens WHERE workspace_id = $1 AND revoked_at IS NULL`, + workspaceID).Scan(&count); err != nil { + log.Printf("admin workspace tokens: count failed for %s: %v", workspaceID, err) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to count tokens"}) + return + } + if count >= maxTokensPerWorkspace { + c.JSON(http.StatusTooManyRequests, gin.H{"error": fmt.Sprintf("maximum %d active tokens per workspace", maxTokensPerWorkspace)}) + return + } + + token, err := wsauth.IssueToken(c.Request.Context(), db.DB, workspaceID) + if err != nil { + log.Printf("admin workspace tokens: issue failed for %s: %v", workspaceID, err) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to create token"}) + return + } + + log.Printf("admin workspace tokens: issued token for workspace %s", workspaceID) + c.JSON(http.StatusCreated, gin.H{ + "auth_token": token, + "workspace_id": workspaceID, + "message": "Save this token now — it cannot be retrieved again.", + }) +} diff --git a/workspace-server/internal/handlers/admin_workspace_tokens_test.go b/workspace-server/internal/handlers/admin_workspace_tokens_test.go new file mode 100644 index 000000000..7d59638b4 --- /dev/null +++ b/workspace-server/internal/handlers/admin_workspace_tokens_test.go @@ -0,0 +1,102 @@ +package handlers + +import ( + "encoding/json" + "errors" + "net/http" + "testing" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/gin-gonic/gin" +) + +func TestAdminWorkspaceTokenHandler_Create_HappyPath(t *testing.T) { + mock, cleanup := withMockDB(t) + defer cleanup() + + mock.ExpectQuery(`SELECT id FROM workspaces WHERE id = \$1 AND status <> 'removed'`). + WithArgs(wsUUID1). + WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(wsUUID1)) + mock.ExpectQuery(`SELECT COUNT\(\*\) FROM workspace_auth_tokens`). + WithArgs(wsUUID1). + WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0)) + mock.ExpectExec(`INSERT INTO workspace_auth_tokens`). + WithArgs(wsUUID1, sqlmock.AnyArg(), sqlmock.AnyArg()). + WillReturnResult(sqlmock.NewResult(1, 1)) + + w := makeReq(t, NewAdminWorkspaceTokenHandler().Create, "POST", + "/admin/workspaces/"+wsUUID1+"/tokens", gin.Params{{Key: "id", Value: wsUUID1}}) + + if w.Code != http.StatusCreated { + t.Fatalf("expected 201, got %d: %s", w.Code, w.Body.String()) + } + var body struct { + AuthToken string `json:"auth_token"` + WorkspaceID string `json:"workspace_id"` + } + if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil { + t.Fatalf("decode: %v", err) + } + if body.AuthToken == "" || body.WorkspaceID != wsUUID1 { + t.Fatalf("unexpected body: %+v", body) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet: %v", err) + } +} + +func TestAdminWorkspaceTokenHandler_Create_MissingWorkspace(t *testing.T) { + mock, cleanup := withMockDB(t) + defer cleanup() + + mock.ExpectQuery(`SELECT id FROM workspaces WHERE id = \$1 AND status <> 'removed'`). + WithArgs(wsUUID1). + WillReturnRows(sqlmock.NewRows([]string{"id"})) + + w := makeReq(t, NewAdminWorkspaceTokenHandler().Create, "POST", + "/admin/workspaces/"+wsUUID1+"/tokens", gin.Params{{Key: "id", Value: wsUUID1}}) + + if w.Code != http.StatusNotFound { + t.Fatalf("expected 404, got %d: %s", w.Code, w.Body.String()) + } +} + +func TestAdminWorkspaceTokenHandler_Create_RateLimited(t *testing.T) { + mock, cleanup := withMockDB(t) + defer cleanup() + + mock.ExpectQuery(`SELECT id FROM workspaces WHERE id = \$1 AND status <> 'removed'`). + WithArgs(wsUUID1). + WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(wsUUID1)) + mock.ExpectQuery(`SELECT COUNT\(\*\) FROM workspace_auth_tokens`). + WithArgs(wsUUID1). + WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(maxTokensPerWorkspace)) + + w := makeReq(t, NewAdminWorkspaceTokenHandler().Create, "POST", + "/admin/workspaces/"+wsUUID1+"/tokens", gin.Params{{Key: "id", Value: wsUUID1}}) + + if w.Code != http.StatusTooManyRequests { + t.Fatalf("expected 429, got %d: %s", w.Code, w.Body.String()) + } +} + +func TestAdminWorkspaceTokenHandler_Create_IssueFails(t *testing.T) { + mock, cleanup := withMockDB(t) + defer cleanup() + + mock.ExpectQuery(`SELECT id FROM workspaces WHERE id = \$1 AND status <> 'removed'`). + WithArgs(wsUUID1). + WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(wsUUID1)) + mock.ExpectQuery(`SELECT COUNT\(\*\) FROM workspace_auth_tokens`). + WithArgs(wsUUID1). + WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0)) + mock.ExpectExec(`INSERT INTO workspace_auth_tokens`). + WillReturnError(errors.New("disk full")) + + w := makeReq(t, NewAdminWorkspaceTokenHandler().Create, "POST", + "/admin/workspaces/"+wsUUID1+"/tokens", gin.Params{{Key: "id", Value: wsUUID1}}) + + if w.Code != http.StatusInternalServerError { + t.Fatalf("expected 500, got %d: %s", w.Code, w.Body.String()) + } +} diff --git a/workspace-server/internal/router/router.go b/workspace-server/internal/router/router.go index f0a61d915..7ce0c1834 100644 --- a/workspace-server/internal/router/router.go +++ b/workspace-server/internal/router/router.go @@ -397,6 +397,8 @@ func Setup(hub *ws.Hub, broadcaster *events.Broadcaster, prov *provisioner.Provi wsAuth.GET("/tokens", tokh.List) wsAuth.POST("/tokens", tokh.Create) wsAuth.DELETE("/tokens/:tokenId", tokh.Revoke) + adminTokH := handlers.NewAdminWorkspaceTokenHandler() + r.POST("/admin/workspaces/:id/tokens", middleware.AdminAuth(db.DB), adminTokH.Create) // Memory memh := handlers.NewMemoryHandler() -- 2.52.0