dbx_tools_auth/
oauth.rs

1use std::{collections::HashMap, net::SocketAddr, time::Duration};
2
3use oauth2::{
4    basic::BasicClient, AuthUrl, AuthorizationCode, ClientId, CsrfToken, EndpointNotSet,
5    EndpointSet, PkceCodeChallenge, RedirectUrl, RefreshToken, Scope, TokenUrl,
6};
7use tokio::{
8    io::{AsyncReadExt, AsyncWriteExt},
9    net::TcpListener,
10};
11use url::Url;
12
13use crate::{token::OAuthTokenResponse, Error, OAuthTemplate, OAuthTemplateContext, Result, Token};
14
15const DEFAULT_PORT: u16 = 8020;
16const MAX_PORT: u16 = 8040;
17
18type OAuthClient =
19    BasicClient<EndpointSet, EndpointNotSet, EndpointNotSet, EndpointNotSet, EndpointSet>;
20
21#[derive(Clone)]
22pub struct OAuthConfig {
23    pub provider: String,
24    pub authorization_endpoint: String,
25    pub token_endpoint: String,
26    pub client_id: String,
27    pub client_secret: Option<String>,
28    pub scopes: Vec<String>,
29    pub extra_token_params: Vec<(String, String)>,
30    pub host: Option<String>,
31}
32
33pub struct OAuthFlow {
34    config: OAuthConfig,
35    http: reqwest::Client,
36    template: OAuthTemplate,
37}
38
39impl OAuthFlow {
40    pub fn new(config: OAuthConfig) -> Result<Self> {
41        for endpoint in [&config.authorization_endpoint, &config.token_endpoint] {
42            let url = Url::parse(endpoint)?;
43            let loopback = url.host_str().is_some_and(|host| {
44                host == "localhost"
45                    || host
46                        .parse::<std::net::IpAddr>()
47                        .is_ok_and(|address| address.is_loopback())
48            });
49            if url.scheme() != "https" && !(url.scheme() == "http" && loopback) {
50                return Err(Error::Config(
51                    "OAuth endpoints must use HTTPS or loopback HTTP".into(),
52                ));
53            }
54        }
55        let http = reqwest::Client::builder()
56            .redirect(reqwest::redirect::Policy::none())
57            .build()?;
58        Ok(Self {
59            config,
60            http,
61            template: OAuthTemplate::default(),
62        })
63    }
64
65    /// Set the branding used by the browser callback page.
66    pub fn with_template(mut self, template: OAuthTemplate) -> Self {
67        self.template = template;
68        self
69    }
70
71    pub async fn login(&self, timeout: Duration) -> Result<Token> {
72        let (listener, address) = bind_callback().await?;
73        let redirect = format!("http://localhost:{}", address.port());
74        let client = self.client(&redirect)?;
75        let (challenge, verifier) = PkceCodeChallenge::new_random_sha256();
76        let mut request = client
77            .authorize_url(CsrfToken::new_random)
78            .set_pkce_challenge(challenge);
79        for scope in self.config.scopes.clone() {
80            request = request.add_scope(Scope::new(scope));
81        }
82        let (authorization_url, csrf) = request.url();
83        if open::that(authorization_url.as_str()).is_err() {
84            eprintln!("Open this URL in a browser:\n{authorization_url}");
85        }
86
87        let callback = tokio::time::timeout(
88            timeout,
89            receive_callback(
90                listener,
91                &redirect,
92                &self.template,
93                self.config.host.as_deref(),
94                csrf.secret(),
95            ),
96        )
97        .await
98        .map_err(|_| Error::OAuth("timed out waiting for browser authorization".into()))??;
99        if callback.state.as_deref() != Some(csrf.secret()) {
100            return Err(Error::OAuth("OAuth state did not match".into()));
101        }
102        if let Some(error) = callback.error {
103            return Err(Error::OAuth(match callback.error_description {
104                Some(description) => format!("{error}: {description}"),
105                None => error,
106            }));
107        }
108        let code = callback
109            .code
110            .ok_or_else(|| Error::OAuth("authorization callback contained no code".into()))?;
111        let response = client
112            .exchange_code(AuthorizationCode::new(code))
113            .set_pkce_verifier(verifier)
114            .request_async(&self.http)
115            .await
116            .map_err(|error| {
117                Error::OAuth(format!("authorization-code exchange failed: {error}"))
118            })?;
119        let mut token = Token::from_response(&response, time::OffsetDateTime::now_utc(), None)?;
120        if token.scopes.is_empty() {
121            token.scopes = self.config.scopes.clone();
122        }
123        Ok(token)
124    }
125
126    pub async fn refresh(&self, token: &Token) -> Result<Token> {
127        let client = self.client("http://localhost:8020")?;
128        let refresh = token
129            .refresh_token()
130            .ok_or_else(|| Error::LoginRequired(self.config.provider.clone()))?;
131        let response: OAuthTokenResponse = client
132            .exchange_refresh_token(&RefreshToken::new(refresh.secret().to_owned()))
133            .request_async(&self.http)
134            .await
135            .map_err(|error| Error::OAuth(format!("refresh-token exchange failed: {error}")))?;
136        Token::from_response(&response, time::OffsetDateTime::now_utc(), Some(token))
137    }
138
139    pub async fn client_credentials(&self) -> Result<Token> {
140        let secret =
141            self.config.client_secret.clone().ok_or_else(|| {
142                Error::Config("client credentials require a client secret".into())
143            })?;
144        let client = self
145            .client("http://localhost:8020")?
146            .set_client_secret(oauth2::ClientSecret::new(secret))
147            .set_auth_type(oauth2::AuthType::BasicAuth);
148        let mut request = client.exchange_client_credentials();
149        for scope in &self.config.scopes {
150            request = request.add_scope(Scope::new(scope.clone()));
151        }
152        for (name, value) in &self.config.extra_token_params {
153            request = request.add_extra_param(name, value);
154        }
155        let response = request.request_async(&self.http).await.map_err(|error| {
156            Error::OAuth(format!("client-credentials exchange failed: {error}"))
157        })?;
158        let mut token = Token::from_response(&response, time::OffsetDateTime::now_utc(), None)?;
159        if token.scopes.is_empty() {
160            token.scopes = self.config.scopes.clone();
161        }
162        Ok(token)
163    }
164
165    fn client(&self, redirect: &str) -> Result<OAuthClient> {
166        let client = BasicClient::new(ClientId::new(self.config.client_id.clone()))
167            .set_auth_uri(
168                AuthUrl::new(self.config.authorization_endpoint.clone())
169                    .map_err(|error| Error::OAuth(error.to_string()))?,
170            )
171            .set_token_uri(
172                TokenUrl::new(self.config.token_endpoint.clone())
173                    .map_err(|error| Error::OAuth(error.to_string()))?,
174            )
175            .set_redirect_uri(
176                RedirectUrl::new(redirect.to_owned())
177                    .map_err(|error| Error::OAuth(error.to_string()))?,
178            );
179        Ok(match &self.config.client_secret {
180            Some(secret) => client.set_client_secret(oauth2::ClientSecret::new(secret.clone())),
181            None => client,
182        })
183    }
184}
185
186#[derive(Debug)]
187struct Callback {
188    code: Option<String>,
189    state: Option<String>,
190    error: Option<String>,
191    error_description: Option<String>,
192}
193
194async fn bind_callback() -> Result<(TcpListener, SocketAddr)> {
195    for port in DEFAULT_PORT..=MAX_PORT {
196        if let Ok(listener) = TcpListener::bind(("127.0.0.1", port)).await {
197            let address = listener.local_addr()?;
198            return Ok((listener, address));
199        }
200    }
201    Err(Error::OAuth(format!(
202        "no callback port available from {DEFAULT_PORT} through {MAX_PORT}"
203    )))
204}
205
206async fn receive_callback(
207    listener: TcpListener,
208    redirect: &str,
209    template: &OAuthTemplate,
210    host: Option<&str>,
211    expected_state: &str,
212) -> Result<Callback> {
213    let (mut stream, _) = listener.accept().await?;
214    let mut buffer = vec![0; 16 * 1024];
215    let read = stream.read(&mut buffer).await?;
216    let request = std::str::from_utf8(&buffer[..read])
217        .map_err(|error| Error::OAuth(format!("invalid callback request: {error}")))?;
218    let target = request
219        .lines()
220        .next()
221        .and_then(|line| line.split_whitespace().nth(1))
222        .ok_or_else(|| Error::OAuth("invalid callback request line".into()))?;
223    let url = Url::parse(&format!("{redirect}{target}"))?;
224    let values: HashMap<String, String> = url
225        .query_pairs()
226        .map(|(key, value)| (key.into_owned(), value.into_owned()))
227        .collect();
228    let callback = Callback {
229        code: values.get("code").cloned(),
230        state: values.get("state").cloned(),
231        error: values.get("error").cloned(),
232        error_description: values.get("error_description").cloned(),
233    };
234    let successful = callback.state.as_deref() == Some(expected_state)
235        && callback.error.is_none()
236        && callback.code.is_some();
237    let body = callback_response(template, host, &callback, successful);
238    let status = if successful {
239        "200 OK"
240    } else {
241        "400 Bad Request"
242    };
243    let response = format!(
244        "HTTP/1.1 {status}\r\nContent-Type: text/html; charset=utf-8\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
245        body.len(),
246    );
247    stream.write_all(response.as_bytes()).await?;
248    Ok(callback)
249}
250
251fn callback_response(
252    template: &OAuthTemplate,
253    host: Option<&str>,
254    callback: &Callback,
255    successful: bool,
256) -> String {
257    let default_error = (!successful && callback.error.is_none()).then_some("authorization_failed");
258    let error = callback.error.as_deref().or(default_error);
259    template.render(OAuthTemplateContext {
260        host,
261        error,
262        error_description: callback.error_description.as_deref(),
263    })
264}
dbx_tools_auth/oauth.rs

dbx_tools_auth/
oauth.rs
