lancedb/python/tests/test_oauth.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors

import asyncio
import importlib.util
import json
import os
import subprocess
import sys
import threading
import urllib.parse
from http.server import BaseHTTPRequestHandler, HTTPServer
from pathlib import Path

import pytest


def _load_oauth_module():
    oauth_path = (
        Path(__file__).parents[1] / "python" / "lancedb" / "remote" / "oauth.py"
    )
    spec = importlib.util.spec_from_file_location("lancedb_remote_oauth", oauth_path)
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def test_oauth_config_repr_redacts_client_secret():
    oauth = _load_oauth_module()

    config = oauth.OAuthConfig(
        issuer_url="https://issuer.example.com",
        client_id="client-id",
        scopes=["scope"],
        client_secret="super-secret",
    )

    rendered = repr(config)
    assert "super-secret" not in rendered
    assert "client_secret" not in rendered


def test_authorization_code_uses_pkce_by_default():
    oauth = _load_oauth_module()

    config = oauth.OAuthConfig(
        issuer_url="https://issuer.example.com",
        client_id="client-id",
        scopes=["openid"],
        flow=oauth.OAuthFlowType.AUTHORIZATION_CODE,
    )

    assert config.use_pkce is True
    assert config.redirect_uri is None
    assert config.callback_port is None


def test_device_code_flow_value():
    oauth = _load_oauth_module()

    assert oauth.OAuthFlowType.DEVICE_CODE.value == "device_code"


def test_client_auth_method_values():
    oauth = _load_oauth_module()

    assert oauth.ClientAuthMethod.NONE.value == "none"
    assert oauth.ClientAuthMethod.CLIENT_SECRET_BASIC.value == "client_secret_basic"
    assert oauth.ClientAuthMethod.CLIENT_SECRET_POST.value == "client_secret_post"


def test_client_auth_method_defaults_to_none():
    oauth = _load_oauth_module()

    config = oauth.OAuthConfig(
        issuer_url="https://issuer.example.com",
        client_id="client-id",
        scopes=["openid"],
        client_auth_method=oauth.ClientAuthMethod.CLIENT_SECRET_POST,
    )

    assert config.client_auth_method is oauth.ClientAuthMethod.CLIENT_SECRET_POST

    default_config = oauth.OAuthConfig(
        issuer_url="https://issuer.example.com",
        client_id="client-id",
        scopes=["openid"],
    )
    assert default_config.client_auth_method is None


def test_token_cache_options_default_to_memory_only():
    oauth = _load_oauth_module()

    config = oauth.OAuthConfig(
        issuer_url="https://issuer.example.com",
        client_id="client-id",
        scopes=["openid"],
    )
    assert config.token_cache is None
    assert config.resource is None
    assert config.audience is None

    options = oauth.TokenCacheOptions()
    assert options.cache_dir is None
    assert options.lock_timeout_secs is None


def _remote_oauth():
    pytest.importorskip("lancedb")
    from lancedb.remote import oauth as remote_oauth

    return remote_oauth


def _device_config(remote_oauth, issuer_url, cache_dir):
    return remote_oauth.OAuthConfig(
        issuer_url=issuer_url,
        client_id="client-id",
        scopes=["openid"],
        flow=remote_oauth.OAuthFlowType.DEVICE_CODE,
        token_cache=remote_oauth.TokenCacheOptions(cache_dir=str(cache_dir)),
    )


def test_oauth_session_status_and_logout_without_cache_entry(tmp_path):
    remote_oauth = _remote_oauth()
    config = _device_config(remote_oauth, "https://issuer.example.com", tmp_path)

    session = remote_oauth.OAuthSession(config)
    status = asyncio.run(session.status())
    assert status.refreshable is False
    assert status.issuer_url == "https://issuer.example.com"
    assert status.client_id == "client-id"
    assert status.scopes == ["openid"]
    assert status.flow == "device_code"
    assert status.obtained_at is None

    logout = asyncio.run(session.logout())
    assert logout.removed is False


class _MockIdpState:
    def __init__(self, port):
        self.port = port
        self.lock = threading.Lock()
        self.device_authorizations = 0
        self.refresh_grants = 0
        self.invalid_grant_rejections = 0
        self.access_tokens_issued = 0
        self.current_refresh = None
        self.requests = []


class _MockIdpHandler(BaseHTTPRequestHandler):
    @property
    def state(self) -> _MockIdpState:
        return self.server.state

    def log_message(self, fmt, *args):
        pass

    def _respond(self, status, payload):
        body = json.dumps(payload).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self):
        if self.path.startswith("/.well-known/openid-configuration"):
            base = f"http://127.0.0.1:{self.state.port}"
            self._respond(
                200,
                {
                    "token_endpoint": f"{base}/token",
                    "device_authorization_endpoint": f"{base}/device",
                },
            )
        else:
            self._respond(404, {})

    def do_POST(self):
        length = int(self.headers.get("Content-Length", 0))
        body = self.rfile.read(length).decode()
        params = urllib.parse.parse_qs(body)
        with self.state.lock:
            self.state.requests.append(params)

        if self.path == "/device":
            with self.state.lock:
                self.state.device_authorizations += 1
            base = f"http://127.0.0.1:{self.state.port}"
            self._respond(
                200,
                {
                    "device_code": "device-code",
                    "user_code": "ABCD-EFGH",
                    "verification_uri": f"{base}/verify",
                    "expires_in": 60,
                    "interval": 1,
                },
            )
            return

        if self.path == "/token":
            grant_type = params.get("grant_type", [""])[0]
            with self.state.lock:
                if grant_type == "refresh_token":
                    self.state.refresh_grants += 1
                    offered = params.get("refresh_token", [""])[0]
                    if offered != self.state.current_refresh:
                        self.state.invalid_grant_rejections += 1
                        self._respond(400, {"error": "invalid_grant"})
                        return
                elif "device_code" not in grant_type:
                    self._respond(400, {"error": "unsupported_grant_type"})
                    return
                self.state.access_tokens_issued += 1
                number = self.state.access_tokens_issued
                refresh = f"refresh-{number}"
                self.state.current_refresh = refresh
            self._respond(
                200,
                {
                    "access_token": f"access-{number}",
                    "refresh_token": refresh,
                    "expires_in": 3600,
                },
            )
            return

        self._respond(404, {})


def _start_mock_idp() -> tuple[_MockIdpState, HTTPServer]:
    server = HTTPServer(("127.0.0.1", 0), _MockIdpHandler)
    state = _MockIdpState(server.server_address[1])
    server.state = state
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()
    return state, server


def _run_subprocess(script: Path, issuer_url: str, cache_dir: Path, target: dict):
    env = dict(os.environ)
    env["LANCEDB_OAUTH_BROWSER"] = "/usr/bin/true"
    result = subprocess.run(
        [sys.executable, str(script), issuer_url, str(cache_dir), json.dumps(target)],
        capture_output=True,
        text=True,
        timeout=120,
        env=env,
    )
    assert result.returncode == 0, (
        f"subprocess failed:\nstdout: {result.stdout}\nstderr: {result.stderr}"
    )
    return result


LOGIN_SCRIPT = """
import asyncio
import json
import sys

from lancedb.remote import OAuthConfig, OAuthFlowType, OAuthSession, TokenCacheOptions

issuer_url, cache_dir = sys.argv[1], sys.argv[2]
config = OAuthConfig(
    issuer_url=issuer_url,
    client_id="client-id",
    scopes=["openid"],
    flow=OAuthFlowType.DEVICE_CODE,
    token_cache=TokenCacheOptions(cache_dir=cache_dir),
    **json.loads(sys.argv[3]),
)
session = OAuthSession(config)
status = asyncio.run(session.login())
assert status.refreshable, "login must cache a refresh token"
assert status.resource == config.resource
assert status.audience == config.audience
print("LOGIN-OK")
"""

REUSE_SCRIPT = """
import asyncio
import json
import sys

import lancedb
from lancedb.remote import OAuthConfig, OAuthFlowType, OAuthSession, TokenCacheOptions

issuer_url, cache_dir = sys.argv[1], sys.argv[2]
config = OAuthConfig(
    issuer_url=issuer_url,
    client_id="client-id",
    scopes=["openid"],
    flow=OAuthFlowType.DEVICE_CODE,
    token_cache=TokenCacheOptions(cache_dir=cache_dir),
    **json.loads(sys.argv[3]),
)

session = OAuthSession(config)
status = asyncio.run(session.status())
assert status.refreshable, "second process must see the cached session"


async def main():
    # Point the database endpoint at a dead port. OAuth headers are fetched
    # before the request is sent, so a successful refresh proves the second
    # process reused the cached session; only the database call fails.
    db = await lancedb.connect_async(
        "db://e2e",
        host_override="http://127.0.0.1:1",
        client_config={"retry_config": {"retries": 0}},
        oauth_config=config,
    )
    try:
        await db.table_names()
    except Exception:
        print("DATABASE-UNREACHABLE-AS-EXPECTED")
    else:
        raise AssertionError("expected the database request to fail")


asyncio.run(main())
print("REUSE-OK")
"""


@pytest.mark.parametrize(
    "target",
    [
        {},
        {
            "resource": "https://api.example.com/a?x=1&y=two",
            "audience": "audience + & / ü",
        },
    ],
)
def test_cross_process_session_reuse_without_new_prompt(tmp_path, target):
    pytest.importorskip("lancedb")
    state, server = _start_mock_idp()
    try:
        issuer_url = f"http://127.0.0.1:{state.port}"
        login_script = tmp_path / "login.py"
        login_script.write_text(LOGIN_SCRIPT)
        reuse_script = tmp_path / "reuse.py"
        reuse_script.write_text(REUSE_SCRIPT)
        cache_dir = tmp_path / "oauth-cache"

        result = _run_subprocess(login_script, issuer_url, cache_dir, target)
        assert "LOGIN-OK" in result.stdout
        assert state.device_authorizations == 1

        result = _run_subprocess(reuse_script, issuer_url, cache_dir, target)
        assert "REUSE-OK" in result.stdout
        assert "DATABASE-UNREACHABLE-AS-EXPECTED" in result.stdout

        # The second process refreshed exactly once and never started a new
        # interactive device flow.
        assert state.refresh_grants == 1
        assert state.device_authorizations == 1
        assert state.invalid_grant_rejections == 0

        assert len(state.requests) == 3
        for params in state.requests:
            for key in ("resource", "audience"):
                assert params.get(key) == ([target[key]] if key in target else None)
        config = _device_config(_remote_oauth(), issuer_url, cache_dir)
        config.resource = target.get("resource")
        config.audience = target.get("audience")
        logout = asyncio.run(_remote_oauth().OAuthSession(config).logout())
        assert logout.removed is True
    finally:
        server.shutdown()
        server.server_close()