mirror of
https://github.com/lancedb/lancedb.git
synced 2026-09-22 21:15:35 +00:00
Colin Patrick McCabe
3a1d3be256
Support configuring resource and audience for OAuth authorization, token exchange, and refresh requests.
353 lines
11 KiB
Python
353 lines
11 KiB
Python
# 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_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()
|