teng-lin--notebooklm-py/tests/unit/mcp/test_host_guard.py

"""Behavioral tests for the MCP loopback Host-header guard (#1869)."""

from __future__ import annotations

import pytest

from notebooklm.mcp._host_guard import LoopbackHostGuardMiddleware


class _Recorder:
    """A trivial ASGI app that records whether it was reached."""

    def __init__(self) -> None:
        self.reached = False

    async def __call__(self, scope, receive, send) -> None:
        self.reached = True
        await send({"type": "http.response.start", "status": 200, "headers": []})
        await send({"type": "http.response.body", "body": b"ok"})


def _http_scope(host: str | None) -> dict:
    headers = [(b"host", host.encode())] if host is not None else []
    return {"type": "http", "headers": headers}


async def _run(mw: LoopbackHostGuardMiddleware, scope: dict) -> int | None:
    sent: list[dict] = []

    async def receive() -> dict:
        return {"type": "http.request", "body": b""}

    async def send(message: dict) -> None:
        sent.append(message)

    await mw(scope, receive, send)
    starts = [m for m in sent if m["type"] == "http.response.start"]
    return starts[0]["status"] if starts else None


@pytest.mark.asyncio
@pytest.mark.parametrize("host", ["evil.example", "attacker.com:9420", "0.0.0.0", "", None])
async def test_non_loopback_host_is_rejected(host: str | None) -> None:
    app = _Recorder()
    status = await _run(LoopbackHostGuardMiddleware(app, allow_external=False), _http_scope(host))
    assert status == 403
    assert app.reached is False  # request never reached the wrapped app


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "host", ["127.0.0.1", "127.0.0.1:9420", "localhost", "[::1]", "[::1]:9420"]
)
async def test_loopback_host_passes_through(host: str) -> None:
    app = _Recorder()
    status = await _run(LoopbackHostGuardMiddleware(app, allow_external=False), _http_scope(host))
    assert status == 200
    assert app.reached is True


@pytest.mark.asyncio
async def test_allow_external_bypasses_the_guard() -> None:
    # When an external bind is opted into (auth is mandatory there), the guard must
    # NOT reject a non-loopback Host — the bearer/OAuth layer owns access control.
    app = _Recorder()
    status = await _run(
        LoopbackHostGuardMiddleware(app, allow_external=True), _http_scope("evil.example")
    )
    assert status == 200
    assert app.reached is True


@pytest.mark.asyncio
async def test_non_http_scope_passes_through() -> None:
    # Lifespan/websocket scopes have no Host to validate and must not be blocked.
    app = _Recorder()
    sent: list[dict] = []

    async def receive() -> dict:
        return {"type": "lifespan.startup"}

    async def send(message: dict) -> None:
        sent.append(message)

    await LoopbackHostGuardMiddleware(app, allow_external=False)(
        {"type": "lifespan"}, receive, send
    )
    assert app.reached is True