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"}
@@ -358,6 +364,39 @@ class GraphServerHandler(http.server.SimpleHTTPRequestHandler):
     def log_message(self, format, *args):
         pass
 
+    def _is_request_safe(self) -> bool:
+        """Return True if the request targets and originates from loopback.
+
+        The Host header guards against DNS rebinding: legitimate widget traffic
+        always carries a loopback name even though the notebook page and this
+        server live on different ports. Origin and Referer, when present, are
+        checked too so a page on a non-loopback site cannot drive the server
+        with a direct cross-site request; their port differs from ours but the
+        hostname is still loopback for real traffic.
+        """
+        host = self.headers.get("Host")
+        if not host:
+            return False
+        try:
+            hostname = urllib.parse.urlsplit(f"//{host.strip()}").hostname
+        except ValueError:
+            return False
+        if hostname not in _LOOPBACK_HOSTS:
+            return False
+
+        for header_name in ("Origin", "Referer"):
+            value = self.headers.get(header_name)
+            if not value:
+                continue
+            try:
+                origin_hostname = urllib.parse.urlsplit(value.strip()).hostname
+            except ValueError:
+                return False
+            if origin_hostname not in _LOOPBACK_HOSTS:
+                return False
+
+        return True
+
     def do_json_response(self, data):
         self.send_response(200)
         self.send_header("Access-Control-Allow-Origin", "*")
@@ -410,10 +449,16 @@ def handle_post_node_expansion(self):
         )
 
     def do_GET(self):
+        if not self._is_request_safe():
+            self.send_error(403, "Forbidden")
+            return
         assert self.path == GraphServer.endpoints["get_ping"]
         self.handle_get_ping()
 
     def do_POST(self):
+        if not self._is_request_safe():
+            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

@@ -750,6 +750,100 @@ def test_convert_schema_shared_label():
     assert set(labels["Person"]["propertyDeclarationNames"]) == {"id", "name"}
 
 
+@pytest.mark.skipif(
+    graph_visualization is None, reason="Requires `spanner-graph-notebook`"
+)
+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"])
+
+    def test_loopback_host_allowed(self):
+        response = requests.get(self._route())
+        self.assertEqual(response.status_code, 200)
+
+    def test_loopback_host_allowed_post(self):
+        route = self.server.build_route(graph_server.GraphServer.endpoints["post_ping"])
+        response = requests.post(route, json={"data": "ping"})
+        self.assertEqual(response.status_code, 200)
+
+    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)
+
+    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)
+
+    def test_missing_host_rejected(self):
+        handler = mock.Mock(spec=graph_server.GraphServerHandler)
+        handler.headers = {}
+        self.assertFalse(graph_server.GraphServerHandler._is_request_safe(handler))
+
+    def test_malformed_host_rejected(self):
+        # An unterminated IPv6 literal makes urlsplit raise ValueError.
+        handler = mock.Mock(spec=graph_server.GraphServerHandler)
+        handler.headers = {"Host": "[::1"}
+        self.assertFalse(graph_server.GraphServerHandler._is_request_safe(handler))
+
+    def test_non_loopback_origin_rejected(self):
+        # Host header is valid (loopback), but Origin is a malicious third-party site.
+        response = requests.get(
+            self._route(),
+            headers={"Host": "localhost", "Origin": "http://attacker.example.com"},
+        )
+        self.assertEqual(response.status_code, 403)
+
+    def test_non_loopback_referer_rejected(self):
+        # Host header is valid (loopback), but Referer reveals a malicious third-party origin.
+        response = requests.get(
+            self._route(),
+            headers={
+                "Host": "127.0.0.1",
+                "Referer": "http://attacker.example.com/exploit-page",
+            },
+        )
+        self.assertEqual(response.status_code, 403)
+
+    def test_malformed_origin_rejected(self):
+        # An unterminated IPv6 literal in Origin makes urlsplit raise ValueError.
+        handler = mock.Mock(spec=graph_server.GraphServerHandler)
+        handler.headers = {"Host": "localhost", "Origin": "http://[::1"}
+        self.assertFalse(graph_server.GraphServerHandler._is_request_safe(handler))
+
+    def test_loopback_origin_and_referer_allowed(self):
+        response = requests.get(
+            self._route(),
+            headers={
+                "Host": "localhost",
+                "Origin": "http://localhost:8080",
+                "Referer": "http://127.0.0.1:8080/notebook",
+            },
+        )
+        self.assertEqual(response.status_code, 200)
+
+    def test_null_origin_rejected(self):
+        # Sandboxed iframes and data: URIs send "Origin: null", where
+        # urlsplit("null").hostname is None.
+        response = requests.get(
+            self._route(),
+            headers={"Host": "localhost", "Origin": "null"},
+        )
+        self.assertEqual(response.status_code, 403)
+
+
 def test_stop_server_immediately_after_init():
     """Verifies that GraphServer.init() waits until the underlying TCP server
     is initialized before returning, preventing a race condition with stop_server().