fix: validate Host header in graph_server request handlers

packages/bigquery-magics/bigquery_magics/graph_server.py

@@ -17,12 +17,18 @@
 import json
 import socketserver
 import threading
+import urllib.parse
 from typing import Any, Dict, List
 
 from google.cloud import bigquery
 
 from bigquery_magics import core
 
+# The graph widget only ever reaches this server over the loopback interface.
+# Requests carrying any other Host are cross-site (e.g. a DNS-rebinding page in
+# the notebook user's browser) and are refused.
+_LOOPBACK_HOSTS = frozenset({"localhost", "127.0.0.1", "::1"})
+
 
 def execute_node_expansion(params, request):
     return {"error": "Node expansion not yet implemented"}
@@ -354,6 +360,22 @@ class GraphServerHandler(http.server.SimpleHTTPRequestHandler):
     def log_message(self, format, *args):
         pass
 
+    def _host_is_loopback(self):
+        """Return True if the request targets the loopback interface.
+
+        Origin cannot be used here because the notebook page and this server
+        live on different ports, so legitimate widget traffic is cross-origin.
+        The Host header, however, is always a loopback name for that traffic.
+        """
+        host = self.headers.get("Host")
+        if not host:
+            return False
+        try:
+            hostname = urllib.parse.urlsplit(f"//{host}").hostname
+        except ValueError:
+            return False
+        return hostname in _LOOPBACK_HOSTS
+
     def do_json_response(self, data):
         self.send_response(200)
         self.send_header("Access-Control-Allow-Origin", "*")
@@ -406,10 +428,16 @@ def handle_post_node_expansion(self):
         )
 
     def do_GET(self):
+        if not self._host_is_loopback():
+            self.send_error(403, "Forbidden")
+            return
         assert self.path == GraphServer.endpoints["get_ping"]
         self.handle_get_ping()
 
     def do_POST(self):
+        if not self._host_is_loopback():
+            self.send_error(403, "Forbidden")
+            return
         if self.path == GraphServer.endpoints["post_ping"]:
             self.handle_post_ping()
         elif self.path == GraphServer.endpoints["post_node_expansion"]:

packages/bigquery-magics/tests/unit/test_graph_server.py

@@ -747,3 +747,42 @@ def test_convert_schema_shared_label():
     labels = {label["name"]: label for label in result["labels"]}
     assert "Person" in labels
     assert set(labels["Person"]["propertyDeclarationNames"]) == {"id", "name"}
+
+
+class TestGraphServerHostHeader(unittest.TestCase):
+    def setUp(self):
+        self.server = graph_server.GraphServer()
+        self.server_thread = self.server.init()
+
+    def tearDown(self):
+        self.server.stop_server()
+        self.server_thread.join()
+
+    def _route(self):
+        return self.server.build_route(graph_server.GraphServer.endpoints["get_ping"])
+
+    @pytest.mark.skipif(
+        graph_visualization is None, reason="Requires `spanner-graph-notebook`"
+    )
+    def test_loopback_host_allowed(self):
+        response = requests.get(self._route())
+        self.assertEqual(response.status_code, 200)
+
+    @pytest.mark.skipif(
+        graph_visualization is None, reason="Requires `spanner-graph-notebook`"
+    )
+    def test_non_loopback_host_rejected(self):
+        # A DNS-rebinding page reaches the loopback socket but carries the
+        # attacker's hostname in the Host header.
+        response = requests.get(self._route(), headers={"Host": "attacker.example"})
+        self.assertEqual(response.status_code, 403)
+
+    @pytest.mark.skipif(
+        graph_visualization is None, reason="Requires `spanner-graph-notebook`"
+    )
+    def test_non_loopback_host_rejected_post(self):
+        route = self.server.build_route(graph_server.GraphServer.endpoints["post_ping"])
+        response = requests.post(
+            route, json={"data": "ping"}, headers={"Host": "evil.example:1234"}
+        )
+        self.assertEqual(response.status_code, 403)