diff --git a/docs/web/enterprise/authentication.mdx b/docs/web/enterprise/authentication.mdx index 920275d16..4cc5d0362 100644 --- a/docs/web/enterprise/authentication.mdx +++ b/docs/web/enterprise/authentication.mdx @@ -98,7 +98,7 @@ client acquires tokens: | --- | --- | --- | --- | --- | | Client Credentials | `OAuthFlowType.CLIENT_CREDENTIALS` | `OAuthFlowType.ClientCredentials` | `OAuthFlow::ClientCredentials` | Service-to-service / machine-to-machine. Requires a client ID and client secret registered with your identity provider. | | Authorization Code | `OAuthFlowType.AUTHORIZATION_CODE` | `OAuthFlowType.AuthorizationCode` | `OAuthFlow::AuthorizationCode(AuthorizationCodeOptions::new())` | Interactive user sign-in through a browser and local callback server. Uses S256 PKCE by default. | -| Device Code | `OAuthFlowType.DEVICE_CODE` | `OAuthFlowType.DeviceCode` | `OAuthFlow::DeviceCode` | CLI or headless sign-in. The client displays a verification URL and user code for the user to enter on another device. | +| Device Code | `OAuthFlowType.DEVICE_CODE` | `OAuthFlowType.DeviceCode` | `OAuthFlow::DeviceCode(DeviceCodeOptions::new())` | CLI or headless sign-in. The client displays a verification URL and user code for the user to enter on another device. | | Azure Managed Identity | `OAuthFlowType.AZURE_MANAGED_IDENTITY` | `OAuthFlowType.AzureManagedIdentity` | `OAuthFlow::AzureManagedIdentity { client_id: None }` | Azure workloads with access to the IMDS endpoint. Tokens are fetched from IMDS, so no client secret is stored on the client. | ### Configure OAuth {#configure-oauth} @@ -290,7 +290,7 @@ const db = await lancedb.connect("db://your-database-uri", { ```rust Rust icon="rust" use lancedb::remote::{AuthorizationCodeOptions, OAuthConfig, OAuthFlow, TokenCacheOptions}; -// Use OAuthFlow::DeviceCode for a CLI without a local browser or callback server. +// Use OAuthFlow::DeviceCode(Default::default()) for a CLI without a local browser or callback server. let oauth_config = OAuthConfig { issuer_url: "https://your-issuer.example.com".into(), client_id: "your-public-client-id".into(), @@ -317,9 +317,16 @@ let db = lancedb::connect("db://your-database-uri") ``` -With `DEVICE_CODE` / `DeviceCode` / `OAuthFlow::DeviceCode`, the same +With `DEVICE_CODE` / `DeviceCode` / `OAuthFlow::DeviceCode(DeviceCodeOptions::new())`, the same configuration and connection calls apply; only the `flow` value changes. +Device Code sends no PKCE parameters by default. For a provider that requires +S256 PKCE on device grants, Rust callers can select +`OAuthFlow::DeviceCode(DeviceCodeOptions::new().use_pkce(true))` (import +`DeviceCodeOptions` from `lancedb::remote`). This option is currently Rust-only. +Enable it only when the provider supports the extension; a failed PKCE request +is not retried without PKCE. + ### Reuse an interactive session {#reuse-an-interactive-session} By default, tokens stay in process memory. Set `token_cache` / `tokenCache` @@ -379,7 +386,8 @@ Credentials and rejected for Azure Managed Identity. ### Configuration reference {#configuration-reference} -The same configuration is available in all three SDKs. Python and Rust use +The shared configuration below is available in all three SDKs; device-code PKCE +is a Rust-only option through `DeviceCodeOptions`. Python and Rust use `snake_case` field names; TypeScript uses `camelCase`. In Rust, set optional `OAuthConfig` fields to `Some(value)` or `None`. Configure `redirect_uri`, `callback_port`, and `use_pkce` through the builder methods on diff --git a/nodejs/src/remote.rs b/nodejs/src/remote.rs index df882334b..5c0517641 100644 --- a/nodejs/src/remote.rs +++ b/nodejs/src/remote.rs @@ -276,7 +276,7 @@ impl TryFrom for lancedb::remote::oauth::OAuthConfig { } OAuthFlow::AuthorizationCode(options) } - "device_code" => OAuthFlow::DeviceCode, + "device_code" => OAuthFlow::DeviceCode(Default::default()), "azure_managed_identity" => OAuthFlow::AzureManagedIdentity { client_id: config.managed_identity_client_id, }, @@ -560,7 +560,7 @@ mod tests { let converted = lancedb::remote::oauth::OAuthConfig::try_from(config).unwrap(); assert!(matches!( converted.flow, - lancedb::remote::oauth::OAuthFlow::DeviceCode + lancedb::remote::oauth::OAuthFlow::DeviceCode(options) if !options.use_pkce )); } diff --git a/python/src/oauth.rs b/python/src/oauth.rs index 1639f239a..08be6ac05 100644 --- a/python/src/oauth.rs +++ b/python/src/oauth.rs @@ -67,7 +67,7 @@ impl TryFrom for OAuthConfig { } OAuthFlow::AuthorizationCode(options) } - "device_code" => OAuthFlow::DeviceCode, + "device_code" => OAuthFlow::DeviceCode(Default::default()), "azure_managed_identity" => OAuthFlow::AzureManagedIdentity { client_id: py.managed_identity_client_id, }, @@ -327,7 +327,7 @@ mod tests { fn test_device_code_conversion() { let config = base_config(); let converted = OAuthConfig::try_from(config).unwrap(); - assert!(matches!(converted.flow, OAuthFlow::DeviceCode)); + assert!(matches!(converted.flow, OAuthFlow::DeviceCode(options) if !options.use_pkce)); } #[test] diff --git a/rust/lancedb/cli/src/auth.rs b/rust/lancedb/cli/src/auth.rs index 090bfd1d8..18eab7c20 100644 --- a/rust/lancedb/cli/src/auth.rs +++ b/rust/lancedb/cli/src/auth.rs @@ -221,7 +221,7 @@ impl ManagedOidc { async fn sdk_config(&self) -> Result { let flow = match self.config.flow { OidcFlow::ClientCredentials => OAuthFlow::ClientCredentials, - OidcFlow::DeviceCode => OAuthFlow::DeviceCode, + OidcFlow::DeviceCode => OAuthFlow::DeviceCode(Default::default()), OidcFlow::AuthorizationCode => { let mut options = AuthorizationCodeOptions::new(); if let Some(uri) = &self.config.redirect_uri { diff --git a/rust/lancedb/src/remote.rs b/rust/lancedb/src/remote.rs index 23b320837..dc8431638 100644 --- a/rust/lancedb/src/remote.rs +++ b/rust/lancedb/src/remote.rs @@ -36,7 +36,8 @@ fn extract_job_id(body: &str) -> Option { pub use client::{ClientConfig, HeaderProvider, RetryConfig, TimeoutConfig, TlsConfig}; pub use db::{RemoteDatabaseOptions, RemoteDatabaseOptionsBuilder}; pub use oauth::{ - AuthorizationCodeOptions, ClientAuthMethod, OAuthConfig, OAuthFlow, OAuthHeaderProvider, + AuthorizationCodeOptions, ClientAuthMethod, DeviceCodeOptions, OAuthConfig, OAuthFlow, + OAuthHeaderProvider, }; pub use token_cache::{OAuthSession, SessionLogout, SessionStatus, TokenCacheOptions}; diff --git a/rust/lancedb/src/remote/oauth.rs b/rust/lancedb/src/remote/oauth.rs index 95d74274e..e6f7cc37d 100644 --- a/rust/lancedb/src/remote/oauth.rs +++ b/rust/lancedb/src/remote/oauth.rs @@ -157,6 +157,41 @@ impl AuthorizationCodeOptions { } } +/// Options for the OAuth Device Authorization flow. +/// +/// PKCE is disabled by default because it is a provider extension to the device +/// grant. Enable it only when the provider supports and requires S256 PKCE, +/// such as a Keycloak client with a PKCE requirement. There is no automatic +/// fallback to a device request without PKCE. +/// +/// # Example +/// +/// ``` +/// use lancedb::remote::{DeviceCodeOptions, OAuthFlow}; +/// +/// let flow = OAuthFlow::DeviceCode(DeviceCodeOptions::new().use_pkce(true)); +/// ``` +#[derive(Debug, Clone, Default)] +#[non_exhaustive] +pub struct DeviceCodeOptions { + /// Whether to protect the device authorization exchange with S256 PKCE. + /// Defaults to `false`. + pub use_pkce: bool, +} + +impl DeviceCodeOptions { + /// Create device-code options with PKCE disabled. + pub fn new() -> Self { + Self::default() + } + + /// Enable or disable the provider-specific S256 PKCE extension. + pub fn use_pkce(mut self, use_pkce: bool) -> Self { + self.use_pkce = use_pkce; + self + } +} + /// OAuth authentication flow configuration. #[derive(Debug, Clone)] pub enum OAuthFlow { @@ -170,8 +205,9 @@ pub enum OAuthFlow { AuthorizationCode(AuthorizationCodeOptions), /// Device Authorization grant for CLI and headless environments. The - /// verification URI and user code are written to stderr. - DeviceCode, + /// verification URI and user code are written to stderr. PKCE is disabled + /// by default; see [`DeviceCodeOptions`] for provider-specific options. + DeviceCode(DeviceCodeOptions), /// Azure Managed Identity via IMDS. /// Works on Azure VMs, AKS, App Service, and Azure Functions. @@ -763,7 +799,9 @@ impl OidcClient { ) -> Result { Self::validate_issuer_transport(&issuer_url)?; + // Cargo embeds the package version at compile time; no runtime variable is needed. let http_client = Client::builder() + .user_agent(concat!("LanceDB-Rust-Client/", env!("CARGO_PKG_VERSION"))) .timeout(Duration::from_secs(30)) .redirect(reqwest::redirect::Policy::custom(|attempt| { if oauth_url_uses_secure_transport(attempt.url()) { @@ -1244,17 +1282,29 @@ struct OAuthErrorResponse { struct DeviceCodeSource { oidc: OidcClient, + use_pkce: bool, +} + +struct DeviceAuthorization { + response: StandardDeviceAuthorizationResponse, + verifier: Option, } impl std::fmt::Debug for DeviceCodeSource { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("DeviceCodeSource") .field("oidc", &self.oidc) + .field("use_pkce", &self.use_pkce) .finish() } } impl DeviceCodeSource { + fn with_pkce(mut self, use_pkce: bool) -> Self { + self.use_pkce = use_pkce; + self + } + fn new( issuer_url: String, client_id: String, @@ -1274,10 +1324,11 @@ impl DeviceCodeSource { resource, audience, )?, + use_pkce: false, }) } - async fn request_device_authorization(&self) -> Result { + async fn request_device_authorization(&self) -> Result { let endpoint = self .oidc .get_discovery() @@ -1295,7 +1346,14 @@ impl DeviceCodeSource { .oidc .base_client() .set_device_authorization_url(device_url); + // Device-flow PKCE is a provider extension; send it only when selected. + let pkce = self.use_pkce.then(PkceCodeChallenge::new_random_sha256); let mut request = client.exchange_device_code(); + if let Some((challenge, _)) = &pkce { + request = request + .add_extra_param("code_challenge", challenge.as_str()) + .add_extra_param("code_challenge_method", "S256"); + } for scope in &self.oidc.scopes { request = request.add_scope(Scope::new(scope.clone())); } @@ -1313,17 +1371,20 @@ impl DeviceCodeSource { if let Some(uri) = device.verification_uri_complete() { validate_oauth_url(uri.secret(), "verification_uri_complete")?; } - Ok(device) + Ok(DeviceAuthorization { + response: device, + verifier: pkce.map(|(_, verifier)| verifier), + }) } - async fn poll_for_token( - &self, - device: &StandardDeviceAuthorizationResponse, - ) -> Result { + async fn poll_for_token(&self, device: &DeviceAuthorization) -> Result { let (client, _) = self.oidc.token_client().await?; let mut request = client - .exchange_device_access_token(device) + .exchange_device_access_token(&device.response) .set_max_backoff_interval(Duration::from_secs(10)); + if let Some(verifier) = &device.verifier { + request = request.add_extra_param("code_verifier", verifier.secret()); + } for (name, value) in self.oidc.target_params() { request = request.add_extra_param(name, value); } @@ -1345,13 +1406,17 @@ impl TokenSource for DeviceCodeSource { async fn fetch_token(&self) -> Result { let device = self.request_device_authorization().await?; show_oauth_prompt(&device_prompt( - device.verification_uri().as_str(), - device.user_code().secret(), + device.response.verification_uri().as_str(), + device.response.user_code().secret(), )); let (browser_url, name) = device + .response .verification_uri_complete() .map(|uri| (uri.secret().as_str(), "verification_uri_complete")) - .unwrap_or((device.verification_uri().as_str(), "verification_uri")); + .unwrap_or(( + device.response.verification_uri().as_str(), + "verification_uri", + )); launch_browser(validate_oauth_url(browser_url, name)?); self.poll_for_token(&device).await } @@ -1613,15 +1678,18 @@ pub(crate) fn build_token_source(config: &OAuthConfig) -> Result Box::new(DeviceCodeSource::new( - config.issuer_url.clone(), - config.client_id.clone(), - config.client_secret.clone(), - client_auth_method, - config.scopes.clone(), - config.resource.clone(), - config.audience.clone(), - )?), + OAuthFlow::DeviceCode(options) => Box::new( + DeviceCodeSource::new( + config.issuer_url.clone(), + config.client_id.clone(), + config.client_secret.clone(), + client_auth_method, + config.scopes.clone(), + config.resource.clone(), + config.audience.clone(), + )? + .with_pkce(options.use_pkce), + ), OAuthFlow::AzureManagedIdentity { client_id } => Box::new(AzureImdsSource::new( config.scopes.clone(), client_id.clone(), @@ -2080,6 +2148,19 @@ mod tests { assert!(output.contains("ABCD-EFGH")); } + #[test] + fn test_device_code_options_default_and_pkce() { + assert!(!DeviceCodeOptions::default().use_pkce); + assert!(!DeviceCodeOptions::new().use_pkce); + assert!(DeviceCodeOptions::new().use_pkce(true).use_pkce); + assert!( + !DeviceCodeOptions::new() + .use_pkce(true) + .use_pkce(false) + .use_pkce + ); + } + #[test] fn test_authorization_code_options_default_to_pkce() { let options = AuthorizationCodeOptions::new(); @@ -2629,30 +2710,33 @@ mod tests { #[tokio::test] async fn test_device_authorization_polls_until_success() { - let (issuer_url, token_requests, server) = spawn_device_server().await; - let source = DeviceCodeSource::new( - issuer_url, - "client-id".to_string(), - Some("secret".to_string()), - ClientAuthMethod::ClientSecretBasic, - vec!["openid".to_string()], - None, - None, - ) - .unwrap(); + for use_pkce in [false, true] { + let (issuer_url, token_requests, server) = spawn_device_server(use_pkce).await; + let source = DeviceCodeSource::new( + issuer_url, + "client-id".to_string(), + Some("secret".to_string()), + ClientAuthMethod::ClientSecretBasic, + vec!["openid".to_string()], + None, + None, + ) + .unwrap() + .with_pkce(use_pkce); - let device = source.request_device_authorization().await.unwrap(); - let response = source.poll_for_token(&device).await.unwrap(); + let device = source.request_device_authorization().await.unwrap(); + let response = source.poll_for_token(&device).await.unwrap(); - assert_eq!(response.access_token.secret(), "device-token"); - assert_eq!( - response - .refresh_token() - .map(|token| token.secret().as_str()), - Some("device-refresh") - ); - assert_eq!(token_requests.load(Ordering::SeqCst), 3); - server.await.unwrap(); + assert_eq!(response.access_token.secret(), "device-token"); + assert_eq!( + response + .refresh_token() + .map(|token| token.secret().as_str()), + Some("device-refresh") + ); + assert_eq!(token_requests.load(Ordering::SeqCst), 3); + server.await.unwrap(); + } } #[tokio::test] @@ -3318,7 +3402,7 @@ mod tests { (issuer_url, server) } - async fn spawn_device_server() -> (String, Arc, JoinHandle<()>) { + async fn spawn_device_server(use_pkce: bool) -> (String, Arc, JoinHandle<()>) { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); let issuer_url = format!("http://{addr}"); @@ -3326,6 +3410,7 @@ mod tests { let server_token_requests = Arc::clone(&token_requests); let server = tokio::spawn(async move { + let mut challenge = None; for _ in 0..5 { let (mut stream, _) = listener.accept().await.unwrap(); let request = read_http_request(&mut stream).await; @@ -3345,6 +3430,18 @@ mod tests { ); assert!(request.body.contains("scope=openid")); assert!(!request.body.contains("client_secret")); + let form: HashMap = + url::form_urlencoded::parse(request.body.as_bytes()) + .into_owned() + .collect(); + if use_pkce { + assert_eq!(form["code_challenge_method"], "S256"); + assert!(!form.contains_key("code_verifier")); + challenge = Some(form["code_challenge"].clone()); + } else { + assert!(!form.contains_key("code_challenge")); + assert!(!form.contains_key("code_challenge_method")); + } let device = format!( r#"{{"device_code":"device-code","user_code":"ABCD-EFGH","verification_uri":"http://{addr}/verify","verification_uri_complete":"http://{addr}/verify?user_code=ABCD-EFGH","expires_in":60,"interval":1}}"# ); @@ -3359,6 +3456,19 @@ mod tests { )); assert!(request.body.contains("device_code=device-code")); assert!(!request.body.contains("client_secret")); + let form: HashMap = + url::form_urlencoded::parse(request.body.as_bytes()) + .into_owned() + .collect(); + if use_pkce { + let verifier = PkceCodeVerifier::new(form["code_verifier"].clone()); + assert_eq!( + PkceCodeChallenge::from_code_verifier_sha256(&verifier).as_str(), + challenge.as_deref().unwrap() + ); + } else { + assert!(!form.contains_key("code_verifier")); + } let request = server_token_requests.fetch_add(1, Ordering::SeqCst); match request { 0 => { @@ -3449,14 +3559,16 @@ mod tests { (issuer_url, token_requests, server) } - fn test_device_authorization_response( - expires_in: u64, - interval: u64, - ) -> StandardDeviceAuthorizationResponse { - serde_json::from_str(&format!( + fn test_device_authorization_response(expires_in: u64, interval: u64) -> DeviceAuthorization { + let response = serde_json::from_str(&format!( r#"{{"device_code":"device-code","user_code":"ABCD-EFGH","verification_uri":"http://127.0.0.1/verify","expires_in":{expires_in},"interval":{interval}}}"# )) - .unwrap() + .unwrap(); + let (_, verifier) = PkceCodeChallenge::new_random_sha256(); + DeviceAuthorization { + response, + verifier: Some(verifier), + } } async fn spawn_device_error_server(error: &'static str) -> (String, JoinHandle<()>) { @@ -3501,6 +3613,12 @@ mod tests { for _ in 0..3 { let (mut stream, _) = listener.accept().await.unwrap(); let request = read_http_request(&mut stream).await; + assert!( + request + .header("user-agent") + .is_some_and(|value| !value.is_empty()), + "OAuth discovery and token requests must identify the client" + ); if request .line diff --git a/rust/lancedb/src/remote/token_cache.rs b/rust/lancedb/src/remote/token_cache.rs index 6f34ab73c..cb6c6bca0 100644 --- a/rust/lancedb/src/remote/token_cache.rs +++ b/rust/lancedb/src/remote/token_cache.rs @@ -43,7 +43,7 @@ //! client_id: "my-app".to_string(), //! client_secret: None, //! scopes: vec!["openid".to_string()], -//! flow: OAuthFlow::DeviceCode, +//! flow: OAuthFlow::DeviceCode(Default::default()), //! client_auth_method: None, //! refresh_buffer_secs: None, //! resource: Some("https://api.example.com".to_string()), @@ -237,7 +237,7 @@ fn canonicalize_scopes(scopes: &[String]) -> Vec { fn flow_key(flow: &OAuthFlow) -> Option<&'static str> { match flow { OAuthFlow::AuthorizationCode(_) => Some("authorization_code"), - OAuthFlow::DeviceCode => Some("device_code"), + OAuthFlow::DeviceCode(_) => Some("device_code"), OAuthFlow::ClientCredentials | OAuthFlow::AzureManagedIdentity { .. } => None, } } @@ -845,7 +845,7 @@ pub struct SessionLogout { /// client_id: "my-app".to_string(), /// client_secret: None, /// scopes: vec!["openid".to_string()], -/// flow: OAuthFlow::DeviceCode, +/// flow: OAuthFlow::DeviceCode(Default::default()), /// client_auth_method: None, /// refresh_buffer_secs: None, /// resource: None, @@ -1012,7 +1012,7 @@ mod tests { client_secret: None, client_auth_method: None, scopes: vec!["openid".to_string()], - flow: OAuthFlow::DeviceCode, + flow: OAuthFlow::DeviceCode(Default::default()), refresh_buffer_secs: None, resource: None, audience: None, @@ -1546,6 +1546,18 @@ mod tests { assert_eq!(canonical.file_stem, key.file_stem); } + #[test] + fn test_device_pkce_preserves_cache_identity() { + let dir = cache_tempdir(); + let mut config = device_config(dir.path()); + let standard = CacheKey::new(&config).unwrap(); + config.flow = OAuthFlow::DeviceCode(crate::remote::DeviceCodeOptions::new().use_pkce(true)); + assert_eq!( + CacheKey::new(&config).unwrap().file_stem, + standard.file_stem + ); + } + #[test] fn test_cache_key_separates_identity_dimensions() { let base = device_config(Path::new("/tmp/cache"));