diff --git a/packages/bigquery-magics/bigquery_magics/graph_server.py b/packages/bigquery-magics/bigquery_magics/graph_server.py index 74d3e937e173..4bdcb86d94a3 100644 --- a/packages/bigquery-magics/bigquery_magics/graph_server.py +++ b/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"]: diff --git a/packages/bigquery-magics/tests/unit/test_graph_server.py b/packages/bigquery-magics/tests/unit/test_graph_server.py index 2f4d10283560..afa7102d8fdd 100644 --- a/packages/bigquery-magics/tests/unit/test_graph_server.py +++ b/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().