diff --git a/workspace-server/internal/handlers/handlers_test.go b/workspace-server/internal/handlers/handlers_test.go index 16351f05f..d67b73493 100644 --- a/workspace-server/internal/handlers/handlers_test.go +++ b/workspace-server/internal/handlers/handlers_test.go @@ -94,11 +94,12 @@ func setupTestDB(t *testing.T) sqlmock.Sqlmock { mockDB.Close() }) - // Disable SSRF checks for the duration of this test only. Restore - // the previous state via t.Cleanup so that TestIsSafeURL_* tests - // (which run with SSRF enabled) are not affected by state leak. - restore := setSSRFCheckForTest(false) - t.Cleanup(restore) + // Preserve SSRF state across this test. Individual tests that need + // SSRF disabled can call setSSRFCheckForTest(false) themselves. + // This prevents setupTestDB from silently overriding SSRF guards in + // regression tests that assert the guard is active (issue #807). + prevSSRF := ssrfCheckEnabled + t.Cleanup(func() { ssrfCheckEnabled = prevSSRF }) // The wsauth.platform_inbound_secret cache (#189) is package-level // state in another package — without a reset between tests, a diff --git a/workspace-server/internal/handlers/registry.go b/workspace-server/internal/handlers/registry.go index 9d7ffd58d..ec3112abc 100644 --- a/workspace-server/internal/handlers/registry.go +++ b/workspace-server/internal/handlers/registry.go @@ -3,6 +3,7 @@ package handlers import ( "context" "database/sql" + "encoding/json" "errors" "fmt" "log" @@ -363,6 +364,23 @@ func (h *RegistryHandler) Register(c *gin.Context) { } agentCardStr := string(reconciledCard) + // SSRF guard: the embedded agent_card.url is attacker-writable and is + // used by the transcript proxy regardless of delivery mode. Reject the + // same unsafe targets that UpdateCard rejects. Empty URL is allowed + // (poll-mode workspaces may omit it). #2130 #2132 + { + var cardURL struct { + URL string `json:"url"` + } + if jsonErr := json.Unmarshal(reconciledCard, &cardURL); jsonErr == nil && cardURL.URL != "" { + if err := isSafeURL(cardURL.URL); err != nil { + log.Printf("Registry register: workspace %s agent_card url rejected: %v", payload.ID, err) + c.JSON(http.StatusBadRequest, gin.H{"error": "workspace URL not allowed"}) + return + } + } + } + // urlForUpsert: poll-mode workspaces don't need a URL. Empty input // becomes NULL via sql.NullString so the row's URL stays clean (the // CASE below also preserves an existing provisioner-set URL, which @@ -873,6 +891,17 @@ func (h *RegistryHandler) UpdateCard(c *gin.Context) { return // response already written } + var card struct { + URL string `json:"url"` + } + if err := json.Unmarshal(payload.AgentCard, &card); err == nil && card.URL != "" { + if err := isSafeURL(card.URL); err != nil { + log.Printf("UpdateCard: workspace %s agent_card url rejected: %v", payload.WorkspaceID, err) + c.JSON(http.StatusBadRequest, gin.H{"error": "workspace URL not allowed"}) + return + } + } + agentCardStr := string(payload.AgentCard) _, err := db.DB.ExecContext(c.Request.Context(), ` UPDATE workspaces SET agent_card = $2::jsonb, updated_at = now() WHERE id = $1 diff --git a/workspace-server/internal/handlers/registry_test.go b/workspace-server/internal/handlers/registry_test.go index ed7f90467..22e380da4 100644 --- a/workspace-server/internal/handlers/registry_test.go +++ b/workspace-server/internal/handlers/registry_test.go @@ -627,6 +627,91 @@ func TestUpdateCard_DBError(t *testing.T) { } } +func TestUpdateCard_RejectsMetadataURL(t *testing.T) { + t.Setenv("MOLECULE_ENV", "production") + mock := setupTestDB(t) + restore := setSSRFCheckForTest(true) + t.Cleanup(restore) + setupTestRedis(t) + broadcaster := newTestBroadcaster() + handler := NewRegistryHandler(broadcaster) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + + body := `{"workspace_id":"ws-card","agent_card":{"name":"bad","url":"http://169.254.169.254/latest/meta-data/"}}` + c.Request = httptest.NewRequest("POST", "/registry/update-card", bytes.NewBufferString(body)) + c.Request.Header.Set("Content-Type", "application/json") + + handler.UpdateCard(c) + + if w.Code != http.StatusBadRequest { + t.Errorf("expected status 400, got %d: %s", w.Code, w.Body.String()) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet sqlmock expectations: %v", err) + } +} + +func TestUpdateCard_RejectsNonHTTPScheme(t *testing.T) { + mock := setupTestDB(t) + restore := setSSRFCheckForTest(true) + t.Cleanup(restore) + setupTestRedis(t) + broadcaster := newTestBroadcaster() + handler := NewRegistryHandler(broadcaster) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + + body := `{"workspace_id":"ws-card","agent_card":{"name":"bad","url":"file:///etc/passwd"}}` + c.Request = httptest.NewRequest("POST", "/registry/update-card", bytes.NewBufferString(body)) + c.Request.Header.Set("Content-Type", "application/json") + + handler.UpdateCard(c) + + if w.Code != http.StatusBadRequest { + t.Errorf("expected status 400, got %d: %s", w.Code, w.Body.String()) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Errorf("unmet sqlmock expectations: %v", err) + } +} + +// TestRegister_RejectsBadEmbeddedURL verifies that the agent_card.url field +// is validated in Register, not just UpdateCard. The URL is attacker-writable +// and used by the transcript proxy, so a metadata/loopback target must be +// rejected before the row is persisted. #2132 +func TestRegister_RejectsBadEmbeddedURL(t *testing.T) { + mock := setupTestDB(t) + restore := setSSRFCheckForTest(true) + t.Cleanup(restore) + setupTestRedis(t) + broadcaster := newTestBroadcaster() + handler := NewRegistryHandler(broadcaster) + + // Bootstrap path — no live tokens. + mock.ExpectQuery("SELECT COUNT\\(\\*\\) FROM workspace_auth_tokens"). + WithArgs("ws-bad-card-url"). + WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0)) + // resolveDeliveryMode: no row yet, default push. + mock.ExpectQuery(`SELECT delivery_mode, runtime FROM workspaces WHERE id`). + WithArgs("ws-bad-card-url"). + WillReturnError(sql.ErrNoRows) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + body := `{"id":"ws-bad-card-url","url":"http://example.com","agent_card":{"name":"bad","url":"http://169.254.169.254/latest/meta-data/"}}` + c.Request = httptest.NewRequest("POST", "/registry/register", bytes.NewBufferString(body)) + c.Request.Header.Set("Content-Type", "application/json") + + handler.Register(c) + + if w.Code != http.StatusBadRequest { + t.Errorf("expected 400 for bad embedded agent_card.url, got %d: %s", w.Code, w.Body.String()) + } +} + // TestRegister_GuardAgainstResurrectingRemovedRow verifies the #73 fix: // the ON CONFLICT UPSERT must carry a `WHERE status IS DISTINCT FROM 'removed'` // clause so that a late heartbeat from a workspace that was just deleted diff --git a/workspace-server/internal/handlers/transcript.go b/workspace-server/internal/handlers/transcript.go index 704842009..2b5287333 100644 --- a/workspace-server/internal/handlers/transcript.go +++ b/workspace-server/internal/handlers/transcript.go @@ -12,19 +12,21 @@ package handlers import ( "context" - "fmt" + "errors" "io" "log" - "net" "net/http" "net/url" - "strings" "time" "git.moleculesai.app/molecule-ai/molecule-core/workspace-server/internal/db" "github.com/gin-gonic/gin" ) +// errTranscriptRedirectBlocked is returned by the transcript proxy's +// CheckRedirect when a redirect target fails the SSRF policy. +var errTranscriptRedirectBlocked = errors.New("redirect target blocked by SSRF policy") + // TranscriptHandler proxies /workspaces/:id/transcript to the workspace agent. type TranscriptHandler struct { httpClient *http.Client @@ -32,7 +34,21 @@ type TranscriptHandler struct { func NewTranscriptHandler() *TranscriptHandler { return &TranscriptHandler{ - httpClient: &http.Client{Timeout: 15 * time.Second}, + httpClient: &http.Client{ + Timeout: 15 * time.Second, + CheckRedirect: func(req *http.Request, via []*http.Request) error { + // Go's default client follows up to 10 redirects and copies + // headers (including Authorization) to each hop. A redirect to + // an internal/metadata target would bypass the front-door + // isSafeURL check because that check only inspected the first + // URL. Re-validate every hop. #2130 #2132 + if err := isSafeURL(req.URL.String()); err != nil { + log.Printf("transcript: redirect to %s rejected: %v", req.URL.String(), err) + return errTranscriptRedirectBlocked + } + return nil + }, + }, } } @@ -66,7 +82,7 @@ func (h *TranscriptHandler) Get(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "invalid workspace URL"}) return } - if err := validateWorkspaceURL(target); err != nil { + if err := isSafeURL(target.String()); err != nil { log.Printf("transcript: workspace %s URL rejected: %v", workspaceID, err) c.JSON(http.StatusBadRequest, gin.H{"error": "workspace URL not allowed"}) return @@ -121,58 +137,3 @@ func (h *TranscriptHandler) Get(c *gin.Context) { } c.Data(resp.StatusCode, resp.Header.Get("Content-Type"), body) } - -// validateWorkspaceURL enforces that the agent_card URL is safe to -// proxy to. agent_card is attacker-writable via /registry/register so -// any workspace-token holder could otherwise point the URL at cloud -// metadata (169.254.169.254), the Docker host, or other internal -// services reachable from the platform container. -// -// Policy: -// - scheme must be http or https (no file://, gopher://, ftp://, etc.) -// - host must be present -// - block cloud metadata endpoints (IMDS, GCP, Azure) -// - block link-local IPs (169.254/16 IPv4, fe80::/10 IPv6) -// - loopback is allowed — local dev runs workspaces on 127.0.0.1 -// - Docker internal hostnames (host.docker.internal, *.molecule-core-net) -// are allowed; the whole threat model assumes the platform already -// trusts peers on that network -func validateWorkspaceURL(u *url.URL) error { - if u.Scheme != "http" && u.Scheme != "https" { - return fmt.Errorf("unsupported scheme %q", u.Scheme) - } - host := u.Hostname() - if host == "" { - return fmt.Errorf("empty host") - } - - // Hostname blocklist (pre-IP-parse — these are usually resolved by - // the HTTP stack, not by us). - lower := strings.ToLower(host) - for _, banned := range []string{ - "metadata.google.internal", - "metadata.azure.com", - "metadata", - } { - if lower == banned { - return fmt.Errorf("metadata hostname blocked: %s", host) - } - } - - // IP-literal checks. - if ip := net.ParseIP(host); ip != nil { - // IMDS / cloud metadata. - if ip.String() == "169.254.169.254" { - return fmt.Errorf("cloud metadata endpoint blocked") - } - // Link-local: IPv4 169.254.0.0/16, IPv6 fe80::/10. - if ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() { - return fmt.Errorf("link-local address blocked: %s", host) - } - // IPv6 unique local fd00::/8 — used by some IMDS implementations. - if ip.To4() == nil && len(ip) == net.IPv6len && ip[0] == 0xfd { - return fmt.Errorf("IPv6 unique-local address blocked: %s", host) - } - } - return nil -} diff --git a/workspace-server/internal/handlers/transcript_test.go b/workspace-server/internal/handlers/transcript_test.go index 5407893bd..19fae49d2 100644 --- a/workspace-server/internal/handlers/transcript_test.go +++ b/workspace-server/internal/handlers/transcript_test.go @@ -5,16 +5,12 @@ import ( "encoding/json" "net/http" "net/http/httptest" - "net/url" "testing" "github.com/DATA-DOG/go-sqlmock" "github.com/gin-gonic/gin" ) -// urlParse is a tiny wrapper so table-driven tests can keep their lines short. -func urlParse(s string) (*url.URL, error) { return url.Parse(s) } - // expectWorkspaceURLLookup programs the sqlmock to answer the SELECT that // TranscriptHandler.Get issues for `agent_card->>'url'`. Tests call this // instead of inserting real rows (we use sqlmock — there's no DB). @@ -50,6 +46,7 @@ func TestTranscript_WorkspaceNotFound(t *testing.T) { } func TestTranscript_ProxyForwardsAndReturnsBody(t *testing.T) { + allowLoopbackForTest(t) mock := setupTestDB(t) setupTestRedis(t) h := NewTranscriptHandler() @@ -90,6 +87,7 @@ func TestTranscript_ProxyForwardsAndReturnsBody(t *testing.T) { } func TestTranscript_ProxyPropagatesAllowlistedQueryParams(t *testing.T) { + allowLoopbackForTest(t) mock := setupTestDB(t) setupTestRedis(t) h := NewTranscriptHandler() @@ -115,12 +113,15 @@ func TestTranscript_ProxyPropagatesAllowlistedQueryParams(t *testing.T) { } // SSRF regression tests — see issue #272. agent_card->>'url' is attacker- -// writable via /registry/register so validateWorkspaceURL must reject +// writable via /registry/register so the production SSRF policy must reject // link-local / cloud-metadata / non-http(s) targets before the outbound // HTTP call fires. func TestTranscript_RejectsCloudMetadataIP(t *testing.T) { + t.Setenv("MOLECULE_ENV", "production") mock := setupTestDB(t) + restore := setSSRFCheckForTest(true) + t.Cleanup(restore) setupTestRedis(t) h := NewTranscriptHandler() @@ -138,6 +139,8 @@ func TestTranscript_RejectsCloudMetadataIP(t *testing.T) { func TestTranscript_RejectsNonHTTPScheme(t *testing.T) { mock := setupTestDB(t) + restore := setSSRFCheckForTest(true) + t.Cleanup(restore) setupTestRedis(t) h := NewTranscriptHandler() @@ -155,6 +158,8 @@ func TestTranscript_RejectsNonHTTPScheme(t *testing.T) { func TestTranscript_RejectsMetadataHostname(t *testing.T) { mock := setupTestDB(t) + restore := setSSRFCheckForTest(true) + t.Cleanup(restore) setupTestRedis(t) h := NewTranscriptHandler() @@ -171,7 +176,10 @@ func TestTranscript_RejectsMetadataHostname(t *testing.T) { } func TestTranscript_RejectsLinkLocalIPv6(t *testing.T) { + t.Setenv("MOLECULE_ENV", "production") mock := setupTestDB(t) + restore := setSSRFCheckForTest(true) + t.Cleanup(restore) setupTestRedis(t) h := NewTranscriptHandler() @@ -187,45 +195,62 @@ func TestTranscript_RejectsLinkLocalIPv6(t *testing.T) { } } -// validateWorkspaceURL unit tests — pure function, no DB/Redis needed. -func TestValidateWorkspaceURL(t *testing.T) { - cases := []struct { - name string - raw string - wantErr bool - }{ - {"http localhost allowed (dev)", "http://127.0.0.1:8000", false}, - {"https public allowed", "https://agent.example.com", false}, - {"docker internal allowed", "http://host.docker.internal:8000", false}, - {"IMDS IP rejected", "http://169.254.169.254", true}, - {"GCP metadata hostname rejected", "http://metadata.google.internal", true}, - {"Azure metadata rejected", "http://metadata.azure.com", true}, - {"file scheme rejected", "file:///etc/passwd", true}, - {"gopher rejected", "gopher://internal:70/", true}, - {"IPv6 link-local rejected", "http://[fe80::1]", true}, - {"IPv4 link-local multicast rejected", "http://224.0.0.1", true}, +func TestTranscript_RejectsLoopbackURL(t *testing.T) { + t.Setenv("MOLECULE_ENV", "production") + mock := setupTestDB(t) + restore := setSSRFCheckForTest(true) + t.Cleanup(restore) + setupTestRedis(t) + h := NewTranscriptHandler() + + wsID := expectWorkspaceURLLookup(mock, "http://127.0.0.1:8000/") + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: wsID}} + c.Request = httptest.NewRequest("GET", "/workspaces/"+wsID+"/transcript", nil) + h.Get(c) + if w.Code != http.StatusBadRequest { + t.Errorf("expected 400 for loopback target, got %d: %s", w.Code, w.Body.String()) } - for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { - u, parseErr := urlParse(tc.raw) - if parseErr != nil && !tc.wantErr { - t.Fatalf("parse error: %v", parseErr) - } - if parseErr != nil { - return // unparseable URLs are rejected upstream; not this function's job - } - err := validateWorkspaceURL(u) - if tc.wantErr && err == nil { - t.Errorf("expected error for %q, got nil", tc.raw) - } - if !tc.wantErr && err != nil { - t.Errorf("expected OK for %q, got %v", tc.raw, err) - } - }) +} + +// TestTranscript_BlocksRedirectToInternalHost verifies that a 302 redirect to +// an SSRF-forbidden target is not followed. The initial URL is safe (a test +// server), but the redirect target points to cloud metadata. Without the +// per-hop check the Authorization bearer would be forwarded to the metadata +// endpoint. #2132 +func TestTranscript_BlocksRedirectToInternalHost(t *testing.T) { + allowLoopbackForTest(t) + mock := setupTestDB(t) + restore := setSSRFCheckForTest(true) + t.Cleanup(restore) + setupTestRedis(t) + h := NewTranscriptHandler() + + // First server is a safe-looking target that redirects to IMDS. + stub := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Location", "http://169.254.169.254/latest/meta-data/") + w.WriteHeader(http.StatusFound) + })) + defer stub.Close() + + wsID := expectWorkspaceURLLookup(mock, stub.URL) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Params = gin.Params{{Key: "id", Value: wsID}} + c.Request = httptest.NewRequest("GET", "/workspaces/"+wsID+"/transcript", nil) + c.Request.Header.Set("Authorization", "Bearer test-token") + h.Get(c) + + if w.Code != http.StatusBadGateway { + t.Errorf("expected 502 when redirect is blocked, got %d: %s", w.Code, w.Body.String()) } } func TestTranscript_UnreachableWorkspaceReturns502(t *testing.T) { + allowLoopbackForTest(t) mock := setupTestDB(t) setupTestRedis(t) h := NewTranscriptHandler() @@ -255,6 +280,7 @@ func TestTranscript_UnreachableWorkspaceReturns502(t *testing.T) { // req.Header.Set("Authorization", c.GetHeader("Authorization")) // This test verifies the fix and acts as a regression guard. func TestTranscript_ForwardsAuthHeader(t *testing.T) { + allowLoopbackForTest(t) mock := setupTestDB(t) setupTestRedis(t) h := NewTranscriptHandler() @@ -308,6 +334,7 @@ func TestTranscript_ForwardsAuthHeader(t *testing.T) { // request. The workspace will return 401 in this case, which the proxy // faithfully relays — no silent upgrade of privilege. func TestTranscript_NoAuthHeader_PassesThrough(t *testing.T) { + allowLoopbackForTest(t) mock := setupTestDB(t) setupTestRedis(t) h := NewTranscriptHandler()