Merge pull request #2095 from Molecule-AI/fix/ssrf-discoverhostpeer-1484

fix(discovery): isSafeURL guard on registered URLs (closes #1484)

workspace-server/internal/handlers/discovery.go

@@ -102,6 +102,19 @@ func discoverHostPeer(ctx context.Context, c *gin.Context, targetID string) {
 		return
 	}
 
+	// #1484 SSRF defense-in-depth: the URL came from the workspaces table
+	// without any per-write validation (workspace runtimes register their
+	// own URLs via /registry/register, and a misbehaving / compromised
+	// runtime could register a 169.254.169.254 metadata URL). Validate
+	// before handing it to the caller, who might dispatch HTTP against it.
+	// Currently gated by the bearer-required Discover handler, but this
+	// guard makes discoverHostPeer safe regardless of upstream auth shape.
+	if err := isSafeURL(url.String); err != nil {
+		log.Printf("Discovery: rejecting unsafe registered URL for %s (#1484): %v", resolvedID, err)
+		c.JSON(http.StatusServiceUnavailable, gin.H{"error": "workspace URL failed safety check", "status": status})
+		return
+	}
+
 	db.CacheURL(ctx, resolvedID, url.String)
 	c.JSON(http.StatusOK, gin.H{
 		"id":     resolvedID,
@@ -172,6 +185,18 @@ func writeExternalWorkspaceURL(ctx context.Context, c *gin.Context, callerID, ta
 		outURL = strings.Replace(outURL, "127.0.0.1", "host.docker.internal", 1)
 		outURL = strings.Replace(outURL, "localhost", "host.docker.internal", 1)
 	}
+
+	// #1484 SSRF defense-in-depth — same rationale as discoverHostPeer.
+	// We validate the post-rewrite URL because the rewrite changes which
+	// host the caller would dispatch against (host.docker.internal is
+	// only reachable inside a docker network; isSafeURL accepts it but
+	// blocks a metadata IP that survived the rewrite untouched).
+	if err := isSafeURL(outURL); err != nil {
+		log.Printf("Discovery: rejecting unsafe external workspace URL for %s (#1484): %v", targetID, err)
+		c.JSON(http.StatusServiceUnavailable, gin.H{"error": "workspace URL failed safety check"})
+		return true
+	}
+
 	c.JSON(http.StatusOK, gin.H{"id": targetID, "url": outURL, "name": wsName})
 	return true
 }

workspace-server/internal/handlers/discovery_test.go

@@ -671,6 +671,94 @@ func TestWriteExternalWorkspaceURL_RewritesLocalhostForDockerCaller(t *testing.T
 	}
 }
 
+// --- #1484 SSRF defense-in-depth regression tests ---
+
+// TestDiscoverHostPeer_RejectsMetadataIPURL pins the #1484 guard:
+// even though discoverHostPeer is currently gated by a bearer-required
+// Discover handler, the function MUST refuse to hand back a URL that
+// resolves into the cloud-metadata range. setupTestDB disables SSRF
+// for normal tests, so we re-enable it here for the duration of the
+// case and provide a literal IP so the check doesn't depend on DNS.
+func TestDiscoverHostPeer_RejectsMetadataIPURL(t *testing.T) {
+	mock := setupTestDB(t)
+	setupTestRedis(t)
+	restoreSSRF := setSSRFCheckForTest(true)
+	t.Cleanup(restoreSSRF)
+
+	mock.ExpectQuery(`SELECT url, status, forwarded_to FROM workspaces WHERE id =`).
+		WithArgs("ws-bad").
+		WillReturnRows(sqlmock.NewRows([]string{"url", "status", "forwarded_to"}).
+			AddRow("http://169.254.169.254/latest/meta-data/", "online", nil))
+
+	w := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(w)
+	c.Request = httptest.NewRequest("GET", "/x", nil)
+
+	discoverHostPeer(context.Background(), c, "ws-bad")
+	if w.Code != http.StatusServiceUnavailable {
+		t.Errorf("expected 503 for metadata-IP URL, got %d: %s", w.Code, w.Body.String())
+	}
+	if !strings.Contains(w.Body.String(), "safety check") {
+		t.Errorf("response should mention 'safety check', got %s", w.Body.String())
+	}
+}
+
+// TestDiscoverHostPeer_AcceptsPublicURL is the positive counterpart —
+// a routable hostname (literal public-range IP, no DNS dependency)
+// passes through the new guard and returns 200. Without it, the
+// rejection test above could pass by falsely failing every URL.
+func TestDiscoverHostPeer_AcceptsPublicURL(t *testing.T) {
+	mock := setupTestDB(t)
+	setupTestRedis(t)
+	restoreSSRF := setSSRFCheckForTest(true)
+	t.Cleanup(restoreSSRF)
+
+	mock.ExpectQuery(`SELECT url, status, forwarded_to FROM workspaces WHERE id =`).
+		WithArgs("ws-good").
+		WillReturnRows(sqlmock.NewRows([]string{"url", "status", "forwarded_to"}).
+			AddRow("http://8.8.8.8/a2a", "online", nil))
+
+	w := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(w)
+	c.Request = httptest.NewRequest("GET", "/x", nil)
+
+	discoverHostPeer(context.Background(), c, "ws-good")
+	if w.Code != http.StatusOK {
+		t.Errorf("expected 200 for public-IP URL, got %d: %s", w.Code, w.Body.String())
+	}
+}
+
+// TestWriteExternalWorkspaceURL_RejectsMetadataIPURL is the parallel
+// guard for the external-runtime path. Same #1484 rationale as the
+// host-peer test above; covers writeExternalWorkspaceURL specifically.
+func TestWriteExternalWorkspaceURL_RejectsMetadataIPURL(t *testing.T) {
+	mock := setupTestDB(t)
+	setupTestRedis(t)
+	restoreSSRF := setSSRFCheckForTest(true)
+	t.Cleanup(restoreSSRF)
+
+	mock.ExpectQuery(`SELECT COALESCE\(url,''\) FROM workspaces WHERE id =`).
+		WithArgs("ws-ext").
+		WillReturnRows(sqlmock.NewRows([]string{"url"}).
+			AddRow("http://169.254.169.254/computeMetadata/v1/"))
+	// callerRuntime lookup happens before isSafeURL — must mock it.
+	mock.ExpectQuery(`SELECT COALESCE\(runtime,'langgraph'\) FROM workspaces WHERE id =`).
+		WithArgs("ws-caller").
+		WillReturnRows(sqlmock.NewRows([]string{"runtime"}).AddRow("langgraph"))
+
+	w := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(w)
+	c.Request = httptest.NewRequest("GET", "/x", nil)
+
+	handled := writeExternalWorkspaceURL(context.Background(), c, "ws-caller", "ws-ext", "Bad WS")
+	if !handled {
+		t.Fatal("expected handled=true (the function did write a 503)")
+	}
+	if w.Code != http.StatusServiceUnavailable {
+		t.Errorf("expected 503 for metadata-IP URL, got %d: %s", w.Code, w.Body.String())
+	}
+}
+
 // --- discoverHostPeer smoke (currently unreachable via Discover) ---
 
 func TestDiscoverHostPeer_Smoke_CacheHit(t *testing.T) {