From 2df644f528d056dfb8aa28d280b7cf817be2e47e Mon Sep 17 00:00:00 2001 From: Hongming Wang Date: Wed, 22 Apr 2026 19:40:06 -0700 Subject: [PATCH] =?UTF-8?q?fix(handlers):=20unblock=20Platform=20(Go)=20CI?= =?UTF-8?q?=20=E2=80=94=20sqlmock=20budget-check=20+=20test=20loopback?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes 14 of the 18 failing tests that have been reddening Platform (Go) CI on main since the 2026-04-18 open-source restructure + 2026-04-21 SSRF-backport. Reduces handlers package failure count 18 → 4 (remaining 4 are unrelated schema/behavior drift — see follow-ups). Three root causes fixed: 1. httptest.NewServer binds to 127.0.0.1; isSafeURL rejects loopback. Tests that stub workspace URLs via httptest therefore 502'd at the SSRF guard before reaching the handler logic they wanted to exercise. Fix: add `testAllowLoopback` var to ssrf.go + `allowLoopbackForTest(t)` helper in handlers_test.go. Only 127.0.0.0/8 and ::1 are relaxed; 169.254 metadata, RFC-1918, TEST-NET, CGNAT, and link-local protections remain active. Flag is paired with t.Cleanup and is never touched by production code. 2. ProxyA2A's checkWorkspaceBudget query (SELECT budget_limit, COALESCE (monthly_spend, 0) FROM workspaces WHERE id = $1) was added with the restructure but the a2a_proxy_test.go sqlmock expectations never caught up, producing "call to Query ... was not expected" on every ProxyA2A-exercising test. Fix: `expectBudgetCheck(mock, workspaceID)` helper that registers an empty-rows expectation (checkWorkspaceBudget fails-open on sql.ErrNoRows, so an empty result = "no budget limit"). Added to each of the 8 affected TestProxyA2A_* tests in the correct position relative to access-control + activity-log expectations. 3. TestAdminMemories_Import_Success + _RedactsSecretsBeforeDedup mocked a 5-arg INSERT when the handler actually issues a 4-arg INSERT (workspace_id, content, scope, namespace) unless the payload carries a created_at override. Removed the spurious 5th AnyArg from both tests; _PreservesCreatedAt is untouched since it legitimately uses the 5-arg form. Also: TestResolveAgentURL_CacheHit and _CacheMissDBHit used bogus `cached.example` / `dbhit.example` hostnames that fail DNS resolution inside isSafeURL (which happens BEFORE the loopback check). Swapped to `127.0.0.1` variants preserving test intent (they never hit the network). Remaining 4 failures — out of scope for this PR, tracked separately: - TestGitHubToken_NoTokenProvider (handler behavior drift — 500 vs 404) - TestWorkspaceList + TestWorkspaceList_WithData (Scan arg count — workspaces table gained a column, mock not updated) - TestRegister_ProvisionerURLPreserved (request body shape drift) Closes the 4 wrong-target PRs (#1710, #1718, #1719, #1664) that all tried to silence the symptom by disabling golangci-lint — which has `continue-on-error: true` in ci.yml and was never the actual blocker. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../internal/handlers/a2a_proxy_test.go | 42 ++++++++++++++++--- .../internal/handlers/admin_memories_test.go | 9 ++-- .../handlers/handlers_additional_test.go | 2 + .../internal/handlers/handlers_test.go | 30 +++++++++++++ workspace-server/internal/handlers/ssrf.go | 19 +++++++-- 5 files changed, 88 insertions(+), 14 deletions(-) diff --git a/workspace-server/internal/handlers/a2a_proxy_test.go b/workspace-server/internal/handlers/a2a_proxy_test.go index 438e4c06..89ca8029 100644 --- a/workspace-server/internal/handlers/a2a_proxy_test.go +++ b/workspace-server/internal/handlers/a2a_proxy_test.go @@ -22,11 +22,13 @@ import ( func TestProxyA2A_InvalidJSON(t *testing.T) { mock := setupTestDB(t) mr := setupTestRedis(t) + allowLoopbackForTest(t) broadcaster := newTestBroadcaster() handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir()) // Cache a URL so the handler doesn't fall back to DB mr.Set(fmt.Sprintf("ws:%s:url", "ws-badjson"), "http://localhost:9999") + expectBudgetCheck(mock, "ws-badjson") w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) @@ -59,6 +61,7 @@ func TestProxyA2A_InvalidJSON(t *testing.T) { func TestProxyA2A_AlreadyWrappedJSONRPC(t *testing.T) { mock := setupTestDB(t) mr := setupTestRedis(t) + allowLoopbackForTest(t) broadcaster := newTestBroadcaster() handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir()) @@ -73,6 +76,7 @@ func TestProxyA2A_AlreadyWrappedJSONRPC(t *testing.T) { defer agentServer.Close() mr.Set(fmt.Sprintf("ws:%s:url", "ws-wrapped"), agentServer.URL) + expectBudgetCheck(mock, "ws-wrapped") // Expect async activity log mock.ExpectExec("INSERT INTO activity_logs"). @@ -114,6 +118,7 @@ func TestProxyA2A_AlreadyWrappedJSONRPC(t *testing.T) { func TestProxyA2A_DBLookupFallback(t *testing.T) { mock := setupTestDB(t) setupTestRedis(t) // empty Redis — no cached URL + allowLoopbackForTest(t) broadcaster := newTestBroadcaster() handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir()) @@ -124,6 +129,9 @@ func TestProxyA2A_DBLookupFallback(t *testing.T) { })) defer agentServer.Close() + // Budget check runs first (before URL resolution) + expectBudgetCheck(mock, "ws-db-fallback") + // Redis miss → DB lookup → returns URL mock.ExpectQuery("SELECT url, status FROM workspaces WHERE id ="). WithArgs("ws-db-fallback"). @@ -162,6 +170,9 @@ func TestProxyA2A_DBLookupError(t *testing.T) { broadcaster := newTestBroadcaster() handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir()) + // Budget check runs first (before URL resolution) + expectBudgetCheck(mock, "ws-dberr") + // Redis miss → DB lookup → error mock.ExpectQuery("SELECT url, status FROM workspaces WHERE id ="). WithArgs("ws-dberr"). @@ -191,6 +202,7 @@ func TestProxyA2A_DBLookupError(t *testing.T) { func TestProxyA2A_AgentReturnsError(t *testing.T) { mock := setupTestDB(t) mr := setupTestRedis(t) + allowLoopbackForTest(t) broadcaster := newTestBroadcaster() handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir()) @@ -202,6 +214,7 @@ func TestProxyA2A_AgentReturnsError(t *testing.T) { defer agentServer.Close() mr.Set(fmt.Sprintf("ws:%s:url", "ws-agent-err"), agentServer.URL) + expectBudgetCheck(mock, "ws-agent-err") // Expect async activity log (with "error" status since agent returned 500) mock.ExpectExec("INSERT INTO activity_logs"). @@ -234,6 +247,7 @@ func TestProxyA2A_AgentReturnsError(t *testing.T) { func TestProxyA2A_MessageIDInjected(t *testing.T) { mock := setupTestDB(t) mr := setupTestRedis(t) + allowLoopbackForTest(t) broadcaster := newTestBroadcaster() handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir()) @@ -246,6 +260,7 @@ func TestProxyA2A_MessageIDInjected(t *testing.T) { defer agentServer.Close() mr.Set(fmt.Sprintf("ws:%s:url", "ws-msgid"), agentServer.URL) + expectBudgetCheck(mock, "ws-msgid") mock.ExpectExec("INSERT INTO activity_logs"). WillReturnResult(sqlmock.NewResult(0, 1)) @@ -284,6 +299,7 @@ func TestProxyA2A_MessageIDInjected(t *testing.T) { func TestProxyA2A_CallerIDPropagated(t *testing.T) { mock := setupTestDB(t) mr := setupTestRedis(t) + allowLoopbackForTest(t) broadcaster := newTestBroadcaster() handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir()) @@ -303,6 +319,8 @@ func TestProxyA2A_CallerIDPropagated(t *testing.T) { WithArgs("ws-target"). WillReturnRows(sqlmock.NewRows([]string{"id", "parent_id"}).AddRow("ws-target", "ws-parent")) + expectBudgetCheck(mock, "ws-target") + // Expect activity log with source_id set mock.ExpectExec("INSERT INTO activity_logs"). WillReturnResult(sqlmock.NewResult(0, 1)) @@ -377,6 +395,7 @@ func TestProxyA2A_AccessDenied_DifferentParents(t *testing.T) { func TestProxyA2A_AllowedSelf_SkipsAccessCheck(t *testing.T) { mock := setupTestDB(t) mr := setupTestRedis(t) + allowLoopbackForTest(t) broadcaster := newTestBroadcaster() handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir()) @@ -386,6 +405,7 @@ func TestProxyA2A_AllowedSelf_SkipsAccessCheck(t *testing.T) { })) defer agentServer.Close() mr.Set(fmt.Sprintf("ws:%s:url", "ws-self"), agentServer.URL) + expectBudgetCheck(mock, "ws-self") mock.ExpectExec("INSERT INTO activity_logs").WillReturnResult(sqlmock.NewResult(0, 1)) @@ -659,6 +679,7 @@ func TestProxyA2AError_BusyShape(t *testing.T) { func TestProxyA2A_BodyReadFailure_DeliveryConfirmed(t *testing.T) { mock := setupTestDB(t) mr := setupTestRedis(t) + allowLoopbackForTest(t) broadcaster := newTestBroadcaster() handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir()) @@ -687,6 +708,7 @@ func TestProxyA2A_BodyReadFailure_DeliveryConfirmed(t *testing.T) { wsID := "ws-bodyreadfail" mr.Set(fmt.Sprintf("ws:%s:url", wsID), agentServer.URL) + expectBudgetCheck(mock, wsID) // Expect async activity log INSERT (logA2ASuccess is called because // delivery_confirmed is true and the handler detected a 2xx status). @@ -941,14 +963,18 @@ func TestNormalizeA2APayload_MissingMethodReturnsEmpty(t *testing.T) { func TestResolveAgentURL_CacheHit(t *testing.T) { setupTestDB(t) mr := setupTestRedis(t) + allowLoopbackForTest(t) handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir()) - mr.Set("ws:ws-cached:url", "http://cached.example/a2a") + // Use loopback IP (unlocked by allowLoopbackForTest) so isSafeURL passes — + // cached.example does not resolve and would trip the DNS guard. + cached := "http://127.0.0.1:9999/a2a" + mr.Set("ws:ws-cached:url", cached) url, perr := handler.resolveAgentURL(context.Background(), "ws-cached") if perr != nil { t.Fatalf("unexpected error: %+v", perr) } - if url != "http://cached.example/a2a" { + if url != cached { t.Errorf("got %q, want cached URL", url) } } @@ -956,21 +982,24 @@ func TestResolveAgentURL_CacheHit(t *testing.T) { func TestResolveAgentURL_CacheMissDBHit(t *testing.T) { mock := setupTestDB(t) mr := setupTestRedis(t) + allowLoopbackForTest(t) handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir()) + // Use loopback IP (unlocked by allowLoopbackForTest) so isSafeURL passes. + dbURL := "http://127.0.0.1:9998" mock.ExpectQuery("SELECT url, status FROM workspaces WHERE id ="). WithArgs("ws-dbhit"). - WillReturnRows(sqlmock.NewRows([]string{"url", "status"}).AddRow("http://dbhit.example", "online")) + WillReturnRows(sqlmock.NewRows([]string{"url", "status"}).AddRow(dbURL, "online")) url, perr := handler.resolveAgentURL(context.Background(), "ws-dbhit") if perr != nil { t.Fatalf("unexpected error: %+v", perr) } - if url != "http://dbhit.example" { - t.Errorf("got %q, want http://dbhit.example", url) + if url != dbURL { + t.Errorf("got %q, want %q", url, dbURL) } // Verify cached now - if v, err := mr.Get("ws:ws-dbhit:url"); err != nil || v != "http://dbhit.example" { + if v, err := mr.Get("ws:ws-dbhit:url"); err != nil || v != dbURL { t.Errorf("expected Redis cache populated; got v=%q err=%v", v, err) } } @@ -1020,6 +1049,7 @@ func TestResolveAgentURL_DockerRewrite(t *testing.T) { // covered by TestResolveAgentURL_DockerRewrite_NilProvisionerNoRewrite. mr := setupTestRedis(t) setupTestDB(t) + allowLoopbackForTest(t) handler := NewWorkspaceHandler(newTestBroadcaster(), nil, "http://localhost:8080", t.TempDir()) mr.Set("ws:ws-dock:url", "http://127.0.0.1:55555") diff --git a/workspace-server/internal/handlers/admin_memories_test.go b/workspace-server/internal/handlers/admin_memories_test.go index 8e3920d7..b2881190 100644 --- a/workspace-server/internal/handlers/admin_memories_test.go +++ b/workspace-server/internal/handlers/admin_memories_test.go @@ -182,9 +182,9 @@ func TestAdminMemories_Import_Success(t *testing.T) { WithArgs("ws-uuid-1", sqlmock.AnyArg(), "LOCAL"). WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(false)) - // Insert succeeds. + // Insert succeeds. Handler uses 4-arg INSERT when created_at is absent. mock.ExpectExec("INSERT INTO agent_memories"). - WithArgs("ws-uuid-1", sqlmock.AnyArg(), "LOCAL", "general", sqlmock.AnyArg()). + WithArgs("ws-uuid-1", sqlmock.AnyArg(), "LOCAL", "general"). WillReturnResult(sqlmock.NewResult(1, 1)) w := adminPost(t, h, []map[string]interface{}{ @@ -326,9 +326,10 @@ func TestAdminMemories_Import_RedactsSecretsBeforeDedup(t *testing.T) { WithArgs("ws-uuid-1", redacted, "LOCAL"). WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(false)) - // Insert — receives the redacted content (not raw). + // Insert — receives the redacted content (not raw). Handler uses the + // 4-arg INSERT when created_at is absent from the payload. mock.ExpectExec("INSERT INTO agent_memories"). - WithArgs("ws-uuid-1", redacted, "LOCAL", "general", sqlmock.AnyArg()). + WithArgs("ws-uuid-1", redacted, "LOCAL", "general"). WillReturnResult(sqlmock.NewResult(1, 1)) w := adminPost(t, h, []map[string]interface{}{ diff --git a/workspace-server/internal/handlers/handlers_additional_test.go b/workspace-server/internal/handlers/handlers_additional_test.go index a2468c0f..1f6b152c 100644 --- a/workspace-server/internal/handlers/handlers_additional_test.go +++ b/workspace-server/internal/handlers/handlers_additional_test.go @@ -399,11 +399,13 @@ func TestProxyA2A_WorkspaceNoURL(t *testing.T) { func TestProxyA2A_AgentUnreachable(t *testing.T) { mock := setupTestDB(t) mr := setupTestRedis(t) + allowLoopbackForTest(t) broadcaster := newTestBroadcaster() handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", t.TempDir()) // Point to an unreachable address mr.Set(fmt.Sprintf("ws:%s:url", "ws-dead"), "http://127.0.0.1:1") + expectBudgetCheck(mock, "ws-dead") // Expect workspace name query for error activity log mock.ExpectQuery("SELECT name FROM workspaces WHERE id ="). diff --git a/workspace-server/internal/handlers/handlers_test.go b/workspace-server/internal/handlers/handlers_test.go index d5a56d19..13441fa5 100644 --- a/workspace-server/internal/handlers/handlers_test.go +++ b/workspace-server/internal/handlers/handlers_test.go @@ -55,6 +55,34 @@ func newTestBroadcaster() *events.Broadcaster { return events.NewBroadcaster(hub) } +// allowLoopbackForTest flips the ssrf.go testAllowLoopback escape hatch +// for the duration of the test, so httptest.NewServer's loopback URLs +// don't trip the SSRF guard. The 169.254 metadata, RFC-1918, TEST-NET, +// CGNAT, and link-local guards stay active — only 127.0.0.0/8 and ::1 +// are relaxed. Always paired with t.Cleanup to restore; multiple +// parallel tests won't race because Go test flips it sequentially per +// test unless t.Parallel() is used, and these tests don't parallelize. +func allowLoopbackForTest(t *testing.T) { + t.Helper() + prev := testAllowLoopback + testAllowLoopback = true + t.Cleanup(func() { testAllowLoopback = prev }) +} + +// expectBudgetCheck adds the sqlmock expectation for the budget-check +// query that ProxyA2A runs before forwarding. checkWorkspaceBudget +// fails-open on sql.ErrNoRows, so we return a deliberately-empty +// result — budget_limit NULL + monthly_spend 0 means "no limit". +// All a2a_proxy_test.go tests that run ProxyA2A (not just +// dispatchA2A unit tests) need this expectation; it was added to the +// handler in the 2026-04-18 restructure but the tests never caught up, +// leaving Platform (Go) CI red for weeks. +func expectBudgetCheck(mock sqlmock.Sqlmock, workspaceID string) { + mock.ExpectQuery(`SELECT budget_limit, COALESCE\(monthly_spend, 0\) FROM workspaces WHERE id = \$1`). + WithArgs(workspaceID). + WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"})) +} + // ---------- TestRegisterHandler ---------- func TestRegisterHandler(t *testing.T) { @@ -385,6 +413,7 @@ func TestWorkspaceList(t *testing.T) { func TestProxyA2A_JSONRPCWrapping(t *testing.T) { mock := setupTestDB(t) mr := setupTestRedis(t) + allowLoopbackForTest(t) broadcaster := newTestBroadcaster() handler := NewWorkspaceHandler(broadcaster, nil, "http://localhost:8080", "/tmp/configs") @@ -400,6 +429,7 @@ func TestProxyA2A_JSONRPCWrapping(t *testing.T) { // Cache the agent URL in Redis so the handler finds it mr.Set(fmt.Sprintf("ws:%s:url", "ws-proxy"), agentServer.URL) + expectBudgetCheck(mock, "ws-proxy") // Expect async activity log INSERT from the LogActivity goroutine mock.ExpectExec("INSERT INTO activity_logs"). diff --git a/workspace-server/internal/handlers/ssrf.go b/workspace-server/internal/handlers/ssrf.go index 67af118d..42e3ff3e 100644 --- a/workspace-server/internal/handlers/ssrf.go +++ b/workspace-server/internal/handlers/ssrf.go @@ -30,7 +30,7 @@ func isSafeURL(rawURL string) error { return fmt.Errorf("empty hostname") } if ip := net.ParseIP(host); ip != nil { - if ip.IsLoopback() || ip.IsUnspecified() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || ip.IsInterfaceLocalMulticast() { + if (ip.IsLoopback() && !testAllowLoopback) || ip.IsUnspecified() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || ip.IsInterfaceLocalMulticast() { return fmt.Errorf("forbidden loopback/unspecified/link-local IP: %s", ip) } if isPrivateOrMetadataIP(ip) { @@ -50,7 +50,7 @@ func isSafeURL(rawURL string) error { if ip == nil { continue } - if ip.IsLoopback() || ip.IsUnspecified() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || ip.IsInterfaceLocalMulticast() { + if (ip.IsLoopback() && !testAllowLoopback) || ip.IsUnspecified() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || ip.IsInterfaceLocalMulticast() { return fmt.Errorf("hostname %s resolves to forbidden link-local/loopback IP: %s", host, ip) } if isPrivateOrMetadataIP(ip) { @@ -60,6 +60,16 @@ func isSafeURL(rawURL string) error { return nil } +// testAllowLoopback is a test-only escape hatch. When true, isSafeURL +// accepts 127.0.0.0/8 and ::1 so unit tests that stub workspace URLs +// with httptest.NewServer (which binds to loopback) can reach their +// own mock backends. Flipped via allowLoopbackForTest(t) in tests — +// never set in production code paths. +// +// The 169.254 metadata, RFC-1918, TEST-NET, CGNAT, and link-local +// guards are NOT relaxed by this flag — only loopback. +var testAllowLoopback = false + // isPrivateOrMetadataIP returns true for IPs that must not be reached via A2A. // // Always blocked (both modes): @@ -106,8 +116,9 @@ func isPrivateOrMetadataIP(ip net.IP) bool { } // IPv6 path — .To4() was nil so this is a real v6 address. - // ::1 (loopback) — treat as blocked here too for defense-in-depth. - if ip.IsLoopback() { + // ::1 (loopback) — treat as blocked here too for defense-in-depth, + // unless tests have opted into loopback via testAllowLoopback. + if ip.IsLoopback() && !testAllowLoopback { return true } // Link-local fe80::/10 — always blocked.