mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-12 23:02:28 +00:00
refactor(rust): split remote sql and oauth modules into submodules
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BgkiPYS2eTmAWVyqn6KBnN
This commit is contained in:
Claude
@@ -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<Self> {
|
||||
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<PkceCodeVerifier>,
|
||||
}
|
||||
|
||||
#[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<String>,
|
||||
client_auth_method: ClientAuthMethod,
|
||||
scopes: Vec<String>,
|
||||
resource: Option<String>,
|
||||
audience: Option<String>,
|
||||
options: AuthorizationCodeOptions,
|
||||
) -> Result<Self> {
|
||||
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<AuthorizationRequest> {
|
||||
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<String> {
|
||||
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<PkceCodeVerifier>,
|
||||
) -> Result<TokenResponse> {
|
||||
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<TokenResponse> {
|
||||
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<RefreshResult> {
|
||||
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<String>,
|
||||
client_auth_method: ClientAuthMethod,
|
||||
scopes: Vec<String>,
|
||||
resource: Option<String>,
|
||||
audience: Option<String>,
|
||||
) -> Result<Self> {
|
||||
Ok(Self {
|
||||
oidc: OidcClient::new(
|
||||
issuer_url,
|
||||
client_id,
|
||||
client_secret,
|
||||
client_auth_method,
|
||||
scopes,
|
||||
resource,
|
||||
audience,
|
||||
)?,
|
||||
})
|
||||
}
|
||||
|
||||
async fn request_device_authorization(&self) -> Result<StandardDeviceAuthorizationResponse> {
|
||||
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<TokenResponse> {
|
||||
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<TokenResponse> {
|
||||
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<RefreshResult> {
|
||||
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<AuthorizationCallback> {
|
||||
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<AuthorizationCallback> {
|
||||
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",
|
||||
"<html><body><h2>Authentication successful</h2><p>You can close this window.</p></body></html>",
|
||||
)
|
||||
} else {
|
||||
(
|
||||
"400 Bad Request",
|
||||
"<html><body><h2>Authentication failed</h2><p>Return to the application for details.</p></body></html>",
|
||||
)
|
||||
};
|
||||
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<String>,
|
||||
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<String>, client_id: Option<String>) -> Result<Self> {
|
||||
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<String> {
|
||||
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<TokenResponse> {
|
||||
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.
|
||||
///
|
||||
|
||||
@@ -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<String>,
|
||||
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<String>, client_id: Option<String>) -> Result<Self> {
|
||||
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<String> {
|
||||
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<TokenResponse> {
|
||||
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}"),
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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<Self> {
|
||||
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<PkceCodeVerifier>,
|
||||
}
|
||||
|
||||
#[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<String>,
|
||||
client_auth_method: ClientAuthMethod,
|
||||
scopes: Vec<String>,
|
||||
resource: Option<String>,
|
||||
audience: Option<String>,
|
||||
options: AuthorizationCodeOptions,
|
||||
) -> Result<Self> {
|
||||
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<AuthorizationRequest> {
|
||||
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<String> {
|
||||
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<PkceCodeVerifier>,
|
||||
) -> Result<TokenResponse> {
|
||||
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<TokenResponse> {
|
||||
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<RefreshResult> {
|
||||
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<String>,
|
||||
client_auth_method: ClientAuthMethod,
|
||||
scopes: Vec<String>,
|
||||
resource: Option<String>,
|
||||
audience: Option<String>,
|
||||
) -> Result<Self> {
|
||||
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<StandardDeviceAuthorizationResponse> {
|
||||
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<TokenResponse> {
|
||||
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<TokenResponse> {
|
||||
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<RefreshResult> {
|
||||
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<AuthorizationCallback> {
|
||||
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<AuthorizationCallback> {
|
||||
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",
|
||||
"<html><body><h2>Authentication successful</h2><p>You can close this window.</p></body></html>",
|
||||
)
|
||||
} else {
|
||||
(
|
||||
"400 Bad Request",
|
||||
"<html><body><h2>Authentication failed</h2><p>Return to the application for details.</p></body></html>",
|
||||
)
|
||||
};
|
||||
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;
|
||||
}
|
||||
@@ -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<Option<chrono::DateTime<chro
|
||||
.transpose()
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
struct SqlTarget {
|
||||
uri: String,
|
||||
tls: bool,
|
||||
}
|
||||
|
||||
fn resolve_sql_host_override(
|
||||
host_override: Option<&str>,
|
||||
sql_host_override: Option<&str>,
|
||||
) -> Result<SqlTarget> {
|
||||
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<SqlTarget> {
|
||||
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<u16> {
|
||||
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<Channel> {
|
||||
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<ClientTlsConfig> {
|
||||
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<Channel>,
|
||||
headers: &HeaderMap,
|
||||
) -> Result<FlightClient> {
|
||||
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<String, String>) -> 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::<f32>() * 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<Duration>,
|
||||
env_name: &str,
|
||||
default: Option<Duration>,
|
||||
) -> Result<Option<Duration>> {
|
||||
if configured.is_some() {
|
||||
return Ok(configured);
|
||||
}
|
||||
match std::env::var(env_name) {
|
||||
Ok(value) => value
|
||||
.parse::<u64>()
|
||||
.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<T>(
|
||||
timeout: Option<Duration>,
|
||||
operation: &str,
|
||||
future: impl std::future::Future<Output = Result<T>>,
|
||||
) -> Result<T> {
|
||||
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"]
|
||||
|
||||
@@ -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<SqlTarget> {
|
||||
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<SqlTarget> {
|
||||
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<u16> {
|
||||
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<Channel> {
|
||||
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<ClientTlsConfig> {
|
||||
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<Channel>,
|
||||
headers: &HeaderMap,
|
||||
) -> Result<FlightClient> {
|
||||
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<String, String>) -> 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::<f32>() * 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<Duration>,
|
||||
env_name: &str,
|
||||
default: Option<Duration>,
|
||||
) -> Result<Option<Duration>> {
|
||||
if configured.is_some() {
|
||||
return Ok(configured);
|
||||
}
|
||||
match std::env::var(env_name) {
|
||||
Ok(value) => value
|
||||
.parse::<u64>()
|
||||
.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<T>(
|
||||
timeout: Option<Duration>,
|
||||
operation: &str,
|
||||
future: impl std::future::Future<Output = Result<T>>,
|
||||
) -> Result<T> {
|
||||
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}"),
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user