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 access_token: None,
62 };
63 let client = Arc::new(
64 AuthClient::new(
65 profile,
66 Arc::new(MemoryStore::new()),
67 AuthOptions::default(),
68 false,
69 )
70 .unwrap(),
71 );
72
73 let (first, second) = tokio::join!(client.token(), client.token());
74 let first = first.unwrap();
75 let second = second.unwrap();
76 assert_eq!(first.access_token, "access");
77 assert_eq!(first.scopes, ["files:read", "jobs"]);
78 assert_eq!(second.access_token, "access");
79 assert_eq!(second.scopes, ["files:read", "jobs"]);
80 assert_eq!(
81 client
82 .refresh_rejected_token("older-rejected-token")
83 .await
84 .unwrap()
85 .access_token,
86 "access"
87 );
88
89 let requests = server.await.unwrap();
90 let request = &requests[0];
91 assert!(request.starts_with("POST /oidc/accounts/account/v1/token HTTP/1.1\r\n"));
92 assert!(request.to_ascii_lowercase().contains(
93 "authorization: basic y2xpzw50onnly3jldA=="
94 .to_ascii_lowercase()
95 .as_str()
96 ));
97 let parameters = url::form_urlencoded::parse(
98 request
99 .split_once("\r\n\r\n")
100 .map(|(_, body)| body)
101 .unwrap_or_default()
102 .as_bytes(),
103 )
104 .into_owned()
105 .collect::<std::collections::HashMap<_, _>>();
106 assert_eq!(
107 parameters.get("grant_type").map(String::as_str),
108 Some("client_credentials")
109 );
110 assert_eq!(
111 parameters.get("scope").map(String::as_str),
112 Some("files:read jobs")
113 );
114 assert_eq!(
115 parameters.get("assume_group").map(String::as_str),
116 Some("group")
117 );
118 }
119
120 #[tokio::test]
121 async fn force_refresh_discovers_a_workspace_token_without_a_cached_credential() {
122 let (host, server) = start_server(2).await;
123 let profile = Profile {
124 name: "service".into(),
125 host: Url::parse(&host).unwrap(),
126 account_id: None,
127 workspace_id: None,
128 client_id: "client".into(),
129 group_id: None,
130 scopes: Vec::new(),
131 target: TargetKind::Workspace,
132 auth_kind: AuthKind::MachineToMachine,
133 client_secret: Some("secret".into()),
134 access_token: None,
135 };
136 let client = AuthClient::new(
137 profile,
138 Arc::new(MemoryStore::new()),
139 AuthOptions::default(),
140 false,
141 )
142 .unwrap();
143
144 let token = client.force_refresh().await.unwrap();
145
146 assert_eq!(token.access_token, "access");
147 assert_eq!(token.scopes, ["all-apis"]);
148 let requests = server.await.unwrap();
149 assert!(requests[0]
150 .starts_with("GET /oidc/.well-known/oauth-authorization-server HTTP/1.1\r\n"));
151 assert!(requests[1].starts_with("POST /token HTTP/1.1\r\n"));
152 let body = requests[1]
153 .split_once("\r\n\r\n")
154 .map(|(_, body)| body)
155 .unwrap_or_default();
156 assert!(body.contains("grant_type=client_credentials"));
157 assert!(body.contains("scope=all-apis"));
158 assert!(!body.contains("offline_access"));
159 }
160
161 async fn start_server(request_count: usize) -> (String, tokio::task::JoinHandle<Vec<String>>) {
162 let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0))
163 .await
164 .unwrap();
165 let host = format!("http://{}", listener.local_addr().unwrap());
166 let token_endpoint = format!("{host}/token");
167 let server = tokio::spawn(async move {
168 let mut requests = Vec::with_capacity(request_count);
169 for _ in 0..request_count {
170 let (mut stream, _) = listener.accept().await.unwrap();
171 let request = read_request(&mut stream).await;
172 let body = if request.starts_with("GET ") {
173 format!(
174 r#"{{"authorization_endpoint":"{token_endpoint}/authorize","token_endpoint":"{token_endpoint}"}}"#
175 )
176 } else {
177 r#"{"access_token":"access","token_type":"Bearer","expires_in":3600}"#.into()
178 };
179 stream
180 .write_all(
181 format!(
182 "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
183 body.len()
184 )
185 .as_bytes(),
186 )
187 .await
188 .unwrap();
189 requests.push(request);
190 }
191 requests
192 });
193 (host, server)
194 }
195
196 async fn read_request(stream: &mut tokio::net::TcpStream) -> String {
197 let mut request = Vec::new();
198 loop {
199 let mut buffer = [0_u8; 4096];
200 let read = stream.read(&mut buffer).await.unwrap();
201 if read == 0 {
202 break;
203 }
204 request.extend_from_slice(&buffer[..read]);
205 let Some(header_end) = request.windows(4).position(|window| window == b"\r\n\r\n")
206 else {
207 continue;
208 };
209 let headers = String::from_utf8_lossy(&request[..header_end]);
210 let content_length = headers
211 .lines()
212 .find_map(|line| {
213 let (name, value) = line.split_once(':')?;
214 name.eq_ignore_ascii_case("content-length")
215 .then(|| value.trim().parse::<usize>().ok())
216 .flatten()
217 })
218 .unwrap_or_default();
219 if request.len() >= header_end + 4 + content_length {
220 break;
221 }
222 }
223 String::from_utf8(request).unwrap()
224 }
225}