dbx_tools_databricks_auth/
m2m.rs1use crate::{oauth_endpoints, Profile, Result, Token};
2pub struct MachineToMachineFlow {
3 profile: Profile,
4 http: reqwest::Client,
5}
6impl MachineToMachineFlow {
7 pub fn new(profile: Profile) -> Result<Self> {
8 Ok(Self {
9 profile,
10 http: reqwest::Client::builder()
11 .redirect(reqwest::redirect::Policy::none())
12 .build()?,
13 })
14 }
15 pub async fn token(&self) -> Result<Token> {
16 let endpoints = oauth_endpoints::resolve(&self.profile, &self.http).await?;
17 dbx_tools_auth::OAuthFlow::new(dbx_tools_auth::OAuthConfig {
18 provider: "databricks".into(),
19 authorization_endpoint: endpoints.authorization_endpoint,
20 token_endpoint: endpoints.token_endpoint,
21 client_id: self.profile.client_id.clone(),
22 client_secret: self.profile.client_secret().map(str::to_owned),
23 scopes: self.profile.machine_scopes(),
24 extra_token_params: self
25 .profile
26 .group_id
27 .iter()
28 .map(|group| ("assume_group".into(), group.clone()))
29 .collect(),
30 host: Some(self.profile.host.to_string()),
31 })?
32 .client_credentials()
33 .await
34 }
35}
36
37#[cfg(test)]
38mod tests {
39 use std::sync::Arc;
40
41 use tokio::io::{AsyncReadExt, AsyncWriteExt};
42 use url::Url;
43
44 use super::*;
45 use crate::{AuthClient, AuthKind, AuthOptions, AuthSession, MemoryStore, TargetKind};
46
47 #[tokio::test]
48 async fn mints_and_caches_an_account_token_with_cli_request_semantics() {
49 let (host, server) = start_server(1).await;
50 let profile = Profile {
51 name: "service".into(),
52 host: Url::parse(&host).unwrap(),
53 account_id: Some("account".into()),
54 workspace_id: None,
55 client_id: "client".into(),
56 group_id: Some("group".into()),
57 scopes: vec!["jobs".into(), "files:read".into()],
58 target: TargetKind::Account,
59 auth_kind: AuthKind::MachineToMachine,
60 client_secret: Some("secret".into()),
61 };
62 let client = Arc::new(
63 AuthClient::new(
64 profile,
65 Arc::new(MemoryStore::new()),
66 AuthOptions::default(),
67 false,
68 )
69 .unwrap(),
70 );
71
72 let (first, second) = tokio::join!(client.token(), client.token());
73 let first = first.unwrap();
74 let second = second.unwrap();
75 assert_eq!(first.access_token, "access");
76 assert_eq!(first.scopes, ["files:read", "jobs"]);
77 assert_eq!(second.access_token, "access");
78 assert_eq!(second.scopes, ["files:read", "jobs"]);
79 assert_eq!(
80 client
81 .refresh_rejected_token("older-rejected-token")
82 .await
83 .unwrap()
84 .access_token,
85 "access"
86 );
87
88 let requests = server.await.unwrap();
89 let request = &requests[0];
90 assert!(request.starts_with("POST /oidc/accounts/account/v1/token HTTP/1.1\r\n"));
91 assert!(request.to_ascii_lowercase().contains(
92 "authorization: basic y2xpzw50onnly3jldA=="
93 .to_ascii_lowercase()
94 .as_str()
95 ));
96 let parameters = url::form_urlencoded::parse(
97 request
98 .split_once("\r\n\r\n")
99 .map(|(_, body)| body)
100 .unwrap_or_default()
101 .as_bytes(),
102 )
103 .into_owned()
104 .collect::<std::collections::HashMap<_, _>>();
105 assert_eq!(
106 parameters.get("grant_type").map(String::as_str),
107 Some("client_credentials")
108 );
109 assert_eq!(
110 parameters.get("scope").map(String::as_str),
111 Some("files:read jobs")
112 );
113 assert_eq!(
114 parameters.get("assume_group").map(String::as_str),
115 Some("group")
116 );
117 }
118
119 #[tokio::test]
120 async fn force_refresh_discovers_a_workspace_token_without_a_cached_credential() {
121 let (host, server) = start_server(2).await;
122 let profile = Profile {
123 name: "service".into(),
124 host: Url::parse(&host).unwrap(),
125 account_id: None,
126 workspace_id: None,
127 client_id: "client".into(),
128 group_id: None,
129 scopes: Vec::new(),
130 target: TargetKind::Workspace,
131 auth_kind: AuthKind::MachineToMachine,
132 client_secret: Some("secret".into()),
133 };
134 let client = AuthClient::new(
135 profile,
136 Arc::new(MemoryStore::new()),
137 AuthOptions::default(),
138 false,
139 )
140 .unwrap();
141
142 let token = client.force_refresh().await.unwrap();
143
144 assert_eq!(token.access_token, "access");
145 assert_eq!(token.scopes, ["all-apis"]);
146 let requests = server.await.unwrap();
147 assert!(requests[0]
148 .starts_with("GET /oidc/.well-known/oauth-authorization-server HTTP/1.1\r\n"));
149 assert!(requests[1].starts_with("POST /token HTTP/1.1\r\n"));
150 let body = requests[1]
151 .split_once("\r\n\r\n")
152 .map(|(_, body)| body)
153 .unwrap_or_default();
154 assert!(body.contains("grant_type=client_credentials"));
155 assert!(body.contains("scope=all-apis"));
156 assert!(!body.contains("offline_access"));
157 }
158
159 async fn start_server(request_count: usize) -> (String, tokio::task::JoinHandle<Vec<String>>) {
160 let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0))
161 .await
162 .unwrap();
163 let host = format!("http://{}", listener.local_addr().unwrap());
164 let token_endpoint = format!("{host}/token");
165 let server = tokio::spawn(async move {
166 let mut requests = Vec::with_capacity(request_count);
167 for _ in 0..request_count {
168 let (mut stream, _) = listener.accept().await.unwrap();
169 let request = read_request(&mut stream).await;
170 let body = if request.starts_with("GET ") {
171 format!(
172 r#"{{"authorization_endpoint":"{token_endpoint}/authorize","token_endpoint":"{token_endpoint}"}}"#
173 )
174 } else {
175 r#"{"access_token":"access","token_type":"Bearer","expires_in":3600}"#.into()
176 };
177 stream
178 .write_all(
179 format!(
180 "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
181 body.len()
182 )
183 .as_bytes(),
184 )
185 .await
186 .unwrap();
187 requests.push(request);
188 }
189 requests
190 });
191 (host, server)
192 }
193
194 async fn read_request(stream: &mut tokio::net::TcpStream) -> String {
195 let mut request = Vec::new();
196 loop {
197 let mut buffer = [0_u8; 4096];
198 let read = stream.read(&mut buffer).await.unwrap();
199 if read == 0 {
200 break;
201 }
202 request.extend_from_slice(&buffer[..read]);
203 let Some(header_end) = request.windows(4).position(|window| window == b"\r\n\r\n")
204 else {
205 continue;
206 };
207 let headers = String::from_utf8_lossy(&request[..header_end]);
208 let content_length = headers
209 .lines()
210 .find_map(|line| {
211 let (name, value) = line.split_once(':')?;
212 name.eq_ignore_ascii_case("content-length")
213 .then(|| value.trim().parse::<usize>().ok())
214 .flatten()
215 })
216 .unwrap_or_default();
217 if request.len() >= header_end + 4 + content_length {
218 break;
219 }
220 }
221 String::from_utf8(request).unwrap()
222 }
223}