diff --git a/CHANGELOG.md b/CHANGELOG.md index af51aa4..9b10d31 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://book.async.rs/overview - `Client::begin_transaction`, `Client::commit` and `Client::rollback` for driving Trino transactions, plus `Client::transaction_id` / `Client::set_transaction_id` to inspect and set the session's transaction at runtime (previously only settable at build time via `ClientBuilder::transaction_id`) - `Error::Transaction` — returned when a transaction operation is attempted in a state that does not allow it (starting one while another is active, or committing/rolling back without one) - `TransactionId::is_active` +- Interactive OAuth2 authentication (`Auth::new_oauth2` / `new_oauth2_with_handler`). On a `401` Bearer challenge the client presents the login URL (browser + stderr by default, or a custom `RedirectHandler`), polls the Trino token endpoint, and retries with the bearer token. Token is cached in-memory for the process. Coordinators with several authentication types enabled (e.g. `http-server.authentication.type=PASSWORD,OAUTH2`) send one `WWW-Authenticate` header per type, so every challenge header is scanned for the Bearer one rather than only the first ### Fixed - **Transactions were unusable.** Trino returns a new transaction's identifier in `X-Trino-Started-Transaction-Id`, but the client parsed that header with a function that recognised only four fixed literals. A real identifier matched none of them and was silently discarded, so `START TRANSACTION` succeeded on the coordinator while every subsequent statement sent `X-Trino-Transaction-Id: NONE` and ran outside the transaction — and `COMMIT`/`ROLLBACK` could not address it. The identifier is now retained and sent on every subsequent request @@ -18,6 +19,7 @@ and this project adheres to [Semantic Versioning](https://book.async.rs/overview ### Changed - **Breaking:** `TransactionId` now models what the `X-Trino-Transaction-Id` header actually carries: `NoTransaction | Id(String)`. The `StartTransaction`, `RollBack` and `Commit` variants are removed — they are SQL statements, not header values, and sending them produced a header Trino does not accept. `to_str` is replaced by `as_header_value(&self) -> &str` and `from_str` by the infallible `from_header_value(&str) -> Self`. `TransactionId` is no longer `Copy` (it now owns a `String`); it is still `Clone`, and now also `PartialEq` and `Eq`. See the [migration guide](MIGRATION.md) +- **Breaking:** `Auth` is now `#[non_exhaustive]` and has a new `OAuth2` variant. Exhaustive `match` on `Auth` must add a wildcard arm ## [0.11.0] - 2026-07-19 diff --git a/CLAUDE.md b/CLAUDE.md index 793eb6b..28aac2a 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -72,3 +72,19 @@ transient failure, query submission (POST) only when definitely not processed. - Tests for `client.rs` internals go in its bottom `mod tests`; integration tests live in `tests/` (wiremock for HTTP, fixtures in `tests/data/models/`). - Releases follow the process in the `release` skill (`.claude/skills/release/`). + +## Manual OAuth2 e2e + +`tests/oauth2.rs::oauth2_real_login` exercises `Auth::new_oauth2()` against a +real, OAuth2-configured Trino coordinator (interactive browser login — not run +in CI). A local Trino + Keycloak stack is committed at +`integration_tests/test_setup/oauth/` (see its README for the one-time +`/etc/hosts` step and setup gotchas): + +```bash +docker compose -f integration_tests/test_setup/oauth/docker-compose.yml up -d +TRINO_OAUTH2_HOST=localhost TRINO_OAUTH2_PORT=8443 TRINO_OAUTH2_NO_VERIFY=1 \ + cargo test --test oauth2 -- --ignored oauth2_real_login +``` + +Or point `TRINO_OAUTH2_HOST` (and `TRINO_OAUTH2_PORT`) at your own Trino + IdP. diff --git a/Cargo.toml b/Cargo.toml index bfd295b..876295e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,6 +12,7 @@ http = {workspace = true} iterable = {workspace = true} lazy_static = {workspace = true} lz4 = {workspace = true, optional = true} +open = {workspace = true} paste = {workspace = true} regex = {workspace = true} # network dependencies @@ -86,6 +87,7 @@ http = "1.4.2" iterable = "0.6" lazy_static = "1.5" lz4 = "1.28" +open = "5" paste = "1.0.15" regex = "1.13.1" reqwest = {version = "0.13.4", default-features = false, features = ["rustls", "json"]} diff --git a/MIGRATION.md b/MIGRATION.md index 9b686af..86ca889 100644 --- a/MIGRATION.md +++ b/MIGRATION.md @@ -81,6 +81,48 @@ match id { It owns a `String`. It is still `Clone`, and now also `PartialEq` and `Eq`. Add `.clone()` where you relied on implicit copies. +### `Auth` is now `#[non_exhaustive]` (OAuth2 support) + +`Auth` gained a new `OAuth2` variant for interactive browser-based +authentication, alongside `Basic` and `Jwt`. To let future variants be added +without another breaking release, `Auth` is now `#[non_exhaustive]` — an +exhaustive `match` no longer compiles and needs a wildcard arm. + +**Before:** + +```rust +match auth { + Auth::Basic(u, p) => ..., + Auth::Jwt(t) => ..., +} +``` + +**After:** + +```rust +match auth { + Auth::Basic(u, p) => ..., + Auth::Jwt(t) => ..., + _ => ..., // required: Auth is now #[non_exhaustive] +} +``` + +To use OAuth2: + +```rust +let client = ClientBuilder::new("user", "coordinator.example.com") + .secure(true) + .auth(Auth::new_oauth2()) + .build()?; +``` + +On a `401` Bearer challenge the client presents the login URL (opens a +browser and prints it to stderr by default; supply a custom `RedirectHandler` +via `Auth::new_oauth2_with_handler` to change that), polls the Trino token +endpoint, and retries the request with the bearer token once the user +completes the login. The token is cached in-memory for the life of the +`Client`. + ## 0.10.x → 0.11.0 ### Error handling (restructured `Error` enum) diff --git a/README.md b/README.md index 23d8176..c7b4fdb 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,7 @@ Fork rationale : ### authn: - Basic Auth - Jwt Auth +- Interactive OAuth2 (browser-based) ### protocols: - Spooling Protocol (for efficient large result set handling) @@ -111,6 +112,39 @@ async fn main() { } ``` +### Interactive OAuth2 example + +Trino's OAuth2 authentication makes the **coordinator** the OAuth client: on a +`401` the client opens the coordinator-supplied login URL in a browser (and +prints it to stderr as a fallback), polls Trino's token endpoint until you finish +the IdP login, then retries with the bearer token. The token is cached in memory +for the life of the `Client`. Requires TLS to the coordinator. + +```rust +use trino_rust_client::auth::Auth; +use trino_rust_client::{ClientBuilder, Row}; + +#[tokio::main] +async fn main() { + let cli = ClientBuilder::new("user", "coordinator.example.com") + .secure(true) + .auth(Auth::new_oauth2()) + .catalog("catalog") + .build() + .unwrap(); + + let data = cli.get_all::("select 1").await.unwrap().into_vec(); + + for r in data { + println!("{:?}", r) + } +} +``` + +Supply a custom presentation strategy (instead of opening a browser) with +`Auth::new_oauth2_with_handler(Arc::new(my_handler))`, and tune the token poll +loop with `.with_poll(max_attempts, timeout)`. + ### Example dealing with fields not known at compile time ```rust use trino_rust_client::{ClientBuilder, Row, Trino}; diff --git a/examples/oauth2.rs b/examples/oauth2.rs new file mode 100644 index 0000000..cf0a51d --- /dev/null +++ b/examples/oauth2.rs @@ -0,0 +1,53 @@ +use std::env::var; + +use dotenvy::dotenv; +use trino_rust_client::auth::Auth; +use trino_rust_client::{ClientBuilder, Row}; + +/// Interactive (browser-based) OAuth2 against a Trino coordinator configured with +/// `http-server.authentication.type=OAUTH2`. +/// +/// On the first request the client receives a `401`, opens the coordinator's +/// login URL in your browser (and prints it to stderr as a fallback for +/// headless/SSH sessions), then polls Trino's token endpoint until you finish +/// the IdP login and retries the request with the bearer token. The token is +/// cached in memory for the life of the `Client`. Trino requires TLS for OAuth2, +/// so `.secure(true)` is mandatory. +/// +/// Run with e.g. a `.env` file providing USERNAME/HOST/PORT/CATALOG/SQL: +/// cargo run --example oauth2 +#[tokio::main] +async fn main() { + dotenv().ok(); + + let user = var("USERNAME").unwrap(); + let host = var("HOST").unwrap(); + let port = var("PORT") + .unwrap_or_else(|_| "8443".into()) + .parse() + .unwrap(); + let catalog = var("CATALOG").unwrap(); + let sql = var("SQL").unwrap(); + + // Default handler: opens the system browser and prints the URL to stderr. + // For a custom presentation strategy use + // `Auth::new_oauth2_with_handler(Arc::new(my_handler))`, and tune the token + // poll loop with `.with_poll(max_attempts, timeout)`. + let auth = Auth::new_oauth2(); + + let cli = ClientBuilder::new(user, host) + .port(port) + .catalog(catalog) + .auth(auth) + .secure(true) // OAuth2 requires HTTPS to the coordinator + // For a self-signed coordinator certificate, also supply its root cert: + // .ssl(Ssl { root_cert: Some(Ssl::read_pem(&"/path/root.pem").unwrap()) }) + .build() + .unwrap(); + + let data = cli.get_all::(sql).await.unwrap().into_vec(); + + for r in data { + println!("{:?}", r) + } +} diff --git a/integration_tests/test_setup/oauth/README.md b/integration_tests/test_setup/oauth/README.md new file mode 100644 index 0000000..c336d7e --- /dev/null +++ b/integration_tests/test_setup/oauth/README.md @@ -0,0 +1,81 @@ +# Manual OAuth2 test stack (Trino + Keycloak) + +A **local, manual** stack for exercising the client's interactive OAuth2 support +against a real Trino coordinator and a real IdP (Keycloak). It backs +`tests/oauth2.rs::oauth2_real_login`. + +> **Not run in CI.** The interactive flow requires a human to complete the +> Keycloak login in a browser — there is no automated login here. The stack has +> been brought up and verified up to that human step: the coordinator starts +> healthy and returns the expected `401` + `WWW-Authenticate: Bearer +> x_redirect_server=..., x_token_server=...` challenge over TLS. Completing the +> browser login (the two gotchas below) is the part you drive yourself. + +## What's in it + +- **Keycloak** (`realm=trino`, confidential client `trino`/`trino-secret`, user + `alice`/`alice`) on `http://keycloak:8080`. +- **Trino 478** coordinator with TLS on `8443` and + `http-server.authentication.type=PASSWORD,OAUTH2`, plus a `memory` catalog. +- A one-shot job that generates a self-signed keystore for the coordinator. + +### Why two authentication types + +A coordinator with several authentication types emits **one `WWW-Authenticate` +header per type**, in configuration order. Verified against Trino 478: + +```console +$ curl -sk -i -X POST https://localhost:8443/v1/statement -H 'X-Trino-User: alice' --data 'SELECT 1' +HTTP/2 401 +www-authenticate: Basic realm="Trino" +www-authenticate: Bearer x_redirect_server="https://localhost:8443/oauth2/token/initiate/...", x_token_server="..." +``` + +`Basic` comes first, so a client that reads only the first header never sees the +Bearer challenge and fails with a bare `401`. `PASSWORD` is listed first here on +purpose to keep the manual e2e run on that hostile ordering. The file-based +password authenticator (`password-authenticator.properties`, `password.db` — +`alice` / `alice`, bcrypt) exists only to make the second type valid; the test +still authenticates via OAuth2. + +## Two gotchas (read before running) + +1. **Trino mandates TLS for OAuth2.** The coordinator serves the client over + `https://localhost:8443`; the client's `auth_http_insecure` can't help here + because it's *Trino* rejecting plain http, not the client. The stack uses a + self-signed cert, so run the test with `TRINO_OAUTH2_NO_VERIFY=1` (or import + the generated cert via `ClientBuilder::ssl`). +2. **Keycloak issuer/hostname must be consistent.** The token `issuer` and + `jwks-url` (used server-to-server by Trino) and the `auth-url` (opened in your + host browser) must all resolve to the *same* Keycloak origin, or issuer + validation fails. The stack pins everything to `http://keycloak:8080`, so add + a hosts entry so your browser can reach it too: + + ```bash + echo "127.0.0.1 keycloak" | sudo tee -a /etc/hosts # one-time + ``` + +## Run + +```bash +docker compose -f integration_tests/test_setup/oauth/docker-compose.yml up -d + +# Wait for Trino to report healthy, then: +TRINO_OAUTH2_HOST=localhost TRINO_OAUTH2_PORT=8443 TRINO_OAUTH2_NO_VERIFY=1 \ + cargo test --test oauth2 -- --ignored oauth2_real_login +``` + +A browser opens for the Keycloak login — sign in as `alice` / `alice`. The test +then completes the poll → bearer → query round-trip and asserts one row. + +The test's Trino session user defaults to `alice` to match the authenticated +principal — Trino denies a query whose session user differs from the OAuth2 +principal (`Access Denied: User alice cannot impersonate user ...`) unless +impersonation is explicitly configured. Override with `TRINO_OAUTH2_USER` for a +coordinator whose principal differs. + +Tear down with: + +```bash +docker compose -f integration_tests/test_setup/oauth/docker-compose.yml down -v +``` diff --git a/integration_tests/test_setup/oauth/docker-compose.yml b/integration_tests/test_setup/oauth/docker-compose.yml new file mode 100644 index 0000000..75a3b97 --- /dev/null +++ b/integration_tests/test_setup/oauth/docker-compose.yml @@ -0,0 +1,83 @@ +# Manual, LOCAL-ONLY stack: Trino (TLS + OAuth2) behind Keycloak, with a ready +# test user. NOT run in CI — the interactive OAuth2 flow needs a human to +# complete the browser login. Pairs with tests/oauth2.rs::oauth2_real_login. +# +# See README.md in this directory for setup, the two common gotchas, and the run +# command. This is a starting point, not turnkey — expect to iterate on the +# Keycloak redirect URIs / issuer hostname on first run. +services: + keycloak: + image: quay.io/keycloak/keycloak:26.0 + command: ["start-dev", "--import-realm", "--http-port=8080"] + environment: + KC_BOOTSTRAP_ADMIN_USERNAME: admin + KC_BOOTSTRAP_ADMIN_PASSWORD: admin + # Pin the issuer/frontend hostname so the token issuer matches what Trino + # validates AND what your host browser reaches. Requires `keycloak` to + # resolve to 127.0.0.1 on the host (see README) so http://keycloak:8080 + # works from both the Trino container and your browser. + KC_HOSTNAME: http://keycloak:8080 + KC_HOSTNAME_BACKCHANNEL_DYNAMIC: "false" + ports: + - "8080:8080" + volumes: + - ./keycloak/trino-realm.json:/opt/keycloak/data/import/trino-realm.json:ro + networks: + - trino-oauth + healthcheck: + # Force bash — the /dev/tcp check is a bash builtin, not POSIX sh. + test: ["CMD", "bash", "-c", "exec 3<>/dev/tcp/localhost/8080"] + interval: 10s + timeout: 5s + retries: 20 + + # Generates a self-signed PKCS12 keystore for the coordinator once. + gen-certs: + image: eclipse-temurin:21-jdk + # Exec-form entrypoint (list) so the shell script is passed as a single + # argument — a folded string here gets mangled by Compose's word-splitting. + entrypoint: + - /bin/sh + - -ec + - | + if [ ! -f /certs/keystore.p12 ]; then + keytool -genkeypair -alias trino -keyalg RSA -keysize 2048 -validity 3650 \ + -storetype PKCS12 -keystore /certs/keystore.p12 -storepass changeit \ + -dname 'CN=localhost' -ext 'SAN=DNS:localhost,DNS:coordinator,IP:127.0.0.1' + fi + volumes: + - certs:/certs + networks: + - trino-oauth + + coordinator: + image: trinodb/trino:478 + depends_on: + keycloak: + condition: service_healthy + gen-certs: + condition: service_completed_successfully + ports: + - "8443:8443" + volumes: + - ./trino/etc:/etc/trino + - certs:/etc/trino/certs:ro + environment: + - JAVA_OPTS=-Xmx1G -XX:+UseG1GC + networks: + - trino-oauth + healthcheck: + # The https endpoint requires auth; the internal http port stays open for + # a simple liveness probe. + test: ["CMD", "curl", "-f", "http://localhost:8080/v1/info"] + interval: 10s + timeout: 5s + retries: 10 + start_period: 30s + +networks: + trino-oauth: + driver: bridge + +volumes: + certs: diff --git a/integration_tests/test_setup/oauth/keycloak/trino-realm.json b/integration_tests/test_setup/oauth/keycloak/trino-realm.json new file mode 100644 index 0000000..a6136c8 --- /dev/null +++ b/integration_tests/test_setup/oauth/keycloak/trino-realm.json @@ -0,0 +1,34 @@ +{ + "realm": "trino", + "enabled": true, + "sslRequired": "none", + "clients": [ + { + "clientId": "trino", + "enabled": true, + "protocol": "openid-connect", + "publicClient": false, + "secret": "trino-secret", + "standardFlowEnabled": true, + "directAccessGrantsEnabled": true, + "redirectUris": [ + "https://localhost:8443/oauth2/callback", + "https://localhost:8443/ui/*" + ], + "webOrigins": ["https://localhost:8443"] + } + ], + "users": [ + { + "username": "alice", + "enabled": true, + "emailVerified": true, + "firstName": "Alice", + "lastName": "Example", + "email": "alice@example.com", + "credentials": [ + { "type": "password", "value": "alice", "temporary": false } + ] + } + ] +} diff --git a/integration_tests/test_setup/oauth/trino/etc/catalog/memory.properties b/integration_tests/test_setup/oauth/trino/etc/catalog/memory.properties new file mode 100644 index 0000000..833abd3 --- /dev/null +++ b/integration_tests/test_setup/oauth/trino/etc/catalog/memory.properties @@ -0,0 +1 @@ +connector.name=memory diff --git a/integration_tests/test_setup/oauth/trino/etc/config.properties b/integration_tests/test_setup/oauth/trino/etc/config.properties new file mode 100644 index 0000000..45d4862 --- /dev/null +++ b/integration_tests/test_setup/oauth/trino/etc/config.properties @@ -0,0 +1,35 @@ +coordinator=true +node-scheduler.include-coordinator=true + +# Required whenever any authentication type is enabled (signs internal requests). +# Any fixed value works for this single-node local stack. +internal-communication.shared-secret=trino-oauth-local-shared-secret + +# Internal discovery/liveness stays on plain http (8080, container-internal); +# the client-facing endpoint is https (8443) where OAuth2 is enforced. This +# avoids self-signed internal-TLS trust issues for a single-node coordinator. +http-server.http.enabled=true +http-server.http.port=8080 +http-server.https.enabled=true +http-server.https.port=8443 +http-server.https.keystore.path=/etc/trino/certs/keystore.p12 +http-server.https.keystore.key=changeit +discovery.uri=http://coordinator:8080 + +# Interactive OAuth2 (Authorization Code flow, coordinator-mediated). +# +# PASSWORD is listed FIRST on purpose: a coordinator with several authentication +# types sends one `WWW-Authenticate` header per type, in this order, so the +# `Basic realm="Trino"` challenge arrives BEFORE the Bearer one. The client must +# scan every header rather than only the first (see `Client::send`). Keeping the +# hostile order here means the manual e2e run exercises that path. +http-server.authentication.type=PASSWORD,OAUTH2 +web-ui.authentication.type=oauth2 +http-server.authentication.oauth2.issuer=http://keycloak:8080/realms/trino +http-server.authentication.oauth2.auth-url=http://keycloak:8080/realms/trino/protocol/openid-connect/auth +http-server.authentication.oauth2.token-url=http://keycloak:8080/realms/trino/protocol/openid-connect/token +http-server.authentication.oauth2.jwks-url=http://keycloak:8080/realms/trino/protocol/openid-connect/certs +http-server.authentication.oauth2.client-id=trino +http-server.authentication.oauth2.client-secret=trino-secret +http-server.authentication.oauth2.principal-field=preferred_username +http-server.authentication.oauth2.scopes=openid diff --git a/integration_tests/test_setup/oauth/trino/etc/jvm.config b/integration_tests/test_setup/oauth/trino/etc/jvm.config new file mode 100644 index 0000000..ba266dd --- /dev/null +++ b/integration_tests/test_setup/oauth/trino/etc/jvm.config @@ -0,0 +1,4 @@ +-server +-Xmx1G +-XX:+UseG1GC +-XX:+ExitOnOutOfMemoryError diff --git a/integration_tests/test_setup/oauth/trino/etc/node.properties b/integration_tests/test_setup/oauth/trino/etc/node.properties new file mode 100644 index 0000000..e70197c --- /dev/null +++ b/integration_tests/test_setup/oauth/trino/etc/node.properties @@ -0,0 +1,3 @@ +node.environment=oauthtest +node.id=ffffffff-ffff-ffff-ffff-fffffffffffe +node.data-dir=/tmp/trino/data diff --git a/integration_tests/test_setup/oauth/trino/etc/password-authenticator.properties b/integration_tests/test_setup/oauth/trino/etc/password-authenticator.properties new file mode 100644 index 0000000..0162f53 --- /dev/null +++ b/integration_tests/test_setup/oauth/trino/etc/password-authenticator.properties @@ -0,0 +1,7 @@ +# File-based PASSWORD authenticator, enabled purely so the coordinator is +# configured with MORE THAN ONE authentication type. Trino then emits one +# `WWW-Authenticate` challenge header per type, in `http-server.authentication.type` +# order — which is exactly the case the client's challenge parsing has to survive +# (see config.properties). +password-authenticator.name=file +file.password-file=/etc/trino/password.db diff --git a/integration_tests/test_setup/oauth/trino/etc/password.db b/integration_tests/test_setup/oauth/trino/etc/password.db new file mode 100644 index 0000000..85606cf --- /dev/null +++ b/integration_tests/test_setup/oauth/trino/etc/password.db @@ -0,0 +1 @@ +alice:$2y$10$74sX248KR8TP3wZqUaj6MuFfh6Vke9aIqiZcXBUEM6a6B6g2p3uDi diff --git a/src/auth.rs b/src/auth.rs deleted file mode 100644 index 365931b..0000000 --- a/src/auth.rs +++ /dev/null @@ -1,31 +0,0 @@ -use std::fmt; - -#[derive(Clone)] -pub enum Auth { - Basic(String, Option), - Jwt(String), -} - -impl Auth { - pub fn new_basic(username: impl ToString, password: Option) -> Auth { - Auth::Basic(username.to_string(), password.map(|p| p.to_string())) - } - - pub fn new_jwt(token: impl ToString) -> Auth { - Auth::Jwt(token.to_string()) - } -} - -impl fmt::Debug for Auth { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Auth::Basic(name, _) => f - .debug_struct("BasicAuth") - .field("username", name) - .field("password", &"******") - .finish(), - - Auth::Jwt(_) => f.debug_struct("JwtAuth").field("token", &"******").finish(), - } - } -} diff --git a/src/auth/mod.rs b/src/auth/mod.rs new file mode 100644 index 0000000..484d649 --- /dev/null +++ b/src/auth/mod.rs @@ -0,0 +1,91 @@ +use std::fmt; +use std::sync::Arc; +use std::time::Duration; + +mod oauth2; +pub(crate) use oauth2::run_flow; +pub use oauth2::{ + parse_www_authenticate, BrowserRedirectHandler, Challenge, OAuth2State, RedirectHandler, +}; + +const DEFAULT_MAX_POLL_ATTEMPTS: usize = 10; +const DEFAULT_POLL_TIMEOUT: Duration = Duration::from_secs(120); + +#[derive(Clone)] +#[non_exhaustive] +pub enum Auth { + Basic(String, Option), + Jwt(String), + OAuth2(Arc), +} + +impl Auth { + pub fn new_basic(username: impl ToString, password: Option) -> Auth { + Auth::Basic(username.to_string(), password.map(|p| p.to_string())) + } + + pub fn new_jwt(token: impl ToString) -> Auth { + Auth::Jwt(token.to_string()) + } + + /// Interactive OAuth2 using the default browser handler. + pub fn new_oauth2() -> Auth { + Auth::new_oauth2_with_handler(Arc::new(BrowserRedirectHandler)) + } + + /// Interactive OAuth2 with a caller-supplied redirect handler. + pub fn new_oauth2_with_handler(handler: Arc) -> Auth { + Auth::OAuth2(Arc::new(OAuth2State::new( + handler, + DEFAULT_MAX_POLL_ATTEMPTS, + DEFAULT_POLL_TIMEOUT, + ))) + } + + /// Override the token-server poll settings. No-op for non-OAuth2 auth. + pub fn with_poll(self, max_attempts: usize, timeout: Duration) -> Auth { + match self { + Auth::OAuth2(state) => Auth::OAuth2(Arc::new(OAuth2State::new( + Arc::clone(&state.handler), + max_attempts, + timeout, + ))), + other => other, + } + } +} + +impl fmt::Debug for Auth { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Auth::Basic(name, _) => f + .debug_struct("BasicAuth") + .field("username", name) + .field("password", &"******") + .finish(), + + Auth::Jwt(_) => f.debug_struct("JwtAuth").field("token", &"******").finish(), + + Auth::OAuth2(_) => f + .debug_struct("OAuth2Auth") + .field("token", &"******") + .finish(), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn oauth2_debug_redacts_token() { + let auth = Auth::new_oauth2(); + if let Auth::OAuth2(state) = &auth { + *state.token.write().unwrap() = Some("super-secret".to_string()); + } + let dbg = format!("{auth:?}"); + assert!(!dbg.contains("super-secret"), "token leaked: {dbg}"); + assert!(dbg.contains("OAuth2Auth")); + } +} diff --git a/src/auth/oauth2.rs b/src/auth/oauth2.rs new file mode 100644 index 0000000..b85b21b --- /dev/null +++ b/src/auth/oauth2.rs @@ -0,0 +1,365 @@ +use std::sync::{Arc, RwLock}; +use std::time::Duration; + +use serde::Deserialize; + +use crate::error::{Error, Result}; + +/// The redirect + token endpoints extracted from a Trino `WWW-Authenticate` +/// Bearer challenge. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Challenge { + pub x_redirect_server: String, + pub x_token_server: String, +} + +/// Parse a Trino OAuth2 `WWW-Authenticate: Bearer ...` challenge. +/// +/// Returns `None` when the header is not a Bearer challenge or lacks +/// `x_token_server` (the one field the flow cannot proceed without). +pub fn parse_www_authenticate(header: &str) -> Option { + let trimmed = header.trim(); + // Must be a Bearer challenge. Use `get` (not indexing) so a non-ASCII or + // malformed header degrades to `None` instead of panicking on a byte slice + // that lands inside a multi-byte char. + match trimmed.get(..6) { + Some(prefix) if prefix.eq_ignore_ascii_case("bearer") => {} + _ => return None, + } + + let mut x_redirect_server = None; + let mut x_token_server = None; + + for part in trimmed.split(',') { + let Some((key, value)) = part.split_once('=') else { + continue; + }; + // The first key may arrive as `bearer x_redirect_server`; take the last + // whitespace-separated token as the real key. + let key = key + .trim() + .rsplit(char::is_whitespace) + .next() + .unwrap_or("") + .trim(); + let value = value.trim().trim_matches('"'); + match key.to_ascii_lowercase().as_str() { + "x_redirect_server" => x_redirect_server = Some(value.to_string()), + "x_token_server" => x_token_server = Some(value.to_string()), + _ => {} + } + } + + Some(Challenge { + // x_redirect_server can legitimately be absent (already-authenticated + // reuse); default to empty so the handler simply has nothing to open. + x_redirect_server: x_redirect_server.unwrap_or_default(), + x_token_server: x_token_server?, + }) +} + +/// Presents the OAuth2 login URL to the user. The client calls this once per +/// authentication; it must return promptly — it only *shows* the URL, it does +/// not wait for the user to finish (the client detects completion by polling). +pub trait RedirectHandler: Send + Sync { + fn redirect(&self, url: &str) -> Result<()>; +} + +/// Default handler: opens the system browser and also prints the URL to stderr +/// so headless / SSH sessions can still complete the flow. +pub struct BrowserRedirectHandler; + +impl RedirectHandler for BrowserRedirectHandler { + fn redirect(&self, url: &str) -> Result<()> { + eprintln!("Open the following URL in a browser to authenticate:\n{url}"); + // Best-effort; failure to launch a browser is not fatal — the URL is + // already on stderr. + let _ = open::that(url); + Ok(()) + } +} + +/// Shared, interior-mutable state behind `Auth::OAuth2`. Cloning the enclosing +/// `Arc` shares one token cache across every clone of the `Client`. +pub struct OAuth2State { + pub(crate) token: RwLock>, + /// Serializes the browser+poll flow so concurrent 401s open one browser. + pub(crate) acquire: tokio::sync::Mutex<()>, + pub(crate) handler: Arc, + pub(crate) max_poll_attempts: usize, + pub(crate) poll_timeout: Duration, +} + +impl OAuth2State { + pub fn new( + handler: Arc, + max_poll_attempts: usize, + poll_timeout: Duration, + ) -> Self { + Self { + token: RwLock::new(None), + acquire: tokio::sync::Mutex::new(()), + handler, + max_poll_attempts, + poll_timeout, + } + } + + /// The currently cached bearer token, if the flow has completed. + pub fn cached_token(&self) -> Option { + self.token.read().unwrap().clone() + } +} + +#[derive(Deserialize)] +struct TokenResponse { + token: Option, + #[serde(rename = "nextUri")] + next_uri: Option, + error: Option, +} + +/// Present the login URL, then poll the token server following `nextUri` until a +/// token is returned, an error is reported, or attempts/timeout are exhausted. +pub(crate) async fn run_flow( + client: &reqwest::Client, + state: &OAuth2State, + challenge: &Challenge, +) -> Result { + if !challenge.x_redirect_server.is_empty() { + state.handler.redirect(&challenge.x_redirect_server)?; + } + + let deadline = tokio::time::Instant::now() + state.poll_timeout; + let mut url = challenge.x_token_server.clone(); + + for _ in 0..state.max_poll_attempts { + let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); + if remaining.is_zero() { + break; + } + let body = match tokio::time::timeout(remaining, async { + let resp = client.get(&url).send().await?; + resp.json::().await + }) + .await + { + Ok(result) => result?, // network/decode error propagates as before + Err(_elapsed) => break, // exceeded poll_timeout on this poll + }; + + if let Some(err) = body.error { + return Err(Error::OAuth2(format!( + "token endpoint returned error: {err}" + ))); + } + if let Some(token) = body.token { + return Ok(token); + } + match body.next_uri { + Some(next) => url = next, + None => { + return Err(Error::OAuth2( + "token endpoint response had neither token nor nextUri".to_string(), + )) + } + } + } + + Err(Error::OAuth2(format!( + "authentication did not complete within {} attempts / {:?}", + state.max_poll_attempts, state.poll_timeout + ))) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_standard_challenge() { + let h = r#"Bearer x_redirect_server="https://c/oauth2/token/initiate/abc", x_token_server="https://c/oauth2/token/abc""#; + let c = parse_www_authenticate(h).expect("should parse"); + assert_eq!(c.x_redirect_server, "https://c/oauth2/token/initiate/abc"); + assert_eq!(c.x_token_server, "https://c/oauth2/token/abc"); + } + + #[test] + fn tolerates_bearer_prefixed_key_quirk() { + // Naive splitting yields the first key as `bearer x_redirect_server`. + let h = r#"Bearer x_redirect_server="https://c/i", x_token_server="https://c/t""#; + let c = parse_www_authenticate(h).expect("should parse"); + assert_eq!(c.x_redirect_server, "https://c/i"); + assert_eq!(c.x_token_server, "https://c/t"); + } + + #[test] + fn ignores_param_order_and_extra_params() { + let h = r#"Bearer realm="trino", x_token_server="https://c/t", x_redirect_server="https://c/i""#; + let c = parse_www_authenticate(h).expect("should parse"); + assert_eq!(c.x_token_server, "https://c/t"); + assert_eq!(c.x_redirect_server, "https://c/i"); + } + + #[test] + fn none_when_no_token_server() { + let h = r#"Bearer x_redirect_server="https://c/i""#; + assert!(parse_www_authenticate(h).is_none()); + } + + #[test] + fn none_when_not_bearer() { + assert!(parse_www_authenticate(r#"Basic realm="trino""#).is_none()); + } + + #[test] + fn none_on_non_ascii_header_without_panicking() { + // A multi-byte char straddling byte index 6 must not panic the byte slice. + assert!(parse_www_authenticate("aaaaaé x_token_server=\"https://c/t\"").is_none()); + } + + use std::sync::Arc; + use std::time::Duration; + + #[test] + fn state_defaults_have_no_token() { + let state = OAuth2State::new( + Arc::new(BrowserRedirectHandler), + 10, + Duration::from_secs(120), + ); + assert!(state.cached_token().is_none()); + } + + #[test] + fn state_stores_and_reads_token() { + let state = OAuth2State::new( + Arc::new(BrowserRedirectHandler), + 10, + Duration::from_secs(120), + ); + *state.token.write().unwrap() = Some("tok".to_string()); + assert_eq!(state.cached_token().as_deref(), Some("tok")); + } + + struct RecordingHandler { + seen: std::sync::Mutex>, + } + impl RedirectHandler for RecordingHandler { + fn redirect(&self, url: &str) -> Result<()> { + self.seen.lock().unwrap().push(url.to_string()); + Ok(()) + } + } + + #[tokio::test] + async fn run_flow_follows_next_uri_then_returns_token() { + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + let server = MockServer::start().await; + // First poll -> keep polling (nextUri to /token/step2). + Mock::given(method("GET")) + .and(path("/token/step1")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "nextUri": format!("{}/token/step2", server.uri()) + }))) + .mount(&server) + .await; + // Second poll -> token ready. + Mock::given(method("GET")) + .and(path("/token/step2")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "token": "final-token" + }))) + .mount(&server) + .await; + + let handler = Arc::new(RecordingHandler { + seen: std::sync::Mutex::new(vec![]), + }); + let state = OAuth2State::new(handler.clone(), 10, Duration::from_secs(30)); + let challenge = Challenge { + x_redirect_server: "https://login.example/redirect".to_string(), + x_token_server: format!("{}/token/step1", server.uri()), + }; + + let token = run_flow(&reqwest::Client::new(), &state, &challenge) + .await + .expect("flow should succeed"); + + assert_eq!(token, "final-token"); + assert_eq!( + handler.seen.lock().unwrap().as_slice(), + &["https://login.example/redirect".to_string()] + ); + } + + #[tokio::test] + async fn run_flow_surfaces_error_field() { + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/token/err")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "error": "access_denied" + }))) + .mount(&server) + .await; + + let state = OAuth2State::new( + Arc::new(BrowserRedirectHandler), + 10, + Duration::from_secs(30), + ); + let challenge = Challenge { + x_redirect_server: String::new(), + x_token_server: format!("{}/token/err", server.uri()), + }; + + let err = run_flow(&reqwest::Client::new(), &state, &challenge) + .await + .unwrap_err(); + match err { + crate::error::Error::OAuth2(msg) => assert!(msg.contains("access_denied")), + other => panic!("expected OAuth2 error, got {other:?}"), + } + } + + #[tokio::test] + async fn run_flow_bounded_by_poll_timeout_when_server_stalls() { + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/token/stall")) + .respond_with( + ResponseTemplate::new(200) + .set_delay(Duration::from_secs(30)) + .set_body_json(serde_json::json!({ "token": "never-arrives-in-time" })), + ) + .mount(&server) + .await; + + let state = OAuth2State::new( + Arc::new(BrowserRedirectHandler), + 10, + Duration::from_millis(150), + ); + let challenge = Challenge { + x_redirect_server: String::new(), + x_token_server: format!("{}/token/stall", server.uri()), + }; + + let err = run_flow(&reqwest::Client::new(), &state, &challenge) + .await + .unwrap_err(); + match err { + crate::error::Error::OAuth2(msg) => assert!(msg.contains("did not complete")), + other => panic!("expected OAuth2 timeout error, got {other:?}"), + } + } +} diff --git a/src/client.rs b/src/client.rs index 195509f..2f9befb 100644 --- a/src/client.rs +++ b/src/client.rs @@ -599,6 +599,15 @@ impl Drop for RowStream<'_, T> { req = match auth { Auth::Basic(u, p) => req.basic_auth(u, p.as_ref()), Auth::Jwt(t) => req.bearer_auth(t), + // Only ever the cached token: a `Drop` must not block on + // an interactive login, so unlike `Client::auth_req` this + // never runs the OAuth2 flow. With no cached token — or an + // expired one — the cancellation is simply lost and the + // coordinator times the query out on its own. + Auth::OAuth2(state) => match state.cached_token() { + Some(t) => req.bearer_auth(t), + None => req, + }, }; } let _ = req.send().await; @@ -1172,7 +1181,6 @@ impl Client { add_session_header(req, &session) }; - let req = self.auth_req(req); self.send(req, StatusCode::OK, |resp| async { let text = resp.text().await?; @@ -1196,7 +1204,6 @@ impl Client { add_prepare_header(req, &session) }; - let req = self.auth_req(req); self.send(req, StatusCode::OK, |resp| async { let text = resp.text().await?; let data: QueryResult = serde_json::from_str(&text) @@ -1216,7 +1223,6 @@ impl Client { add_prepare_header(req, &session) }; - let req = self.auth_req(req); self.send(req, StatusCode::NO_CONTENT, |_| async { Ok(()) }) .await } @@ -1226,6 +1232,13 @@ impl Client { match auth { Auth::Basic(u, p) => req.basic_auth(u, p.as_ref()), Auth::Jwt(t) => req.bearer_auth(t), + // Tokens are acquired lazily: with nothing cached the request + // goes out unauthenticated, and `send` runs the login flow on + // the resulting 401 challenge before retrying once. + Auth::OAuth2(state) => match state.cached_token() { + Some(t) => req.bearer_auth(t), + None => req, + }, } } else { req @@ -1242,7 +1255,59 @@ impl Client { F: FnOnce(Response) -> Fut, Fut: std::future::Future>, { - let resp = req.send().await?; + // Capture the token we are about to authenticate with (if any) so the + // single-flight refresh can tell whether another task already rotated it. + let sent_token = match self.auth.as_ref() { + Some(Auth::OAuth2(state)) => state.cached_token(), + _ => None, + }; + // Clone the UN-authed builder up front so an OAuth2 401 can be retried + // with exactly one fresh token header (`bearer_auth` appends, so auth is + // applied only after the clone is taken). Bodies here are `String`s, so + // `try_clone` always succeeds. + let retry_req = req.try_clone(); + let resp = self.auth_req(req).send().await?; + + if resp.status() == StatusCode::UNAUTHORIZED { + if let (Some(Auth::OAuth2(state)), Some(retry_req)) = (self.auth.as_ref(), retry_req) { + // A coordinator with several authentication types configured + // (e.g. `http-server.authentication.type=PASSWORD,OAUTH2`) sends + // one `WWW-Authenticate` header per type, in configuration + // order — so `Basic realm="Trino"` may well precede the Bearer + // challenge. Scan all of them for the OAuth2 one. + if let Some(challenge) = resp + .headers() + .get_all(reqwest::header::WWW_AUTHENTICATE) + .iter() + .filter_map(|v| v.to_str().ok()) + .find_map(crate::auth::parse_www_authenticate) + { + self.acquire_oauth2_token(state, &challenge, sent_token) + .await?; + let resp = self.auth_req(retry_req).send().await?; + return self + .finish_send(resp, expected_status, handle_response) + .await; + } + } + } + + self.finish_send(resp, expected_status, handle_response) + .await + } + + /// Shared response-status handling (extracted so both the first and the + /// retried OAuth2 request go through the same path). + async fn finish_send( + &self, + resp: Response, + expected_status: StatusCode, + handle_response: F, + ) -> Result + where + F: FnOnce(Response) -> Fut, + Fut: std::future::Future>, + { let status = resp.status(); if status != expected_status { let data = resp.text().await.unwrap_or("".to_string()); @@ -1253,6 +1318,26 @@ impl Client { } } + /// Acquire an OAuth2 token under a single-flight lock: if another task + /// already refreshed while we waited, reuse that token instead of opening a + /// second browser. + async fn acquire_oauth2_token( + &self, + state: &std::sync::Arc, + challenge: &crate::auth::Challenge, + sent_token: Option, + ) -> Result<()> { + let _guard = state.acquire.lock().await; + // Someone else finished the flow (rotating the token this request was + // sent with) while we waited for the lock — reuse it. + if state.cached_token() != sent_token { + return Ok(()); + } + let token = crate::auth::run_flow(&self.client, state, challenge).await?; + *state.token.write().unwrap() = Some(token); + Ok(()) + } + async fn update_session(&self, resp: &Response) { let mut session = self.session.write().await; diff --git a/src/error.rs b/src/error.rs index c53d65c..31e1f23 100644 --- a/src/error.rs +++ b/src/error.rs @@ -42,6 +42,10 @@ pub enum Error { /// or rolling back without one. #[error("transaction error: {0}")] Transaction(String), + /// The interactive OAuth2 authentication flow failed (no token server in the + /// challenge, the token endpoint returned an error, or it timed out). + #[error("oauth2 error: {0}")] + OAuth2(String), #[error("inconsistent data")] InconsistentData, #[error("reach max attempt: {0}")] diff --git a/tests/oauth2.rs b/tests/oauth2.rs new file mode 100644 index 0000000..732b668 --- /dev/null +++ b/tests/oauth2.rs @@ -0,0 +1,350 @@ +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use trino_rust_client::auth::{Auth, RedirectHandler}; +use trino_rust_client::client::ClientBuilder; +use trino_rust_client::error::Result as TrinoResult; +use trino_rust_client::{Client, Row}; + +use wiremock::matchers::{body_string_contains, header, method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +struct RecordingHandler { + seen: Mutex>, +} +impl RedirectHandler for RecordingHandler { + fn redirect(&self, url: &str) -> TrinoResult<()> { + self.seen.lock().unwrap().push(url.to_string()); + Ok(()) + } +} + +/// wiremock 0.6 has no built-in "header absent" combinator. +struct HeaderAbsent(&'static str); +impl wiremock::Match for HeaderAbsent { + fn matches(&self, req: &wiremock::Request) -> bool { + !req.headers.contains_key(self.0) + } +} + +/// Matches iff the request carries EXACTLY ONE `Authorization` header whose +/// value is `Bearer `. Guards against `bearer_auth` appending a second +/// header on the OAuth2 re-auth retry. +struct SingleBearer(&'static str); +impl wiremock::Match for SingleBearer { + fn matches(&self, req: &wiremock::Request) -> bool { + let vals: Vec<_> = req + .headers + .get_all(reqwest::header::AUTHORIZATION) + .into_iter() + .collect(); + vals.len() == 1 && vals[0].to_str().ok() == Some(&format!("Bearer {}", self.0)) + } +} + +fn client_for(server: &MockServer, handler: Arc) -> Client { + let host = server.uri().replace("http://", ""); + let (host, port) = host.split_once(':').unwrap(); + ClientBuilder::new("test-user", host) + .port(port.parse().unwrap()) + .secure(false) + .auth_http_insecure(true) + .auth(Auth::new_oauth2_with_handler(handler)) + .build() + .unwrap() +} + +/// Minimal terminal (FINISHED, no nextUri) statement response, reusing the +/// exact field names from `tests/data/models/query_result_empty`. +fn finished_query_json() -> String { + std::fs::read_to_string("tests/data/models/query_result_empty").unwrap() +} + +#[tokio::test] +async fn oauth2_happy_path_authenticates_and_caches() { + let server = MockServer::start().await; + let challenge = format!( + r#"Bearer x_redirect_server="https://login/redirect", x_token_server="{}/oauth2/token/abc""#, + server.uri() + ); + + // First statement submission with no token -> 401 with challenge. + Mock::given(method("POST")) + .and(path("/v1/statement")) + .and(HeaderAbsent("authorization")) + .respond_with( + ResponseTemplate::new(401).insert_header("WWW-Authenticate", challenge.as_str()), + ) + .up_to_n_times(1) + .mount(&server) + .await; + + // Token endpoint -> token immediately. + Mock::given(method("GET")) + .and(path("/oauth2/token/abc")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "token": "test-token" + }))) + .mount(&server) + .await; + + // Authenticated submission -> a terminal QueryResult fixture. + Mock::given(method("POST")) + .and(path("/v1/statement")) + .and(header("authorization", "Bearer test-token")) + .respond_with(ResponseTemplate::new(200).set_body_string(finished_query_json())) + .mount(&server) + .await; + + let handler = Arc::new(RecordingHandler { + seen: Mutex::new(vec![]), + }); + let client = client_for(&server, handler.clone()); + + client.get_all::("SELECT 1").await.expect("query ok"); + // Browser presented exactly once. + assert_eq!(handler.seen.lock().unwrap().len(), 1); + + // Second query reuses the cached token: no new 401 mock is needed, and the + // handler is not called again. + client.get_all::("SELECT 2").await.expect("query ok"); + assert_eq!(handler.seen.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn oauth2_single_flight_opens_browser_once() { + let server = MockServer::start().await; + // Include x_redirect_server (unlike a bare-minimum challenge) so the + // handler is actually invoked and the "opens once" assertion is meaningful. + let challenge = format!( + r#"Bearer x_redirect_server="https://login/redirect", x_token_server="{}/oauth2/token/abc""#, + server.uri() + ); + Mock::given(method("POST")) + .and(path("/v1/statement")) + .and(HeaderAbsent("authorization")) + .respond_with( + ResponseTemplate::new(401).insert_header("WWW-Authenticate", challenge.as_str()), + ) + .mount(&server) + .await; + // Token endpoint is slow enough that both requests race into the flow. + Mock::given(method("GET")) + .and(path("/oauth2/token/abc")) + .respond_with( + ResponseTemplate::new(200) + .set_delay(Duration::from_millis(200)) + .set_body_json(serde_json::json!({ "token": "test-token" })), + ) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/v1/statement")) + .and(header("authorization", "Bearer test-token")) + .respond_with(ResponseTemplate::new(200).set_body_string(finished_query_json())) + .mount(&server) + .await; + + let handler = Arc::new(RecordingHandler { + seen: Mutex::new(vec![]), + }); + let client = Arc::new(client_for(&server, handler.clone())); + + let c1 = client.clone(); + let c2 = client.clone(); + let (r1, r2) = tokio::join!(c1.get_all::("SELECT 1"), c2.get_all::("SELECT 2"),); + r1.expect("q1 ok"); + r2.expect("q2 ok"); + + // Both queries authenticated, but the browser was presented once. + assert_eq!(handler.seen.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn oauth2_401_without_challenge_is_http_not_ok() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/statement")) + .respond_with(ResponseTemplate::new(401)) // no WWW-Authenticate + .mount(&server) + .await; + + let handler = Arc::new(RecordingHandler { + seen: Mutex::new(vec![]), + }); + let client = client_for(&server, handler.clone()); + + let err = client.get_all::("SELECT 1").await.unwrap_err(); + assert!(matches!( + err, + trino_rust_client::error::Error::HttpNotOk(code, _) if code == reqwest::StatusCode::UNAUTHORIZED + )); + assert_eq!(handler.seen.lock().unwrap().len(), 0); +} + +/// A coordinator with several authentication types configured +/// (`http-server.authentication.type=PASSWORD,OAUTH2`) answers a 401 with one +/// `WWW-Authenticate` header per type, in configuration order — verified against +/// Trino 478, which puts `Basic realm="Trino"` FIRST. Reading only the first +/// header would miss the Bearer challenge and surface a bare `HttpNotOk(401)`. +#[tokio::test] +async fn oauth2_finds_bearer_challenge_after_basic_challenge() { + let server = MockServer::start().await; + let challenge = format!( + r#"Bearer x_redirect_server="https://login/redirect", x_token_server="{}/oauth2/token/abc""#, + server.uri() + ); + + Mock::given(method("POST")) + .and(path("/v1/statement")) + .and(HeaderAbsent("authorization")) + .respond_with( + ResponseTemplate::new(401) + .append_header("WWW-Authenticate", r#"Basic realm="Trino""#) + .append_header("WWW-Authenticate", challenge.as_str()), + ) + .up_to_n_times(1) + .mount(&server) + .await; + + Mock::given(method("GET")) + .and(path("/oauth2/token/abc")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "token": "test-token" + }))) + .mount(&server) + .await; + + Mock::given(method("POST")) + .and(path("/v1/statement")) + .and(header("authorization", "Bearer test-token")) + .respond_with(ResponseTemplate::new(200).set_body_string(finished_query_json())) + .mount(&server) + .await; + + let handler = Arc::new(RecordingHandler { + seen: Mutex::new(vec![]), + }); + let client = client_for(&server, handler.clone()); + + client.get_all::("SELECT 1").await.expect("query ok"); + assert_eq!(handler.seen.lock().unwrap().len(), 1); +} + +/// Simulates a cached OAuth2 token expiring between two queries on the same +/// client. The re-auth retry must carry exactly ONE `Authorization` header — +/// `bearer_auth` APPENDS, so if auth is applied before the retry clone is taken +/// the retried request goes out with two headers and the coordinator rejects it. +#[tokio::test] +async fn oauth2_reauth_on_expiry_sends_single_authorization_header() { + let server = MockServer::start().await; + + let challenge1 = format!( + r#"Bearer x_redirect_server="https://login/redirect", x_token_server="{}/oauth2/token/1""#, + server.uri() + ); + let challenge2 = format!( + r#"Bearer x_redirect_server="https://login/redirect", x_token_server="{}/oauth2/token/2""#, + server.uri() + ); + + // --- Query 1: no cached token -> 401 challenge -> acquire token-v1 -> 200. + Mock::given(method("POST")) + .and(path("/v1/statement")) + .and(body_string_contains("SELECT 1")) + .and(HeaderAbsent("authorization")) + .respond_with( + ResponseTemplate::new(401).insert_header("WWW-Authenticate", challenge1.as_str()), + ) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/oauth2/token/1")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "token": "token-v1" + }))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/v1/statement")) + .and(body_string_contains("SELECT 1")) + .and(header("authorization", "Bearer token-v1")) + .respond_with(ResponseTemplate::new(200).set_body_string(finished_query_json())) + .mount(&server) + .await; + + // --- Query 2: cache attaches (now-expired) token-v1 -> coordinator 401s -> + // acquire token-v2 -> retry must send a SINGLE Bearer token-v2 header. + Mock::given(method("POST")) + .and(path("/v1/statement")) + .and(body_string_contains("SELECT 2")) + .and(header("authorization", "Bearer token-v1")) + .respond_with( + ResponseTemplate::new(401).insert_header("WWW-Authenticate", challenge2.as_str()), + ) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/oauth2/token/2")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "token": "token-v2" + }))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/v1/statement")) + .and(body_string_contains("SELECT 2")) + .and(SingleBearer("token-v2")) + .respond_with(ResponseTemplate::new(200).set_body_string(finished_query_json())) + .mount(&server) + .await; + + let handler = Arc::new(RecordingHandler { + seen: Mutex::new(vec![]), + }); + let client = client_for(&server, handler.clone()); + + client.get_all::("SELECT 1").await.expect("query 1 ok"); + client.get_all::("SELECT 2").await.expect("query 2 ok"); +} + +/// End-to-end against a real Trino coordinator configured for OAuth2. Not run +/// in CI — the interactive flow needs a human to complete the browser login. +/// +/// Run against the bundled local stack (Trino + Keycloak) at +/// `integration_tests/test_setup/oauth/` (see its README for the one-time +/// `/etc/hosts` step and the two setup gotchas): +/// +/// docker compose -f integration_tests/test_setup/oauth/docker-compose.yml up -d +/// TRINO_OAUTH2_HOST=localhost TRINO_OAUTH2_PORT=8443 TRINO_OAUTH2_NO_VERIFY=1 \ +/// cargo test --test oauth2 -- --ignored oauth2_real_login +/// +/// A browser opens for the Keycloak login (user `alice` / `alice`); complete it +/// to let the test pass. Point it at your own coordinator by setting just +/// `TRINO_OAUTH2_HOST` (and `TRINO_OAUTH2_PORT` if not 443). +#[ignore = "requires a real OAuth2-configured Trino coordinator and an interactive browser login"] +#[tokio::test] +async fn oauth2_real_login() { + let host = std::env::var("TRINO_OAUTH2_HOST").expect("set TRINO_OAUTH2_HOST"); + // The Trino session user must match the authenticated OAuth2 principal + // (Keycloak `preferred_username`), otherwise Trino rejects the query as + // impersonation. The bundled stack's user is `alice`. + let user = std::env::var("TRINO_OAUTH2_USER").unwrap_or_else(|_| "alice".to_string()); + let mut builder = ClientBuilder::new(user, host) + .secure(true) + .auth(Auth::new_oauth2()); + if let Ok(port) = std::env::var("TRINO_OAUTH2_PORT") { + builder = builder.port(port.parse().expect("TRINO_OAUTH2_PORT must be a u16")); + } + // The bundled local stack uses a self-signed certificate; set + // TRINO_OAUTH2_NO_VERIFY=1 to skip TLS verification against it. + if std::env::var("TRINO_OAUTH2_NO_VERIFY").is_ok() { + builder = builder.no_verify(true); + } + let client = builder.build().unwrap(); + let ds = client + .get_all::("SELECT 1") + .await + .expect("query ok after login"); + assert_eq!(ds.len(), 1); +}