molecule-ai/molecule-core
35
0
Fork 2
Code Issues 16 Pull Requests 2 Actions 49 Packages Projects Releases Wiki Activity

fix(#611): remove budget_limit from PATCH /workspaces/:id and strip financial fields from GET

Browse Source 
Security Auditor findings on PR #611:

Fix 1 (BLOCKING): Remove budget_limit handling from Update() entirely.
PATCH /workspaces/:id uses ValidateAnyToken — any enrolled workspace bearer
could self-clear its own spending ceiling. The dedicated AdminAuth-gated
PATCH /workspaces/:id/budget is the only authorised write path.

Fix 2 (MEDIUM): Strip budget_limit and monthly_spend from Get() response
before c.JSON(). GET /workspaces/:id is on the open router — any caller
with a valid UUID must not read billing data.

Also updates four existing tests in workspace_budget_test.go that encoded
the old (insecure) behaviour, and adds three new regression tests.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Molecule AI Backend Engineer 2026-04-17 06:11:11 +00:00
parent dd0b282c79
commit fce0be30fd
3 changed files with 178 additions and 50 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

35
platform/internal/handlers/workspace.go
View File

@ -419,6 +419,12 @@ func (h *WorkspaceHandler) Get(c *gin.Context) {
return
}
// Strip financial fields — GET /workspaces/:id is on the open router.
// Any caller with a valid UUID would otherwise read billing data.
// The dedicated budget/spend endpoints are AdminAuth-gated. (#611)
delete(ws, "budget_limit")
delete(ws, "monthly_spend")
c.JSON(http.StatusOK, ws)
}
@ -519,7 +525,10 @@ var sensitiveUpdateFields = map[string]struct{}{
"parent_id": {},
"runtime": {},
"workspace_dir": {},
"budget_limit": {}, // cost-control ceiling — requires admin auth to change
// budget_limit is intentionally NOT here. The dedicated
// PATCH /workspaces/:id/budget (AdminAuth) is the only write path.
// Accepting it here — even behind ValidateAnyToken — lets workspace agents
// self-clear their own spending ceiling. (#611 Security Auditor finding)
}
// Update handles PATCH /workspaces/:id
@ -617,26 +626,10 @@ func (h *WorkspaceHandler) Update(c *gin.Context) {
}
needsRestart = true
}
if budgetLimitVal, ok := body["budget_limit"]; ok {
// Allow null to clear (remove) the budget ceiling.
// Non-null values come in as JSON float64 from map[string]interface{}
// — convert to int64 for storage (USD cents).
var budgetArg interface{}
if budgetLimitVal != nil {
switch v := budgetLimitVal.(type) {
case float64:
budgetArg = int64(v)
case int64:
budgetArg = v
default:
c.JSON(http.StatusBadRequest, gin.H{"error": "budget_limit must be an integer (USD cents) or null"})
return
}
}
if _, err := db.DB.ExecContext(ctx, `UPDATE workspaces SET budget_limit = $2, updated_at = now() WHERE id = $1`, id, budgetArg); err != nil {
log.Printf("Update budget_limit error for %s: %v", id, err)
}
}
// NOTE: budget_limit is intentionally NOT handled here. The dedicated
// PATCH /workspaces/:id/budget (AdminAuth) is the only write path.
// This endpoint uses ValidateAnyToken — any enrolled workspace bearer
// could otherwise self-clear its own spending ceiling. (#611 Security Auditor)
// Update canvas position if both x and y provided
if x, xOk := body["x"]; xOk {

64
platform/internal/handlers/workspace_budget_test.go
View File

@ -34,10 +34,11 @@ var wsColumns = []string{
"budget_limit", "monthly_spend",
}
// ==================== GET — budget_limit serialisation ====================
// ==================== GET — financial fields stripped from open endpoint ====================
// TestWorkspaceBudget_Get_NilLimit verifies that budget_limit is null in the
// JSON response when the DB column IS NULL (no ceiling configured).
// TestWorkspaceBudget_Get_NilLimit verifies that budget_limit and monthly_spend
// are NOT present in GET /workspaces/:id. The endpoint is on the open router —
// any caller with a valid UUID must not read billing data. (#611 Security Auditor)
func TestWorkspaceBudget_Get_NilLimit(t *testing.T) {
mock := setupTestDB(t)
setupTestRedis(t)
@ -66,19 +67,21 @@ func TestWorkspaceBudget_Get_NilLimit(t *testing.T) {
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("failed to parse response: %v", err)
}
if resp["budget_limit"] != nil {
t.Errorf("expected budget_limit=nil, got %v", resp["budget_limit"])
// #611: financial fields must NOT appear on the open GET endpoint.
if _, present := resp["budget_limit"]; present {
t.Errorf("budget_limit must not appear in open GET /workspaces/:id response")
}
if resp["monthly_spend"] != float64(0) {
t.Errorf("expected monthly_spend=0, got %v", resp["monthly_spend"])
if _, present := resp["monthly_spend"]; present {
t.Errorf("monthly_spend must not appear in open GET /workspaces/:id response")
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("sqlmock expectations not met: %v", err)
}
}
// TestWorkspaceBudget_Get_WithLimit verifies that a non-NULL budget_limit is
// returned as the correct integer value (USD cents) in the response.
// TestWorkspaceBudget_Get_WithLimit verifies that budget_limit and monthly_spend
// are stripped from the open GET /workspaces/:id even when the DB has non-zero
// values. Financial reads go through the AdminAuth-gated budget endpoint. (#611)
func TestWorkspaceBudget_Get_WithLimit(t *testing.T) {
mock := setupTestDB(t)
setupTestRedis(t)
@ -91,8 +94,8 @@ func TestWorkspaceBudget_Get_WithLimit(t *testing.T) {
[]byte(`{}`), "http://localhost:9002",
nil, 0, 0.0, "", 0, "", "langgraph", "",
0.0, 0.0, false,
int64(500), // budget_limit = $5.00
int64(123))) // monthly_spend = $1.23
int64(500), // budget_limit = $5.00 in DB
int64(123))) // monthly_spend = $1.23 in DB
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
@ -107,11 +110,17 @@ func TestWorkspaceBudget_Get_WithLimit(t *testing.T) {
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("failed to parse response: %v", err)
}
if resp["budget_limit"] != float64(500) {
t.Errorf("expected budget_limit=500, got %v", resp["budget_limit"])
// #611: financial fields must NOT appear on the open GET endpoint even when
// the DB has non-zero values — they're stripped before c.JSON().
if _, present := resp["budget_limit"]; present {
t.Errorf("budget_limit must not appear in open GET /workspaces/:id response (got %v)", resp["budget_limit"])
}
if resp["monthly_spend"] != float64(123) {
t.Errorf("expected monthly_spend=123, got %v", resp["monthly_spend"])
if _, present := resp["monthly_spend"]; present {
t.Errorf("monthly_spend must not appear in open GET /workspaces/:id response (got %v)", resp["monthly_spend"])
}
// Confirm non-financial fields are still present.
if resp["name"] != "Capped Agent" {
t.Errorf("expected name 'Capped Agent', got %v", resp["name"])
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("sqlmock expectations not met: %v", err)
@ -163,23 +172,21 @@ func TestWorkspaceBudget_Create_WithLimit(t *testing.T) {
}
}
// ==================== PATCH — update budget_limit ====================
// ==================== PATCH — budget_limit silently ignored on general update ====================
// TestWorkspaceBudget_Update_SetLimit verifies that PATCH /workspaces/:id with
// budget_limit=500 issues an UPDATE workspaces SET budget_limit = 500.
// budget_limit=500 does NOT issue any DB write for budget_limit. The only write
// path is the AdminAuth-gated PATCH /workspaces/:id/budget endpoint. (#611)
func TestWorkspaceBudget_Update_SetLimit(t *testing.T) {
mock := setupTestDB(t)
setupTestRedis(t)
handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir())
// Existence probe
// Only the existence probe fires; no UPDATE for budget_limit.
mock.ExpectQuery("SELECT EXISTS.*workspaces WHERE id").
WithArgs("ws-upd-budget").
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true))
// budget_limit UPDATE
mock.ExpectExec("UPDATE workspaces SET budget_limit").
WithArgs("ws-upd-budget", int64(500)).
WillReturnResult(sqlmock.NewResult(0, 1))
// No ExpectExec for budget_limit — sqlmock will fail if one is issued.
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
@ -192,25 +199,24 @@ func TestWorkspaceBudget_Update_SetLimit(t *testing.T) {
if w.Code != http.StatusOK {
t.Errorf("expected 200, got %d: %s", w.Code, w.Body.String())
}
// If a budget_limit UPDATE was issued, sqlmock would have an unexpected call.
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("sqlmock expectations not met: %v", err)
t.Errorf("unexpected DB activity — budget_limit must not be written via general Update: %v", err)
}
}
// TestWorkspaceBudget_Update_ClearLimit verifies that PATCH /workspaces/:id
// with budget_limit=null issues an UPDATE with NULL, clearing the ceiling.
// with budget_limit=null does NOT issue any DB write for budget_limit. (#611)
func TestWorkspaceBudget_Update_ClearLimit(t *testing.T) {
mock := setupTestDB(t)
setupTestRedis(t)
handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir())
// Only the existence probe fires; no UPDATE for budget_limit.
mock.ExpectQuery("SELECT EXISTS.*workspaces WHERE id").
WithArgs("ws-clear-budget").
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true))
// NULL clears the budget ceiling
mock.ExpectExec("UPDATE workspaces SET budget_limit").
WithArgs("ws-clear-budget", nil).
WillReturnResult(sqlmock.NewResult(0, 1))
// No ExpectExec — a budget_limit write here would re-open the vulnerability.
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
@ -224,7 +230,7 @@ func TestWorkspaceBudget_Update_ClearLimit(t *testing.T) {
t.Errorf("expected 200, got %d: %s", w.Code, w.Body.String())
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("sqlmock expectations not met: %v", err)
t.Errorf("unexpected DB activity — budget_limit must not be written via general Update: %v", err)
}
}

129
platform/internal/handlers/workspace_test.go
View File

@ -860,3 +860,132 @@ func TestWorkspaceUpdate_SensitiveField_NoTokensYet_FailOpen(t *testing.T) {
t.Errorf("bootstrap fail-open: got %d, want 200 (%s)", w.Code, w.Body.String())
}
}
// ==================== #611 Security Auditor regressions ====================
// TestWorkspaceGet_FinancialFieldsStripped verifies that GET /workspaces/:id
// does NOT expose budget_limit or monthly_spend. The endpoint is on the open
// router — any caller with a UUID would otherwise read billing data. (#611 Fix 2)
func TestWorkspaceGet_FinancialFieldsStripped(t *testing.T) {
mock := setupTestDB(t)
setupTestRedis(t)
broadcaster := newTestBroadcaster()
handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir())
columns := []string{
"id", "name", "role", "tier", "status", "agent_card", "url",
"parent_id", "active_tasks", "last_error_rate", "last_sample_error",
"uptime_seconds", "current_task", "runtime", "workspace_dir", "x", "y", "collapsed",
"budget_limit", "monthly_spend",
}
// Populate with non-zero financial values to confirm they are stripped.
mock.ExpectQuery("SELECT w.id, w.name").
WithArgs("ws-fin-1").
WillReturnRows(sqlmock.NewRows(columns).
AddRow("ws-fin-1", "Finance Test", "worker", 1, "online", []byte(`{}`),
"http://localhost:9001", nil, 0, 0.0, "", 0, "", "langgraph",
"", 0.0, 0.0, false,
int64(50000), int64(12500))) // budget_limit=500 USD, spend=125 USD
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Params = gin.Params{{Key: "id", Value: "ws-fin-1"}}
c.Request = httptest.NewRequest("GET", "/workspaces/ws-fin-1", nil)
handler.Get(c)
if w.Code != http.StatusOK {
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
}
var resp map[string]interface{}
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("failed to parse response: %v", err)
}
if _, present := resp["budget_limit"]; present {
t.Errorf("budget_limit must not appear in GET /workspaces/:id response (got %v)", resp["budget_limit"])
}
if _, present := resp["monthly_spend"]; present {
t.Errorf("monthly_spend must not appear in GET /workspaces/:id response (got %v)", resp["monthly_spend"])
}
// Sanity-check that normal fields are still present.
if resp["name"] != "Finance Test" {
t.Errorf("expected name 'Finance Test', got %v", resp["name"])
}
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("unmet sqlmock expectations: %v", err)
}
}
// TestWorkspaceUpdate_BudgetLimitIgnored verifies that including budget_limit
// in a PATCH /workspaces/:id body does NOT trigger a DB write. The only write
// path for budget_limit is PATCH /workspaces/:id/budget (AdminAuth-gated).
// Any workspace bearer must not be able to self-clear its spending ceiling.
// (#611 Fix 1)
func TestWorkspaceUpdate_BudgetLimitIgnored(t *testing.T) {
mock := setupTestDB(t)
setupTestRedis(t)
broadcaster := newTestBroadcaster()
handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir())
// Only the existence probe fires — no UPDATE for budget_limit.
mock.ExpectQuery("SELECT EXISTS.*workspaces WHERE id").
WithArgs("ws-budget-test").
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true))
// name update is the only expected write
mock.ExpectExec("UPDATE workspaces SET name").
WithArgs("ws-budget-test", "Safe Name").
WillReturnResult(sqlmock.NewResult(0, 1))
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Params = gin.Params{{Key: "id", Value: "ws-budget-test"}}
// Send budget_limit alongside an innocuous field.
body := `{"name":"Safe Name","budget_limit":null}`
c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-budget-test",
bytes.NewBufferString(body))
c.Request.Header.Set("Content-Type", "application/json")
handler.Update(c)
if w.Code != http.StatusOK {
t.Errorf("expected 200, got %d: %s", w.Code, w.Body.String())
}
// sqlmock will fail if any unexpected DB call was made (e.g. for budget_limit).
if err := mock.ExpectationsWereMet(); err != nil {
t.Errorf("unexpected DB call — budget_limit must not be written via Update: %v", err)
}
}
// TestWorkspaceUpdate_BudgetLimitOnly_Ignored verifies that a body containing
// ONLY budget_limit results in no DB writes at all (besides the existence probe).
func TestWorkspaceUpdate_BudgetLimitOnly_Ignored(t *testing.T) {
mock := setupTestDB(t)
setupTestRedis(t)
broadcaster := newTestBroadcaster()
handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir())
mock.ExpectQuery("SELECT EXISTS.*workspaces WHERE id").
WithArgs("ws-budget-only").
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true))
// No UPDATE expected — budget_limit must be silently skipped.
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Params = gin.Params{{Key: "id", Value: "ws-budget-only"}}
c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-budget-only",
bytes.NewBufferString(`{"budget_limit":999999}`))
c.Request.Header.Set("Content-Type", "application/json")
handler.Update(c)
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("unexpected DB call for budget_limit: %v", err)
}
}