use crate::{oauth_endpoints, Profile, Result, Token};
pub struct MachineToMachineFlow {
profile: Profile,
http: reqwest::Client,
}
impl MachineToMachineFlow {
pub fn new(profile: Profile) -> Result<Self> {
Ok(Self {
profile,
http: reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.build()?,
})
}
pub async fn token(&self) -> Result<Token> {
let endpoints = oauth_endpoints::resolve(&self.profile, &self.http).await?;
dbx_tools_auth::OAuthFlow::new(dbx_tools_auth::OAuthConfig {
provider: "databricks".into(),
authorization_endpoint: endpoints.authorization_endpoint,
token_endpoint: endpoints.token_endpoint,
client_id: self.profile.client_id.clone(),
client_secret: self.profile.client_secret().map(str::to_owned),
scopes: self.profile.machine_scopes(),
extra_token_params: self
.profile
.group_id
.iter()
.map(|group| ("assume_group".into(), group.clone()))
.collect(),
host: Some(self.profile.host.to_string()),
})?
.client_credentials()
.await
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use url::Url;
use super::*;
use crate::{AuthClient, AuthKind, AuthOptions, AuthSession, MemoryStore, TargetKind};
#[tokio::test]
async fn mints_and_caches_an_account_token_with_cli_request_semantics() {
let (host, server) = start_server(1).await;
let profile = Profile {
name: "service".into(),
host: Url::parse(&host).unwrap(),
account_id: Some("account".into()),
workspace_id: None,
client_id: "client".into(),
group_id: Some("group".into()),
scopes: vec!["jobs".into(), "files:read".into()],
target: TargetKind::Account,
auth_kind: AuthKind::MachineToMachine,
client_secret: Some("secret".into()),
access_token: None,
};
let client = Arc::new(
AuthClient::new(
profile,
Arc::new(MemoryStore::new()),
AuthOptions::default(),
false,
)
.unwrap(),
);
let (first, second) = tokio::join!(client.token(), client.token());
let first = first.unwrap();
let second = second.unwrap();
assert_eq!(first.access_token, "access");
assert_eq!(first.scopes, ["files:read", "jobs"]);
assert_eq!(second.access_token, "access");
assert_eq!(second.scopes, ["files:read", "jobs"]);
assert_eq!(
client
.refresh_rejected_token("older-rejected-token")
.await
.unwrap()
.access_token,
"access"
);
let requests = server.await.unwrap();
let request = &requests[0];
assert!(request.starts_with("POST /oidc/accounts/account/v1/token HTTP/1.1\r\n"));
assert!(request.to_ascii_lowercase().contains(
"authorization: basic y2xpzw50onnly3jldA=="
.to_ascii_lowercase()
.as_str()
));
let parameters = url::form_urlencoded::parse(
request
.split_once("\r\n\r\n")
.map(|(_, body)| body)
.unwrap_or_default()
.as_bytes(),
)
.into_owned()
.collect::<std::collections::HashMap<_, _>>();
assert_eq!(
parameters.get("grant_type").map(String::as_str),
Some("client_credentials")
);
assert_eq!(
parameters.get("scope").map(String::as_str),
Some("files:read jobs")
);
assert_eq!(
parameters.get("assume_group").map(String::as_str),
Some("group")
);
}
#[tokio::test]
async fn force_refresh_discovers_a_workspace_token_without_a_cached_credential() {
let (host, server) = start_server(2).await;
let profile = Profile {
name: "service".into(),
host: Url::parse(&host).unwrap(),
account_id: None,
workspace_id: None,
client_id: "client".into(),
group_id: None,
scopes: Vec::new(),
target: TargetKind::Workspace,
auth_kind: AuthKind::MachineToMachine,
client_secret: Some("secret".into()),
access_token: None,
};
let client = AuthClient::new(
profile,
Arc::new(MemoryStore::new()),
AuthOptions::default(),
false,
)
.unwrap();
let token = client.force_refresh().await.unwrap();
assert_eq!(token.access_token, "access");
assert_eq!(token.scopes, ["all-apis"]);
let requests = server.await.unwrap();
assert!(requests[0]
.starts_with("GET /oidc/.well-known/oauth-authorization-server HTTP/1.1\r\n"));
assert!(requests[1].starts_with("POST /token HTTP/1.1\r\n"));
let body = requests[1]
.split_once("\r\n\r\n")
.map(|(_, body)| body)
.unwrap_or_default();
assert!(body.contains("grant_type=client_credentials"));
assert!(body.contains("scope=all-apis"));
assert!(!body.contains("offline_access"));
}
async fn start_server(request_count: usize) -> (String, tokio::task::JoinHandle<Vec<String>>) {
let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0))
.await
.unwrap();
let host = format!("http://{}", listener.local_addr().unwrap());
let token_endpoint = format!("{host}/token");
let server = tokio::spawn(async move {
let mut requests = Vec::with_capacity(request_count);
for _ in 0..request_count {
let (mut stream, _) = listener.accept().await.unwrap();
let request = read_request(&mut stream).await;
let body = if request.starts_with("GET ") {
format!(
r#"{{"authorization_endpoint":"{token_endpoint}/authorize","token_endpoint":"{token_endpoint}"}}"#
)
} else {
r#"{"access_token":"access","token_type":"Bearer","expires_in":3600}"#.into()
};
stream
.write_all(
format!(
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
body.len()
)
.as_bytes(),
)
.await
.unwrap();
requests.push(request);
}
requests
});
(host, server)
}
async fn read_request(stream: &mut tokio::net::TcpStream) -> String {
let mut request = Vec::new();
loop {
let mut buffer = [0_u8; 4096];
let read = stream.read(&mut buffer).await.unwrap();
if read == 0 {
break;
}
request.extend_from_slice(&buffer[..read]);
let Some(header_end) = request.windows(4).position(|window| window == b"\r\n\r\n")
else {
continue;
};
let headers = String::from_utf8_lossy(&request[..header_end]);
let content_length = headers
.lines()
.find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length")
.then(|| value.trim().parse::<usize>().ok())
.flatten()
})
.unwrap_or_default();
if request.len() >= header_end + 4 + content_length {
break;
}
}
String::from_utf8(request).unwrap()
}
}