diff --git a/rust/lancedb/src/remote/oauth.rs b/rust/lancedb/src/remote/oauth.rs index e0c12b5c1..931f033dc 100644 --- a/rust/lancedb/src/remote/oauth.rs +++ b/rust/lancedb/src/remote/oauth.rs @@ -40,6 +40,12 @@ use url::Url; use crate::error::{Error, Result}; use crate::remote::client::HeaderProvider; +mod azure; +use azure::*; + +mod interactive; +use interactive::*; + const DEFAULT_REFRESH_BUFFER_SECS: u64 = 300; const DEFAULT_TOKEN_TTL_SECS: u64 = 3600; const DEFAULT_CALLBACK_PORT: u16 = 8400; @@ -963,616 +969,6 @@ impl TokenSource for ClientCredentialsSource { } } -#[derive(Debug)] -struct ResolvedRedirect { - uri: String, - bind_addr: SocketAddr, - callback_path: String, -} - -impl ResolvedRedirect { - fn new(options: &AuthorizationCodeOptions) -> Result { - let uri = options.redirect_uri.clone().unwrap_or_else(|| { - format!( - "http://127.0.0.1:{}/callback", - options.callback_port.unwrap_or(DEFAULT_CALLBACK_PORT) - ) - }); - let parsed = Url::parse(&uri).map_err(|e| Error::InvalidInput { - message: format!("Invalid OAuth redirect_uri: {e}"), - })?; - - if parsed.scheme() != "http" { - return Err(Error::InvalidInput { - message: "OAuth redirect_uri must use http with a loopback host".to_string(), - }); - } - if parsed.query().is_some() || parsed.fragment().is_some() { - return Err(Error::InvalidInput { - message: "OAuth redirect_uri must not contain a query or fragment".to_string(), - }); - } - - let ip = match parsed.host() { - Some(url::Host::Domain(host)) if host.eq_ignore_ascii_case("localhost") => { - IpAddr::V4(std::net::Ipv4Addr::LOCALHOST) - } - Some(url::Host::Ipv4(ip)) if ip.is_loopback() => IpAddr::V4(ip), - Some(url::Host::Ipv6(ip)) if ip.is_loopback() => IpAddr::V6(ip), - Some(_) => { - return Err(Error::InvalidInput { - message: "OAuth redirect_uri must use a loopback host".to_string(), - }); - } - None => { - return Err(Error::InvalidInput { - message: "OAuth redirect_uri must include a loopback host".to_string(), - }); - } - }; - let port = parsed.port().ok_or(Error::InvalidInput { - message: "OAuth redirect_uri must include a port".to_string(), - })?; - if port == 0 { - return Err(Error::InvalidInput { - message: "OAuth redirect_uri port must be greater than zero".to_string(), - }); - } - if let Some(callback_port) = options.callback_port - && callback_port != port - { - return Err(Error::InvalidInput { - message: format!( - "OAuth callback_port {callback_port} does not match redirect_uri port {port}" - ), - }); - } - - Ok(Self { - uri, - bind_addr: SocketAddr::new(ip, port), - callback_path: parsed.path().to_string(), - }) - } -} - -#[derive(Debug)] -struct AuthorizationRequest { - url: Url, - state: String, - code_verifier: Option, -} - -#[derive(Debug, PartialEq)] -enum AuthorizationCallback { - Code(String), - ProviderError(String), -} - -struct AuthorizationCodeSource { - oidc: OidcClient, - options: AuthorizationCodeOptions, - redirect: ResolvedRedirect, -} - -impl std::fmt::Debug for AuthorizationCodeSource { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("AuthorizationCodeSource") - .field("oidc", &self.oidc) - .field("options", &self.options) - .field("redirect", &self.redirect) - .finish() - } -} - -impl AuthorizationCodeSource { - #[allow(clippy::too_many_arguments)] - fn new( - issuer_url: String, - client_id: String, - client_secret: Option, - client_auth_method: ClientAuthMethod, - scopes: Vec, - resource: Option, - audience: Option, - options: AuthorizationCodeOptions, - ) -> Result { - let redirect = ResolvedRedirect::new(&options)?; - Ok(Self { - oidc: OidcClient::new( - issuer_url, - client_id, - client_secret, - client_auth_method, - scopes, - resource, - audience, - )?, - options, - redirect, - }) - } - - async fn build_authorization_request(&self) -> Result { - let endpoint = self - .oidc - .get_discovery() - .await? - .authorization_endpoint - .ok_or(Error::Runtime { - message: "OIDC discovery did not provide authorization_endpoint".to_string(), - })?; - let auth_url = AuthUrl::new(endpoint).map_err(|e| Error::InvalidInput { - message: format!("Invalid OAuth authorization_endpoint: {e}"), - })?; - let redirect_url = - RedirectUrl::new(self.redirect.uri.clone()).map_err(|e| Error::InvalidInput { - message: format!("Invalid OAuth redirect_uri: {e}"), - })?; - - let pkce = self - .options - .use_pkce - .then(PkceCodeChallenge::new_random_sha256); - - let client = self - .oidc - .base_client() - .set_auth_uri(auth_url) - .set_redirect_uri(redirect_url); - let mut request = client.authorize_url(CsrfToken::new_random); - for scope in &self.oidc.scopes { - request = request.add_scope(Scope::new(scope.clone())); - } - for (name, value) in self.oidc.target_params() { - request = request.add_extra_param(name, value); - } - if let Some((challenge, _)) = pkce.as_ref() { - request = request.set_pkce_challenge(challenge.clone()); - } - let (url, state) = request.url(); - - Ok(AuthorizationRequest { - url, - state: state.secret().clone(), - code_verifier: pkce.map(|(_, verifier)| verifier), - }) - } - - async fn wait_for_callback( - &self, - listener: &TcpListener, - expected_state: &str, - ) -> Result { - let deadline = - TokioInstant::now() + Duration::from_secs(AUTHORIZATION_CALLBACK_TIMEOUT_SECS); - loop { - let (mut stream, _) = tokio::time::timeout_at(deadline, listener.accept()) - .await - .map_err(|_| Error::Runtime { - message: "Timed out waiting for the OAuth authorization callback".to_string(), - })? - .map_err(|e| Error::Runtime { - message: format!("Failed to accept OAuth callback connection: {e}"), - })?; - match read_authorization_callback( - &mut stream, - &self.redirect.callback_path, - expected_state, - deadline, - ) - .await - { - Ok(AuthorizationCallback::Code(code)) => { - write_callback_response(&mut stream, true).await; - return Ok(code); - } - Ok(AuthorizationCallback::ProviderError(message)) => { - write_callback_response(&mut stream, false).await; - return Err(Error::Runtime { message }); - } - Err(error) => { - if TokioInstant::now() >= deadline { - return Err(Error::Runtime { - message: "Timed out waiting for the OAuth authorization callback" - .to_string(), - }); - } - debug!("Ignoring unrelated OAuth callback connection: {error}"); - write_callback_response(&mut stream, false).await; - } - } - } - } - - async fn exchange_code( - &self, - code: &str, - code_verifier: Option, - ) -> Result { - let (client, endpoint) = self.oidc.token_client().await?; - let redirect_url = - RedirectUrl::new(self.redirect.uri.clone()).map_err(|e| Error::InvalidInput { - message: format!("Invalid OAuth redirect_uri: {e}"), - })?; - let mut request = client - .exchange_code(oauth2::AuthorizationCode::new(code.to_string())) - .set_redirect_uri(Cow::Owned(redirect_url)); - for (name, value) in self.oidc.target_params() { - request = request.add_extra_param(name, value); - } - if let Some(verifier) = code_verifier { - request = request.set_pkce_verifier(verifier); - } - - request - .request_async(&self.oidc.http_client) - .await - .map_err(|e| map_token_error(e, &format!("Token request to {endpoint}"))) - } -} - -#[async_trait] -impl TokenSource for AuthorizationCodeSource { - async fn fetch_token(&self) -> Result { - let listener = TcpListener::bind(self.redirect.bind_addr) - .await - .map_err(|e| Error::Runtime { - message: format!( - "Failed to bind OAuth callback server at {}: {e}", - self.redirect.bind_addr - ), - })?; - let request = self.build_authorization_request().await?; - show_oauth_prompt(&authorization_prompt(&request.url)); - launch_browser(request.url.clone()); - let code = self.wait_for_callback(&listener, &request.state).await?; - self.exchange_code(&code, request.code_verifier).await - } - - async fn refresh_token(&self, refresh_token: &str) -> Result { - self.oidc.refresh_token(refresh_token).await - } -} - -/// A minimal OAuth error body, used to sniff retryable `temporarily_unavailable` -/// responses in [`OAuthHttpClient`]. Unknown fields are ignored. -#[derive(Debug, Deserialize)] -struct OAuthErrorResponse { - error: String, -} - -struct DeviceCodeSource { - oidc: OidcClient, -} - -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) - .finish() - } -} - -impl DeviceCodeSource { - fn new( - issuer_url: String, - client_id: String, - client_secret: Option, - client_auth_method: ClientAuthMethod, - scopes: Vec, - resource: Option, - audience: Option, - ) -> Result { - Ok(Self { - oidc: OidcClient::new( - issuer_url, - client_id, - client_secret, - client_auth_method, - scopes, - resource, - audience, - )?, - }) - } - - async fn request_device_authorization(&self) -> Result { - let endpoint = self - .oidc - .get_discovery() - .await? - .device_authorization_endpoint - .ok_or(Error::Runtime { - message: "OIDC discovery did not provide device_authorization_endpoint".to_string(), - })?; - let device_url = - DeviceAuthorizationUrl::new(endpoint.clone()).map_err(|e| Error::InvalidInput { - message: format!("Invalid OAuth device_authorization_endpoint: {e}"), - })?; - - let client = self - .oidc - .base_client() - .set_device_authorization_url(device_url); - let mut request = client.exchange_device_code(); - for scope in &self.oidc.scopes { - request = request.add_scope(Scope::new(scope.clone())); - } - for (name, value) in self.oidc.target_params() { - request = request.add_extra_param(name, value); - } - let device: StandardDeviceAuthorizationResponse = request - .request_async(&self.oidc.http_client) - .await - .map_err(|e| { - map_token_error(e, &format!("Device authorization request to {endpoint}")) - })?; - - validate_oauth_url(device.verification_uri().as_str(), "verification_uri")?; - if let Some(uri) = device.verification_uri_complete() { - validate_oauth_url(uri.secret(), "verification_uri_complete")?; - } - Ok(device) - } - - async fn poll_for_token( - &self, - device: &StandardDeviceAuthorizationResponse, - ) -> Result { - let (client, _) = self.oidc.token_client().await?; - let mut request = client - .exchange_device_access_token(device) - .set_max_backoff_interval(Duration::from_secs(10)); - for (name, value) in self.oidc.target_params() { - request = request.add_extra_param(name, value); - } - request - .request_async( - &self.oidc.http_client, - // RFC 8628: poll slowly; never spin faster than once a second - // even if a misbehaving server reports a zero interval. - |interval: Duration| tokio::time::sleep(interval.max(Duration::from_secs(1))), - None, - ) - .await - .map_err(map_device_token_error) - } -} - -#[async_trait] -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(), - )); - let (browser_url, name) = device - .verification_uri_complete() - .map(|uri| (uri.secret().as_str(), "verification_uri_complete")) - .unwrap_or((device.verification_uri().as_str(), "verification_uri")); - launch_browser(validate_oauth_url(browser_url, name)?); - self.poll_for_token(&device).await - } - - async fn refresh_token(&self, refresh_token: &str) -> Result { - self.oidc.refresh_token(refresh_token).await - } -} - -fn launch_browser(url: Url) { - drop(tokio::task::spawn_blocking(move || { - if let Some(browser) = std::env::var_os("LANCEDB_OAUTH_BROWSER") { - match Command::new(browser).arg(url.as_str()).status() { - Ok(status) if !status.success() => { - warn!("OAuth browser helper exited with status {status}"); - } - Err(error) => warn!("Could not run the OAuth browser helper: {error}"), - Ok(_) => {} - } - } else if let Err(error) = webbrowser::open(url.as_str()) { - warn!("Could not open an OAuth browser automatically: {error}"); - } - })); -} - -async fn read_authorization_callback( - stream: &mut TcpStream, - expected_path: &str, - expected_state: &str, - overall_deadline: TokioInstant, -) -> Result { - const MAX_CALLBACK_REQUEST_BYTES: usize = 16 * 1024; - let deadline = std::cmp::min( - overall_deadline, - TokioInstant::now() + Duration::from_secs(10), - ); - let mut request = Vec::with_capacity(1024); - loop { - let mut buffer = [0; 1024]; - let count = tokio::time::timeout_at(deadline, stream.read(&mut buffer)) - .await - .map_err(|_| Error::Runtime { - message: "Timed out reading the OAuth authorization callback".to_string(), - })? - .map_err(|e| Error::Runtime { - message: format!("Failed to read OAuth authorization callback: {e}"), - })?; - if count == 0 { - return Err(Error::Runtime { - message: "OAuth authorization callback closed before sending a request".to_string(), - }); - } - request.extend_from_slice(&buffer[..count]); - if request.windows(4).any(|window| window == b"\r\n\r\n") { - break; - } - if request.len() >= MAX_CALLBACK_REQUEST_BYTES { - return Err(Error::Runtime { - message: "OAuth authorization callback request was too large".to_string(), - }); - } - } - let request = std::str::from_utf8(&request).map_err(|e| Error::Runtime { - message: format!("OAuth authorization callback was not valid UTF-8: {e}"), - })?; - parse_authorization_callback(request, expected_path, expected_state) -} - -fn parse_authorization_callback( - request: &str, - expected_path: &str, - expected_state: &str, -) -> Result { - let request_target = request - .lines() - .next() - .and_then(|line| { - let mut parts = line.split_whitespace(); - (parts.next() == Some("GET")) - .then(|| parts.next()) - .flatten() - }) - .ok_or(Error::Runtime { - message: "OAuth authorization callback was not a valid HTTP GET request".to_string(), - })?; - let callback = - Url::parse(&format!("http://loopback{request_target}")).map_err(|e| Error::Runtime { - message: format!("OAuth authorization callback URL was invalid: {e}"), - })?; - if callback.path() != expected_path { - return Err(Error::Runtime { - message: format!( - "OAuth authorization callback used unexpected path {}", - callback.path() - ), - }); - } - let params: HashMap<_, _> = callback.query_pairs().into_owned().collect(); - if params.get("state").map(String::as_str) != Some(expected_state) { - return Err(Error::Runtime { - message: "OAuth authorization callback state did not match".to_string(), - }); - } - if let Some(error) = params.get("error") { - let description = params - .get("error_description") - .map(String::as_str) - .unwrap_or(error); - return Ok(AuthorizationCallback::ProviderError(format!( - "OAuth authorization failed: {description}" - ))); - } - params - .get("code") - .cloned() - .map(AuthorizationCallback::Code) - .ok_or(Error::Runtime { - message: "OAuth authorization callback did not contain a code".to_string(), - }) -} - -async fn write_callback_response(stream: &mut TcpStream, success: bool) { - let (status, body) = if success { - ( - "200 OK", - "

Authentication successful

You can close this window.

", - ) - } else { - ( - "400 Bad Request", - "

Authentication failed

Return to the application for details.

", - ) - }; - let response = format!( - "HTTP/1.1 {status}\r\nContent-Type: text/html; charset=utf-8\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", - body.len() - ); - let _ = stream.write_all(response.as_bytes()).await; -} - -struct AzureImdsSource { - client_id: Option, - resource: String, - http_client: Client, -} - -impl std::fmt::Debug for AzureImdsSource { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("AzureImdsSource") - .field("client_id", &self.client_id) - .field("resource", &self.resource) - .finish() - } -} - -impl AzureImdsSource { - fn new(scopes: Vec, client_id: Option) -> Result { - let resource = Self::resource_from_scopes(&scopes)?; - let http_client = Client::builder() - .timeout(Duration::from_secs(30)) - .no_proxy() - .build() - .map_err(|e| Error::Runtime { - message: format!("Failed to create HTTP client for Azure IMDS OAuth: {e}"), - })?; - - Ok(Self { - client_id, - resource, - http_client, - }) - } - - fn resource_from_scopes(scopes: &[String]) -> Result { - let [scope] = scopes else { - return Err(Error::InvalidInput { - message: "AzureManagedIdentity flow requires exactly one OAuth scope or resource" - .to_string(), - }); - }; - - Ok(scope.strip_suffix("/.default").unwrap_or(scope).to_string()) - } -} - -#[async_trait] -impl TokenSource for AzureImdsSource { - async fn fetch_token(&self) -> Result { - let mut url = format!( - "{AZURE_IMDS_ENDPOINT}?api-version={AZURE_IMDS_API_VERSION}&resource={}", - urlencoding::encode(&self.resource), - ); - if let Some(cid) = self.client_id.as_deref() { - url.push_str(&format!("&client_id={}", urlencoding::encode(cid))); - } - - let resp = self - .http_client - .get(&url) - .header("Metadata", "true") - .send() - .await - .map_err(|e| Error::Runtime { - message: format!("Azure IMDS request failed: {e}"), - })?; - - if !resp.status().is_success() { - return Err(Error::Runtime { - message: format!( - "Azure IMDS returned status {}: {}", - resp.status(), - resp.text().await.unwrap_or_default() - ), - }); - } - - resp.json().await.map_err(|e| Error::Runtime { - message: format!("Failed to parse IMDS token response: {e}"), - }) - } -} /// Build the token source for a configuration. /// diff --git a/rust/lancedb/src/remote/oauth/azure.rs b/rust/lancedb/src/remote/oauth/azure.rs new file mode 100644 index 000000000..ad92c6cf4 --- /dev/null +++ b/rust/lancedb/src/remote/oauth/azure.rs @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use super::*; + +pub(super) struct AzureImdsSource { + client_id: Option, + resource: String, + http_client: Client, +} + +impl std::fmt::Debug for AzureImdsSource { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AzureImdsSource") + .field("client_id", &self.client_id) + .field("resource", &self.resource) + .finish() + } +} + +impl AzureImdsSource { + pub(super) fn new(scopes: Vec, client_id: Option) -> Result { + let resource = Self::resource_from_scopes(&scopes)?; + let http_client = Client::builder() + .timeout(Duration::from_secs(30)) + .no_proxy() + .build() + .map_err(|e| Error::Runtime { + message: format!("Failed to create HTTP client for Azure IMDS OAuth: {e}"), + })?; + + Ok(Self { + client_id, + resource, + http_client, + }) + } + + pub(super) fn resource_from_scopes(scopes: &[String]) -> Result { + let [scope] = scopes else { + return Err(Error::InvalidInput { + message: "AzureManagedIdentity flow requires exactly one OAuth scope or resource" + .to_string(), + }); + }; + + Ok(scope.strip_suffix("/.default").unwrap_or(scope).to_string()) + } +} + +#[async_trait] +impl TokenSource for AzureImdsSource { + async fn fetch_token(&self) -> Result { + let mut url = format!( + "{AZURE_IMDS_ENDPOINT}?api-version={AZURE_IMDS_API_VERSION}&resource={}", + urlencoding::encode(&self.resource), + ); + if let Some(cid) = self.client_id.as_deref() { + url.push_str(&format!("&client_id={}", urlencoding::encode(cid))); + } + + let resp = self + .http_client + .get(&url) + .header("Metadata", "true") + .send() + .await + .map_err(|e| Error::Runtime { + message: format!("Azure IMDS request failed: {e}"), + })?; + + if !resp.status().is_success() { + return Err(Error::Runtime { + message: format!( + "Azure IMDS returned status {}: {}", + resp.status(), + resp.text().await.unwrap_or_default() + ), + }); + } + + resp.json().await.map_err(|e| Error::Runtime { + message: format!("Failed to parse IMDS token response: {e}"), + }) + } +} diff --git a/rust/lancedb/src/remote/oauth/interactive.rs b/rust/lancedb/src/remote/oauth/interactive.rs new file mode 100644 index 000000000..c8fe1df50 --- /dev/null +++ b/rust/lancedb/src/remote/oauth/interactive.rs @@ -0,0 +1,533 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use super::*; + +#[derive(Debug)] +pub(super) struct ResolvedRedirect { + pub(super) uri: String, + pub(super) bind_addr: SocketAddr, + pub(super) callback_path: String, +} + +impl ResolvedRedirect { + pub(super) fn new(options: &AuthorizationCodeOptions) -> Result { + let uri = options.redirect_uri.clone().unwrap_or_else(|| { + format!( + "http://127.0.0.1:{}/callback", + options.callback_port.unwrap_or(DEFAULT_CALLBACK_PORT) + ) + }); + let parsed = Url::parse(&uri).map_err(|e| Error::InvalidInput { + message: format!("Invalid OAuth redirect_uri: {e}"), + })?; + + if parsed.scheme() != "http" { + return Err(Error::InvalidInput { + message: "OAuth redirect_uri must use http with a loopback host".to_string(), + }); + } + if parsed.query().is_some() || parsed.fragment().is_some() { + return Err(Error::InvalidInput { + message: "OAuth redirect_uri must not contain a query or fragment".to_string(), + }); + } + + let ip = match parsed.host() { + Some(url::Host::Domain(host)) if host.eq_ignore_ascii_case("localhost") => { + IpAddr::V4(std::net::Ipv4Addr::LOCALHOST) + } + Some(url::Host::Ipv4(ip)) if ip.is_loopback() => IpAddr::V4(ip), + Some(url::Host::Ipv6(ip)) if ip.is_loopback() => IpAddr::V6(ip), + Some(_) => { + return Err(Error::InvalidInput { + message: "OAuth redirect_uri must use a loopback host".to_string(), + }); + } + None => { + return Err(Error::InvalidInput { + message: "OAuth redirect_uri must include a loopback host".to_string(), + }); + } + }; + let port = parsed.port().ok_or(Error::InvalidInput { + message: "OAuth redirect_uri must include a port".to_string(), + })?; + if port == 0 { + return Err(Error::InvalidInput { + message: "OAuth redirect_uri port must be greater than zero".to_string(), + }); + } + if let Some(callback_port) = options.callback_port + && callback_port != port + { + return Err(Error::InvalidInput { + message: format!( + "OAuth callback_port {callback_port} does not match redirect_uri port {port}" + ), + }); + } + + Ok(Self { + uri, + bind_addr: SocketAddr::new(ip, port), + callback_path: parsed.path().to_string(), + }) + } +} + +#[derive(Debug)] +pub(super) struct AuthorizationRequest { + pub(super) url: Url, + state: String, + pub(super) code_verifier: Option, +} + +#[derive(Debug, PartialEq)] +pub(super) enum AuthorizationCallback { + Code(String), + ProviderError(String), +} + +pub(super) struct AuthorizationCodeSource { + pub(super) oidc: OidcClient, + options: AuthorizationCodeOptions, + redirect: ResolvedRedirect, +} + +impl std::fmt::Debug for AuthorizationCodeSource { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AuthorizationCodeSource") + .field("oidc", &self.oidc) + .field("options", &self.options) + .field("redirect", &self.redirect) + .finish() + } +} + +impl AuthorizationCodeSource { + #[allow(clippy::too_many_arguments)] + pub(super) fn new( + issuer_url: String, + client_id: String, + client_secret: Option, + client_auth_method: ClientAuthMethod, + scopes: Vec, + resource: Option, + audience: Option, + options: AuthorizationCodeOptions, + ) -> Result { + let redirect = ResolvedRedirect::new(&options)?; + Ok(Self { + oidc: OidcClient::new( + issuer_url, + client_id, + client_secret, + client_auth_method, + scopes, + resource, + audience, + )?, + options, + redirect, + }) + } + + pub(super) async fn build_authorization_request(&self) -> Result { + let endpoint = self + .oidc + .get_discovery() + .await? + .authorization_endpoint + .ok_or(Error::Runtime { + message: "OIDC discovery did not provide authorization_endpoint".to_string(), + })?; + let auth_url = AuthUrl::new(endpoint).map_err(|e| Error::InvalidInput { + message: format!("Invalid OAuth authorization_endpoint: {e}"), + })?; + let redirect_url = + RedirectUrl::new(self.redirect.uri.clone()).map_err(|e| Error::InvalidInput { + message: format!("Invalid OAuth redirect_uri: {e}"), + })?; + + let pkce = self + .options + .use_pkce + .then(PkceCodeChallenge::new_random_sha256); + + let client = self + .oidc + .base_client() + .set_auth_uri(auth_url) + .set_redirect_uri(redirect_url); + let mut request = client.authorize_url(CsrfToken::new_random); + for scope in &self.oidc.scopes { + request = request.add_scope(Scope::new(scope.clone())); + } + for (name, value) in self.oidc.target_params() { + request = request.add_extra_param(name, value); + } + if let Some((challenge, _)) = pkce.as_ref() { + request = request.set_pkce_challenge(challenge.clone()); + } + let (url, state) = request.url(); + + Ok(AuthorizationRequest { + url, + state: state.secret().clone(), + code_verifier: pkce.map(|(_, verifier)| verifier), + }) + } + + pub(super) async fn wait_for_callback( + &self, + listener: &TcpListener, + expected_state: &str, + ) -> Result { + let deadline = + TokioInstant::now() + Duration::from_secs(AUTHORIZATION_CALLBACK_TIMEOUT_SECS); + loop { + let (mut stream, _) = tokio::time::timeout_at(deadline, listener.accept()) + .await + .map_err(|_| Error::Runtime { + message: "Timed out waiting for the OAuth authorization callback".to_string(), + })? + .map_err(|e| Error::Runtime { + message: format!("Failed to accept OAuth callback connection: {e}"), + })?; + match read_authorization_callback( + &mut stream, + &self.redirect.callback_path, + expected_state, + deadline, + ) + .await + { + Ok(AuthorizationCallback::Code(code)) => { + write_callback_response(&mut stream, true).await; + return Ok(code); + } + Ok(AuthorizationCallback::ProviderError(message)) => { + write_callback_response(&mut stream, false).await; + return Err(Error::Runtime { message }); + } + Err(error) => { + if TokioInstant::now() >= deadline { + return Err(Error::Runtime { + message: "Timed out waiting for the OAuth authorization callback" + .to_string(), + }); + } + debug!("Ignoring unrelated OAuth callback connection: {error}"); + write_callback_response(&mut stream, false).await; + } + } + } + } + + pub(super) async fn exchange_code( + &self, + code: &str, + code_verifier: Option, + ) -> Result { + let (client, endpoint) = self.oidc.token_client().await?; + let redirect_url = + RedirectUrl::new(self.redirect.uri.clone()).map_err(|e| Error::InvalidInput { + message: format!("Invalid OAuth redirect_uri: {e}"), + })?; + let mut request = client + .exchange_code(oauth2::AuthorizationCode::new(code.to_string())) + .set_redirect_uri(Cow::Owned(redirect_url)); + for (name, value) in self.oidc.target_params() { + request = request.add_extra_param(name, value); + } + if let Some(verifier) = code_verifier { + request = request.set_pkce_verifier(verifier); + } + + request + .request_async(&self.oidc.http_client) + .await + .map_err(|e| map_token_error(e, &format!("Token request to {endpoint}"))) + } +} + +#[async_trait] +impl TokenSource for AuthorizationCodeSource { + async fn fetch_token(&self) -> Result { + let listener = TcpListener::bind(self.redirect.bind_addr) + .await + .map_err(|e| Error::Runtime { + message: format!( + "Failed to bind OAuth callback server at {}: {e}", + self.redirect.bind_addr + ), + })?; + let request = self.build_authorization_request().await?; + show_oauth_prompt(&authorization_prompt(&request.url)); + launch_browser(request.url.clone()); + let code = self.wait_for_callback(&listener, &request.state).await?; + self.exchange_code(&code, request.code_verifier).await + } + + async fn refresh_token(&self, refresh_token: &str) -> Result { + self.oidc.refresh_token(refresh_token).await + } +} + +/// A minimal OAuth error body, used to sniff retryable `temporarily_unavailable` +/// responses in [`OAuthHttpClient`]. Unknown fields are ignored. +#[derive(Debug, Deserialize)] +pub(super) struct OAuthErrorResponse { + pub(super) error: String, +} + +pub(super) struct DeviceCodeSource { + oidc: OidcClient, +} + +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) + .finish() + } +} + +impl DeviceCodeSource { + pub(super) fn new( + issuer_url: String, + client_id: String, + client_secret: Option, + client_auth_method: ClientAuthMethod, + scopes: Vec, + resource: Option, + audience: Option, + ) -> Result { + Ok(Self { + oidc: OidcClient::new( + issuer_url, + client_id, + client_secret, + client_auth_method, + scopes, + resource, + audience, + )?, + }) + } + + pub(super) async fn request_device_authorization(&self) -> Result { + let endpoint = self + .oidc + .get_discovery() + .await? + .device_authorization_endpoint + .ok_or(Error::Runtime { + message: "OIDC discovery did not provide device_authorization_endpoint".to_string(), + })?; + let device_url = + DeviceAuthorizationUrl::new(endpoint.clone()).map_err(|e| Error::InvalidInput { + message: format!("Invalid OAuth device_authorization_endpoint: {e}"), + })?; + + let client = self + .oidc + .base_client() + .set_device_authorization_url(device_url); + let mut request = client.exchange_device_code(); + for scope in &self.oidc.scopes { + request = request.add_scope(Scope::new(scope.clone())); + } + for (name, value) in self.oidc.target_params() { + request = request.add_extra_param(name, value); + } + let device: StandardDeviceAuthorizationResponse = request + .request_async(&self.oidc.http_client) + .await + .map_err(|e| { + map_token_error(e, &format!("Device authorization request to {endpoint}")) + })?; + + validate_oauth_url(device.verification_uri().as_str(), "verification_uri")?; + if let Some(uri) = device.verification_uri_complete() { + validate_oauth_url(uri.secret(), "verification_uri_complete")?; + } + Ok(device) + } + + pub(super) async fn poll_for_token( + &self, + device: &StandardDeviceAuthorizationResponse, + ) -> Result { + let (client, _) = self.oidc.token_client().await?; + let mut request = client + .exchange_device_access_token(device) + .set_max_backoff_interval(Duration::from_secs(10)); + for (name, value) in self.oidc.target_params() { + request = request.add_extra_param(name, value); + } + request + .request_async( + &self.oidc.http_client, + // RFC 8628: poll slowly; never spin faster than once a second + // even if a misbehaving server reports a zero interval. + |interval: Duration| tokio::time::sleep(interval.max(Duration::from_secs(1))), + None, + ) + .await + .map_err(map_device_token_error) + } +} + +#[async_trait] +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(), + )); + let (browser_url, name) = device + .verification_uri_complete() + .map(|uri| (uri.secret().as_str(), "verification_uri_complete")) + .unwrap_or((device.verification_uri().as_str(), "verification_uri")); + launch_browser(validate_oauth_url(browser_url, name)?); + self.poll_for_token(&device).await + } + + async fn refresh_token(&self, refresh_token: &str) -> Result { + self.oidc.refresh_token(refresh_token).await + } +} + +pub(super) fn launch_browser(url: Url) { + drop(tokio::task::spawn_blocking(move || { + if let Some(browser) = std::env::var_os("LANCEDB_OAUTH_BROWSER") { + match Command::new(browser).arg(url.as_str()).status() { + Ok(status) if !status.success() => { + warn!("OAuth browser helper exited with status {status}"); + } + Err(error) => warn!("Could not run the OAuth browser helper: {error}"), + Ok(_) => {} + } + } else if let Err(error) = webbrowser::open(url.as_str()) { + warn!("Could not open an OAuth browser automatically: {error}"); + } + })); +} + +pub(super) async fn read_authorization_callback( + stream: &mut TcpStream, + expected_path: &str, + expected_state: &str, + overall_deadline: TokioInstant, +) -> Result { + const MAX_CALLBACK_REQUEST_BYTES: usize = 16 * 1024; + let deadline = std::cmp::min( + overall_deadline, + TokioInstant::now() + Duration::from_secs(10), + ); + let mut request = Vec::with_capacity(1024); + loop { + let mut buffer = [0; 1024]; + let count = tokio::time::timeout_at(deadline, stream.read(&mut buffer)) + .await + .map_err(|_| Error::Runtime { + message: "Timed out reading the OAuth authorization callback".to_string(), + })? + .map_err(|e| Error::Runtime { + message: format!("Failed to read OAuth authorization callback: {e}"), + })?; + if count == 0 { + return Err(Error::Runtime { + message: "OAuth authorization callback closed before sending a request".to_string(), + }); + } + request.extend_from_slice(&buffer[..count]); + if request.windows(4).any(|window| window == b"\r\n\r\n") { + break; + } + if request.len() >= MAX_CALLBACK_REQUEST_BYTES { + return Err(Error::Runtime { + message: "OAuth authorization callback request was too large".to_string(), + }); + } + } + let request = std::str::from_utf8(&request).map_err(|e| Error::Runtime { + message: format!("OAuth authorization callback was not valid UTF-8: {e}"), + })?; + parse_authorization_callback(request, expected_path, expected_state) +} + +pub(super) fn parse_authorization_callback( + request: &str, + expected_path: &str, + expected_state: &str, +) -> Result { + let request_target = request + .lines() + .next() + .and_then(|line| { + let mut parts = line.split_whitespace(); + (parts.next() == Some("GET")) + .then(|| parts.next()) + .flatten() + }) + .ok_or(Error::Runtime { + message: "OAuth authorization callback was not a valid HTTP GET request".to_string(), + })?; + let callback = + Url::parse(&format!("http://loopback{request_target}")).map_err(|e| Error::Runtime { + message: format!("OAuth authorization callback URL was invalid: {e}"), + })?; + if callback.path() != expected_path { + return Err(Error::Runtime { + message: format!( + "OAuth authorization callback used unexpected path {}", + callback.path() + ), + }); + } + let params: HashMap<_, _> = callback.query_pairs().into_owned().collect(); + if params.get("state").map(String::as_str) != Some(expected_state) { + return Err(Error::Runtime { + message: "OAuth authorization callback state did not match".to_string(), + }); + } + if let Some(error) = params.get("error") { + let description = params + .get("error_description") + .map(String::as_str) + .unwrap_or(error); + return Ok(AuthorizationCallback::ProviderError(format!( + "OAuth authorization failed: {description}" + ))); + } + params + .get("code") + .cloned() + .map(AuthorizationCallback::Code) + .ok_or(Error::Runtime { + message: "OAuth authorization callback did not contain a code".to_string(), + }) +} + +pub(super) async fn write_callback_response(stream: &mut TcpStream, success: bool) { + let (status, body) = if success { + ( + "200 OK", + "

Authentication successful

You can close this window.

", + ) + } else { + ( + "400 Bad Request", + "

Authentication failed

Return to the application for details.

", + ) + }; + let response = format!( + "HTTP/1.1 {status}\r\nContent-Type: text/html; charset=utf-8\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ); + let _ = stream.write_all(response.as_bytes()).await; +} diff --git a/rust/lancedb/src/remote/sql.rs b/rust/lancedb/src/remote/sql.rs index e8ddcf6ff..edf9c2aee 100644 --- a/rust/lancedb/src/remote/sql.rs +++ b/rust/lancedb/src/remote/sql.rs @@ -33,6 +33,9 @@ use crate::remote::client::{ClientConfig, TlsConfig}; use crate::remote::retry::ResolvedRetryConfig; use crate::sql::{Query, QueryDescription, QueryHandle, QueryStatus}; +mod transport; +use transport::*; + const DEFAULT_SQL_PORT: u16 = 10025; const DEFAULT_SQL_TLS_PORT: u16 = 10026; const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(120); @@ -1213,297 +1216,6 @@ fn query_expiration(poll_info: &PollInfo) -> Result, - sql_host_override: Option<&str>, -) -> Result { - if let Some(uri) = sql_host_override { - return normalize_sql_host_override(uri); - } - let host_override = host_override.ok_or_else(|| Error::InvalidInput { - message: "sql_host_override is required when the SQL service endpoint cannot be derived from host_override".to_string(), - })?; - let parsed = url::Url::parse(host_override).map_err(|err| Error::InvalidInput { - message: format!("Invalid host_override: {err}"), - })?; - if parsed.scheme() != "http" { - return Err(Error::InvalidInput { - message: "sql_host_override is required for TLS or non-HTTP host overrides".to_string(), - }); - } - validate_endpoint_url(&parsed, "host_override")?; - let port = match parsed.port().or(explicit_port(host_override)) { - Some(u16::MAX) => { - return Err(Error::InvalidInput { - message: "sql_host_override is required when host_override uses port 65535" - .to_string(), - }); - } - Some(port) => port + 1, - None => DEFAULT_SQL_PORT, - }; - Ok(SqlTarget { - uri: endpoint_uri("http", parsed.host_str().unwrap(), port), - tls: false, - }) -} - -fn normalize_sql_host_override(uri: &str) -> Result { - let parsed = url::Url::parse(uri).map_err(|err| Error::InvalidInput { - message: format!("Invalid sql_host_override: {err}"), - })?; - validate_endpoint_url(&parsed, "sql_host_override")?; - let tls = match parsed.scheme().to_ascii_lowercase().as_str() { - "grpc" | "grpc+tcp" | "http" => false, - "grpc+tls" | "grpcs" | "https" => true, - _ => { - return Err(Error::InvalidInput { - message: - "sql_host_override must use grpc, grpc+tcp, grpc+tls, grpcs, http, or https" - .to_string(), - }); - } - }; - let port = parsed.port().or(explicit_port(uri)).unwrap_or(if tls { - DEFAULT_SQL_TLS_PORT - } else { - DEFAULT_SQL_PORT - }); - if port == 0 { - return Err(Error::InvalidInput { - message: "sql_host_override port must be greater than zero".to_string(), - }); - } - Ok(SqlTarget { - uri: endpoint_uri( - if tls { "https" } else { "http" }, - parsed.host_str().unwrap(), - port, - ), - tls, - }) -} - -fn explicit_port(uri: &str) -> Option { - let authority = uri.split_once("://")?.1.split(['/', '?', '#']).next()?; - let suffix = if authority.starts_with('[') { - authority.split_once(']')?.1.strip_prefix(':')? - } else { - authority.rsplit_once(':')?.1 - }; - suffix.parse().ok() -} - -fn validate_endpoint_url(parsed: &url::Url, name: &str) -> Result<()> { - if parsed.host_str().is_none() { - return Err(Error::InvalidInput { - message: format!("{name} must include a hostname"), - }); - } - if !parsed.username().is_empty() || parsed.password().is_some() { - return Err(Error::InvalidInput { - message: format!("{name} must not include user information"), - }); - } - if !matches!(parsed.path(), "" | "/") || parsed.query().is_some() || parsed.fragment().is_some() - { - return Err(Error::InvalidInput { - message: format!("{name} must not include a path, query, or fragment"), - }); - } - Ok(()) -} - -fn endpoint_uri(scheme: &str, host: &str, port: u16) -> String { - if host.contains(':') { - let host = host - .strip_prefix('[') - .and_then(|host| host.strip_suffix(']')) - .unwrap_or(host); - format!("{scheme}://[{host}]:{port}") - } else { - format!("{scheme}://{host}:{port}") - } -} - -async fn connect_channel( - target: &SqlTarget, - config: &ClientConfig, - request_id: &str, -) -> Result { - let connect_timeout = resolve_timeout( - config.timeout_config.connect_timeout, - "LANCE_CLIENT_CONNECT_TIMEOUT", - Some(DEFAULT_CONNECT_TIMEOUT), - )? - .unwrap(); - let mut endpoint = Endpoint::from_shared(target.uri.clone()) - .map_err(|err| sql_error(request_id, err))? - .connect_timeout(connect_timeout); - if target.tls { - endpoint = endpoint - .tls_config(tls_config(config.tls_config.as_ref())?) - .map_err(|err| sql_error(request_id, err))?; - } - tokio::time::timeout(connect_timeout, endpoint.connect()) - .await - .map_err(|_| sql_error(request_id, "SQL connection timed out"))? - .map_err(|err| sql_error(request_id, err)) -} - -fn tls_config(config: Option<&TlsConfig>) -> Result { - let mut tls = ClientTlsConfig::new().with_enabled_roots(); - if let Some(config) = config { - if !config.assert_hostname { - return Err(Error::InvalidInput { - message: "SQL cannot disable TLS hostname verification".to_string(), - }); - } - if let Some(path) = &config.ssl_ca_cert { - let pem = fs::read(path).map_err(|err| Error::InvalidInput { - message: format!("Failed to read SQL CA certificate {path}: {err}"), - })?; - tls = tls.ca_certificate(Certificate::from_pem(pem)); - } - match (&config.cert_file, &config.key_file) { - (Some(cert), Some(key)) => { - let cert_pem = fs::read(cert).map_err(|err| Error::InvalidInput { - message: format!("Failed to read SQL client certificate {cert}: {err}"), - })?; - let key_pem = fs::read(key).map_err(|err| Error::InvalidInput { - message: format!("Failed to read SQL client key {key}: {err}"), - })?; - tls = tls.identity(Identity::from_pem(cert_pem, key_pem)); - } - (None, None) => {} - _ => { - return Err(Error::InvalidInput { - message: "SQL mTLS requires both cert_file and key_file".to_string(), - }); - } - } - } - Ok(tls) -} - -fn client_with_headers( - client: FlightServiceClient, - headers: &HeaderMap, -) -> Result { - let mut client = FlightClient::new_from_inner(client); - for (key, value) in headers { - let value = value.to_str().map_err(|err| Error::InvalidInput { - message: format!("Invalid SQL metadata value for {key:?}: {err}"), - })?; - client - .add_header(key.as_str(), value) - .map_err(|err| Error::InvalidInput { - message: format!("Invalid SQL metadata header {key:?}: {err}"), - })?; - } - Ok(client) -} - -fn merge_headers(destination: &mut HeaderMap, source: &HashMap) -> Result<()> { - for (key, value) in source { - insert_header(destination, key, value)?; - } - Ok(()) -} - -fn insert_header(headers: &mut HeaderMap, key: &str, value: &str) -> Result<()> { - let key = HeaderName::from_bytes(key.as_bytes()).map_err(|err| Error::InvalidInput { - message: format!("Invalid SQL metadata key {key:?}: {err}"), - })?; - let value = HeaderValue::try_from(value).map_err(|err| Error::InvalidInput { - message: format!("Invalid SQL metadata value for {key:?}: {err}"), - })?; - headers.insert(key, value); - Ok(()) -} - -fn validate_namespace_path(path: &[String]) -> Result<()> { - for component in path { - if component.is_empty() - || !component.is_ascii() - || component.contains('$') - || component.bytes().any(|byte| !(0x20..=0x7e).contains(&byte)) - { - return Err(Error::InvalidInput { - message: "default_namespace_path components must be non-empty printable ASCII strings without '$'".to_string(), - }); - } - } - Ok(()) -} - -fn poll_retry_delay(config: &ResolvedRetryConfig, retry_count: u8) -> Duration { - let exponent = i32::from(retry_count.saturating_sub(1).min(16)); - let backoff = config.backoff_factor * 2.0_f32.powi(exponent); - let jitter = rand::random::() * config.backoff_jitter; - Duration::from_secs_f32((backoff + jitter).clamp(MIN_POLL_INTERVAL.as_secs_f32(), 60.0)) -} - -fn cancellation_status_is_ambiguous(code: tonic::Code) -> bool { - matches!( - code, - tonic::Code::Cancelled - | tonic::Code::Unknown - | tonic::Code::DeadlineExceeded - | tonic::Code::Internal - | tonic::Code::Unavailable - | tonic::Code::DataLoss - ) -} - -fn resolve_timeout( - configured: Option, - env_name: &str, - default: Option, -) -> Result> { - if configured.is_some() { - return Ok(configured); - } - match std::env::var(env_name) { - Ok(value) => value - .parse::() - .map(Duration::from_secs) - .map(Some) - .map_err(|_| Error::InvalidInput { - message: format!("Invalid value for {env_name} environment variable: {value:?}"), - }), - Err(_) => Ok(default), - } -} - -async fn with_overall_timeout( - timeout: Option, - operation: &str, - future: impl std::future::Future>, -) -> Result { - match timeout { - Some(timeout) => { - tokio::time::timeout(timeout, future) - .await - .map_err(|_| Error::Runtime { - message: format!("{operation} timed out"), - })? - } - None => future.await, - } -} - -fn sql_error(request_id: &str, error: impl std::fmt::Display) -> Error { - Error::Runtime { - message: format!("SQL error (request_id={request_id}): {error}"), - } -} #[cfg(test)] #[path = "sql_test.rs"] diff --git a/rust/lancedb/src/remote/sql/transport.rs b/rust/lancedb/src/remote/sql/transport.rs new file mode 100644 index 000000000..8ba8e1964 --- /dev/null +++ b/rust/lancedb/src/remote/sql/transport.rs @@ -0,0 +1,296 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use super::*; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(super) struct SqlTarget { + pub(super) uri: String, + pub(super) tls: bool, +} + +pub(super) fn resolve_sql_host_override( + host_override: Option<&str>, + sql_host_override: Option<&str>, +) -> Result { + if let Some(uri) = sql_host_override { + return normalize_sql_host_override(uri); + } + let host_override = host_override.ok_or_else(|| Error::InvalidInput { + message: "sql_host_override is required when the SQL service endpoint cannot be derived from host_override".to_string(), + })?; + let parsed = url::Url::parse(host_override).map_err(|err| Error::InvalidInput { + message: format!("Invalid host_override: {err}"), + })?; + if parsed.scheme() != "http" { + return Err(Error::InvalidInput { + message: "sql_host_override is required for TLS or non-HTTP host overrides".to_string(), + }); + } + validate_endpoint_url(&parsed, "host_override")?; + let port = match parsed.port().or(explicit_port(host_override)) { + Some(u16::MAX) => { + return Err(Error::InvalidInput { + message: "sql_host_override is required when host_override uses port 65535" + .to_string(), + }); + } + Some(port) => port + 1, + None => DEFAULT_SQL_PORT, + }; + Ok(SqlTarget { + uri: endpoint_uri("http", parsed.host_str().unwrap(), port), + tls: false, + }) +} + +pub(super) fn normalize_sql_host_override(uri: &str) -> Result { + let parsed = url::Url::parse(uri).map_err(|err| Error::InvalidInput { + message: format!("Invalid sql_host_override: {err}"), + })?; + validate_endpoint_url(&parsed, "sql_host_override")?; + let tls = match parsed.scheme().to_ascii_lowercase().as_str() { + "grpc" | "grpc+tcp" | "http" => false, + "grpc+tls" | "grpcs" | "https" => true, + _ => { + return Err(Error::InvalidInput { + message: + "sql_host_override must use grpc, grpc+tcp, grpc+tls, grpcs, http, or https" + .to_string(), + }); + } + }; + let port = parsed.port().or(explicit_port(uri)).unwrap_or(if tls { + DEFAULT_SQL_TLS_PORT + } else { + DEFAULT_SQL_PORT + }); + if port == 0 { + return Err(Error::InvalidInput { + message: "sql_host_override port must be greater than zero".to_string(), + }); + } + Ok(SqlTarget { + uri: endpoint_uri( + if tls { "https" } else { "http" }, + parsed.host_str().unwrap(), + port, + ), + tls, + }) +} + +pub(super) fn explicit_port(uri: &str) -> Option { + let authority = uri.split_once("://")?.1.split(['/', '?', '#']).next()?; + let suffix = if authority.starts_with('[') { + authority.split_once(']')?.1.strip_prefix(':')? + } else { + authority.rsplit_once(':')?.1 + }; + suffix.parse().ok() +} + +pub(super) fn validate_endpoint_url(parsed: &url::Url, name: &str) -> Result<()> { + if parsed.host_str().is_none() { + return Err(Error::InvalidInput { + message: format!("{name} must include a hostname"), + }); + } + if !parsed.username().is_empty() || parsed.password().is_some() { + return Err(Error::InvalidInput { + message: format!("{name} must not include user information"), + }); + } + if !matches!(parsed.path(), "" | "/") || parsed.query().is_some() || parsed.fragment().is_some() + { + return Err(Error::InvalidInput { + message: format!("{name} must not include a path, query, or fragment"), + }); + } + Ok(()) +} + +pub(super) fn endpoint_uri(scheme: &str, host: &str, port: u16) -> String { + if host.contains(':') { + let host = host + .strip_prefix('[') + .and_then(|host| host.strip_suffix(']')) + .unwrap_or(host); + format!("{scheme}://[{host}]:{port}") + } else { + format!("{scheme}://{host}:{port}") + } +} + +pub(super) async fn connect_channel( + target: &SqlTarget, + config: &ClientConfig, + request_id: &str, +) -> Result { + let connect_timeout = resolve_timeout( + config.timeout_config.connect_timeout, + "LANCE_CLIENT_CONNECT_TIMEOUT", + Some(DEFAULT_CONNECT_TIMEOUT), + )? + .unwrap(); + let mut endpoint = Endpoint::from_shared(target.uri.clone()) + .map_err(|err| sql_error(request_id, err))? + .connect_timeout(connect_timeout); + if target.tls { + endpoint = endpoint + .tls_config(tls_config(config.tls_config.as_ref())?) + .map_err(|err| sql_error(request_id, err))?; + } + tokio::time::timeout(connect_timeout, endpoint.connect()) + .await + .map_err(|_| sql_error(request_id, "SQL connection timed out"))? + .map_err(|err| sql_error(request_id, err)) +} + +pub(super) fn tls_config(config: Option<&TlsConfig>) -> Result { + let mut tls = ClientTlsConfig::new().with_enabled_roots(); + if let Some(config) = config { + if !config.assert_hostname { + return Err(Error::InvalidInput { + message: "SQL cannot disable TLS hostname verification".to_string(), + }); + } + if let Some(path) = &config.ssl_ca_cert { + let pem = fs::read(path).map_err(|err| Error::InvalidInput { + message: format!("Failed to read SQL CA certificate {path}: {err}"), + })?; + tls = tls.ca_certificate(Certificate::from_pem(pem)); + } + match (&config.cert_file, &config.key_file) { + (Some(cert), Some(key)) => { + let cert_pem = fs::read(cert).map_err(|err| Error::InvalidInput { + message: format!("Failed to read SQL client certificate {cert}: {err}"), + })?; + let key_pem = fs::read(key).map_err(|err| Error::InvalidInput { + message: format!("Failed to read SQL client key {key}: {err}"), + })?; + tls = tls.identity(Identity::from_pem(cert_pem, key_pem)); + } + (None, None) => {} + _ => { + return Err(Error::InvalidInput { + message: "SQL mTLS requires both cert_file and key_file".to_string(), + }); + } + } + } + Ok(tls) +} + +pub(super) fn client_with_headers( + client: FlightServiceClient, + headers: &HeaderMap, +) -> Result { + let mut client = FlightClient::new_from_inner(client); + for (key, value) in headers { + let value = value.to_str().map_err(|err| Error::InvalidInput { + message: format!("Invalid SQL metadata value for {key:?}: {err}"), + })?; + client + .add_header(key.as_str(), value) + .map_err(|err| Error::InvalidInput { + message: format!("Invalid SQL metadata header {key:?}: {err}"), + })?; + } + Ok(client) +} + +pub(super) fn merge_headers(destination: &mut HeaderMap, source: &HashMap) -> Result<()> { + for (key, value) in source { + insert_header(destination, key, value)?; + } + Ok(()) +} + +pub(super) fn insert_header(headers: &mut HeaderMap, key: &str, value: &str) -> Result<()> { + let key = HeaderName::from_bytes(key.as_bytes()).map_err(|err| Error::InvalidInput { + message: format!("Invalid SQL metadata key {key:?}: {err}"), + })?; + let value = HeaderValue::try_from(value).map_err(|err| Error::InvalidInput { + message: format!("Invalid SQL metadata value for {key:?}: {err}"), + })?; + headers.insert(key, value); + Ok(()) +} + +pub(super) fn validate_namespace_path(path: &[String]) -> Result<()> { + for component in path { + if component.is_empty() + || !component.is_ascii() + || component.contains('$') + || component.bytes().any(|byte| !(0x20..=0x7e).contains(&byte)) + { + return Err(Error::InvalidInput { + message: "default_namespace_path components must be non-empty printable ASCII strings without '$'".to_string(), + }); + } + } + Ok(()) +} + +pub(super) fn poll_retry_delay(config: &ResolvedRetryConfig, retry_count: u8) -> Duration { + let exponent = i32::from(retry_count.saturating_sub(1).min(16)); + let backoff = config.backoff_factor * 2.0_f32.powi(exponent); + let jitter = rand::random::() * config.backoff_jitter; + Duration::from_secs_f32((backoff + jitter).clamp(MIN_POLL_INTERVAL.as_secs_f32(), 60.0)) +} + +pub(super) fn cancellation_status_is_ambiguous(code: tonic::Code) -> bool { + matches!( + code, + tonic::Code::Cancelled + | tonic::Code::Unknown + | tonic::Code::DeadlineExceeded + | tonic::Code::Internal + | tonic::Code::Unavailable + | tonic::Code::DataLoss + ) +} + +pub(super) fn resolve_timeout( + configured: Option, + env_name: &str, + default: Option, +) -> Result> { + if configured.is_some() { + return Ok(configured); + } + match std::env::var(env_name) { + Ok(value) => value + .parse::() + .map(Duration::from_secs) + .map(Some) + .map_err(|_| Error::InvalidInput { + message: format!("Invalid value for {env_name} environment variable: {value:?}"), + }), + Err(_) => Ok(default), + } +} + +pub(super) async fn with_overall_timeout( + timeout: Option, + operation: &str, + future: impl std::future::Future>, +) -> Result { + match timeout { + Some(timeout) => { + tokio::time::timeout(timeout, future) + .await + .map_err(|_| Error::Runtime { + message: format!("{operation} timed out"), + })? + } + None => future.await, + } +} + +pub(super) fn sql_error(request_id: &str, error: impl std::fmt::Display) -> Error { + Error::Runtime { + message: format!("SQL error (request_id={request_id}): {error}"), + } +}