dbx_tools_databricks_auth/
m2m.rs

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

dbx_tools_databricks_auth/
m2m.rs
