fix: validate Host header in graph_server request handlers

[Samin061]

GraphServerHandler.do_GET/do_POST served any request that reached the loopback port, so a page in the notebook user's browser can reach the server through DNS rebinding and drive its /post_query handler, which routes into convert_graph_params and builds a BigQuery client with the user's credentials to call list_rows against a caller-supplied api endpoint and table. The widget is legitimately cross-origin because the notebook page and this server run on different ports, so Origin cannot gate the request, but real widget traffic always carries a loopback Host. The handlers now check the Host header and refuse anything that is not localhost, 127.0.0.1, or ::1.

• Make sure to open an issue as a bug/issue before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea
• Ensure the tests and linter pass
• Code coverage does not decrease (if any source code was changed)
• Appropriate docs were updated (if necessary)

[gemini-code-assist]

Code Review

This pull request introduces host header validation to ensure that the graph server only accepts requests targeting the loopback interface, mitigating potential DNS-rebinding attacks. This is implemented via a new _host_is_loopback helper method and applied to both GET and POST requests, with corresponding unit tests added. Feedback suggests wrapping the urllib.parse.urlsplit call in a try...except ValueError block to gracefully handle malformed Host headers and prevent unhandled exceptions.

[parthea]

Hi @Samin061,

Thanks for opening a PR! We'll need to add this to each new test

@pytest.mark.skipif(
    graph_visualization is None, reason="Requires `spanner-graph-notebook`"
)

google-cloud-python/packages/bigquery-magics/bigquery_magics/bigquery.py

Line 695 in abf1178

"Use of --graph requires the spanner-graph-notebook package to be installed. Install it with `pip install 'bigquery-magics[spanner-graph-notebook]'`."

[parthea]

I'm going to switch this to draft but please mark it ready for review once tests are green

[Samin061]

Added the skipif decorator to each of the three new host-header tests, matching the rest of the file. Also guarded the urlsplit call against a ValueError on a malformed Host. Marking ready for review.

[parthea]

Hi @Samin061, Please could you resolve the minor coverage gap?

bigquery_magics/graph_server.py                   202      3     52      1    98%   372, 375-376

[Samin061]

Added two direct tests for _host_is_loopback, one with no Host header and one with a malformed [::1 value, which cover lines 372 and 375-376. Should close the gap.

[parthea]

Gemini raised a concern that we should also inspect origin and referer too.

Checking the Host header only prevents DNS rebinding, but it does not stop a malicious website from making direct browser requests (like fetch('http://localhost')) to your server. To block those, you must also validate the Origin and Referer headers to verify that the request actually initiated from a trusted local page and not an external site.

Could we add these 2 tests?

    @pytest.mark.skipif(
        graph_visualization is None, reason="Requires `spanner-graph-notebook`"
    )
    def test_non_loopback_origin_rejected(self):
        # Host header is valid (loopback), but Origin is a malicious third-party site.
        # Without Origin checks, the host-only check is bypassed and CORS wildcard headers are returned.
        response = requests.get(
            self._route(),
            headers={
                "Host": "localhost",
                "Origin": "http://attacker.example.com"
            }
        )
        # Without the fix, this evaluates to 200 OK. We want it to be rejected.
        self.assertEqual(response.status_code, 403)

    @pytest.mark.skipif(
        graph_visualization is None, reason="Requires `spanner-graph-notebook`"
    )
    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"
            }
        )
        # Without the fix, this evaluates to 200 OK. We want it to be rejected.
        self.assertEqual(response.status_code, 403)

(py3142) partheniou@partheniou-vm-3:~/google-cloud-python-20260814$ pytest packages/bigquery-magics/tests/unit/test_graph_server.py::TestGraphServerHostHeader::test_non_loopback_origin_rejected
========================================================================================== test session starts ==========================================================================================
platform linux -- Python 3.14.2, pytest-9.1.1, pluggy-1.6.0
rootdir: /usr/local/google/home/partheniou/google-cloud-python-20260814/packages/bigquery-magics
plugins: cov-7.1.0, snapshot-0.9.0, anyio-4.14.2, asyncio-1.3.0, pyfakefs-6.2.0
asyncio: mode=Mode.STRICT, debug=False, asyncio_default_fixture_loop_scope=None, asyncio_default_test_loop_scope=function
collected 1 item                                                                                                                                                                                        

packages/bigquery-magics/tests/unit/test_graph_server.py F                                                                                                                                        [100%]

=============================================================================================== FAILURES ================================================================================================
______________________________________________________________________ TestGraphServerHostHeader.test_non_loopback_origin_rejected ______________________________________________________________________

self = <unit.test_graph_server.TestGraphServerHostHeader testMethod=test_non_loopback_origin_rejected>

    @pytest.mark.skipif(
        graph_visualization is None, reason="Requires `spanner-graph-notebook`"
    )
    def test_non_loopback_origin_rejected(self):
        # Host header is valid (loopback), but Origin is a malicious third-party site.
        # Without Origin checks, the host-only check is bypassed and CORS wildcard headers are returned.
        response = requests.get(
            self._route(),
            headers={
                "Host": "localhost",
                "Origin": "http://attacker.example.com"
            }
        )
        # Without the fix, this evaluates to 200 OK. We want it to be rejected.
>       self.assertEqual(response.status_code, 403)
E       AssertionError: 200 != 403

packages/bigquery-magics/tests/unit/test_graph_server.py:821: AssertionError
----------------------------------------------------------------------------------------- Captured stdout call ------------------------------------------------------------------------------------------
BigQuery-magics graph server shutting down...
======================================================================================== short test summary info ========================================================================================
FAILED packages/bigquery-magics/tests/unit/test_graph_server.py::TestGraphServerHostHeader::test_non_loopback_origin_rejected - AssertionError: 200 != 403
=========================================================================================== 1 failed in 5.62s ===========================================================================================
(py3142) partheniou@partheniou-vm-3:~/google-cloud-python-20260814$ pytest packages/bigquery-magics/tests/unit/test_graph_server.py::TestGraphServerHostHeader::test_non_loopback_referer_rejected -v
========================================================================================== test session starts ==========================================================================================
platform linux -- Python 3.14.2, pytest-9.1.1, pluggy-1.6.0 -- /usr/local/google/home/partheniou/.pyenv/versions/3.14.2/envs/py3142/bin/python3.14
cachedir: .pytest_cache
rootdir: /usr/local/google/home/partheniou/google-cloud-python-20260814/packages/bigquery-magics
plugins: cov-7.1.0, snapshot-0.9.0, anyio-4.14.2, asyncio-1.3.0, pyfakefs-6.2.0
asyncio: mode=Mode.STRICT, debug=False, asyncio_default_fixture_loop_scope=None, asyncio_default_test_loop_scope=function
collected 1 item                                                                                                                                                                                        

packages/bigquery-magics/tests/unit/test_graph_server.py::TestGraphServerHostHeader::test_non_loopback_referer_rejected FAILED                                                                    [100%]

=============================================================================================== FAILURES ================================================================================================
_____________________________________________________________________ TestGraphServerHostHeader.test_non_loopback_referer_rejected ______________________________________________________________________

self = <unit.test_graph_server.TestGraphServerHostHeader testMethod=test_non_loopback_referer_rejected>

    @pytest.mark.skipif(
        graph_visualization is None, reason="Requires `spanner-graph-notebook`"
    )
    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"
            }
        )
        # Without the fix, this evaluates to 200 OK. We want it to be rejected.
>       self.assertEqual(response.status_code, 403)
E       AssertionError: 200 != 403

packages/bigquery-magics/tests/unit/test_graph_server.py:836: AssertionError
----------------------------------------------------------------------------------------- Captured stdout call ------------------------------------------------------------------------------------------
BigQuery-magics graph server shutting down...
======================================================================================== short test summary info ========================================================================================
FAILED packages/bigquery-magics/tests/unit/test_graph_server.py::TestGraphServerHostHeader::test_non_loopback_referer_rejected - AssertionError: 200 != 403
=========================================================================================== 1 failed in 5.76s ===========================================================================================
(py3142) partheniou@partheniou-vm-3:~/google-cloud-python-20260814$ 

Gemini suggested this fix

Suggested implementation for _host_is_loopback (renamed to reflect both checks):

    def _is_request_safe(self):
        """Return True if the request headers target and originate from loopback.

        Verifies the Host header to prevent DNS rebinding, and validates the
        Origin and Referer headers (if present) to prevent direct cross-site requests.
        """
        # 1. Validate Host header (DNS Rebinding protection)
        host = self.headers.get("Host")
        if not host:
            return False
        try:
            # strip() prevents whitespace injection tricks
            hostname = urllib.parse.urlsplit(f"//{host.strip()}").hostname
        except ValueError:
            return False
        if hostname not in _LOOPBACK_HOSTS:
            return False

        # 2. Validate Origin/Referer headers if present (Cross-Origin protection)
        for header_name in ("Origin", "Referer"):
            value = self.headers.get(header_name)
            if value:
                try:
                    origin_hostname = urllib.parse.urlsplit(value.strip()).hostname
                except ValueError:
                    return False
                # If the header has a parsed hostname, it must be loopback
                if origin_hostname and origin_hostname not in _LOOPBACK_HOSTS:
                    return False
                    
        return True

[Samin061]

Done. Folded the Origin/Referer check into the helper (now _is_request_safe): Host still guards against DNS rebinding, and if Origin or Referer is present its hostname has to be loopback too, so a page on a non-loopback site can't drive the server with a direct fetch. Real widget traffic is on a different port but the same loopback hostname, so it still passes. Added both tests you posted plus the port-mismatch case, and lint/black are clean.

[Samin061]

any update?