Refresh token support for OAuth 2 clients

src/agent.rs

@@ -431,16 +431,33 @@ impl Agent {
         }
 
         // =====
-        if session.token.server.aws_sso.is_some() {
-            if let Some(new) = refresh_token_using_awssso(&session).await {
-                return self.session_manager.add(new).unwrap(); // XXX:
-            }
+        let maybe_new = if session.token.server.aws_sso.is_some() {
+            refresh_token_using_awssso(&session).await
+        } else {
+            refresh_token_using_oauth2(&session).await
+        };
+        if let Some(new) = maybe_new {
+            self.session_manager.add(new).unwrap() // XXX: unwrap()
+        } else {
+            session
         }
-        // TODO: general refresh_token implementation
+    }
+}
 
-        // =====
-        tracing::warn!( server_id = session.token.server.id(), url = %session.token.server.url, "Skipped token refresh due to missing implementation to refresh this session");
-        session
+async fn refresh_token_using_oauth2(
+    session: &crate::session_manager::Session,
+) -> Option<crate::token::ServerToken> {
+    match crate::oauth_refresh_token::refresh_token(&session.token).await {
+        Ok(token) => Some(token),
+        Err(e) => {
+            tracing::warn!(
+                server_id = session.token.server.id(),
+                url = %session.token.server.url,
+                err = ?e,
+                "Failed to refresh session using refresh_token"
+            );
+            None
+        }
     }
 }
 

src/ext_oauth2.rs

@@ -112,6 +112,21 @@ where
     }
 }
 
+pub(crate) fn reqwest_give_client_auth(
+    req: reqwest::RequestBuilder,
+    oauth: &crate::config::ServerOAuth,
+    params: &[(&str, &str)],
+) -> reqwest::RequestBuilder {
+    match oauth.client_secret {
+        Some(ref secret) => req.basic_auth(&oauth.client_id, Some(secret)).form(params),
+        None => {
+            let auth_param = [("client_id", oauth.client_id.as_ref())];
+            let new_params: Vec<_> = auth_param.iter().chain(params.iter()).collect();
+            req.form(&new_params)
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/lib.rs

@@ -17,6 +17,7 @@ pub mod utils;
 pub mod oauth_awssso;
 pub mod oauth_code;
 pub mod oauth_device_code;
+pub mod oauth_refresh_token;
 
 pub mod api_client;
 pub mod awssso_client;

src/oauth_device_code.rs

@@ -37,21 +37,23 @@ impl OAuthDeviceCodeFlow {
 
         // oauth2 crate doesn't allow non-standard response type
         let scopes = oauth.scope.join(" ");
-        let resp = crate::client::http()
+        let req = crate::client::http()
             .post(grant.device_authorization_endpoint.clone().ok_or_else(|| {
                 crate::Error::ConfigError(format!(
                     "{} is missing device_authorization_endpoint",
                     server.id()
                 ))
             })?)
-            .header(reqwest::header::ACCEPT, "application/json")
-            .basic_auth(&oauth.client_id, oauth.client_secret.as_ref())
-            .form(&[
+            .header(reqwest::header::ACCEPT, "application/json");
+        let req = crate::ext_oauth2::reqwest_give_client_auth(
+            req,
+            oauth,
+            &[
                 ("grant_type", "urn:ietf:params:oauth:grant-type:device_code"),
                 ("scope", &scopes),
-            ])
-            .send()
-            .await?;
+            ],
+        );
+        let resp = req.send().await?;
 
         let status = resp.status();
         if status != reqwest::StatusCode::OK {
@@ -96,15 +98,16 @@ impl OAuthDeviceCodeFlow {
             .post(oauth.token_endpoint.clone().ok_or_else(|| {
                 crate::Error::ConfigError(format!("{} is missing token_endpoint", self.server.id()))
             })?)
-            .header(reqwest::header::ACCEPT, "application/json")
-            .basic_auth(&oauth.client_id, oauth.client_secret.as_ref())
-            .form(&[
+            .header(reqwest::header::ACCEPT, "application/json");
+        let req = crate::ext_oauth2::reqwest_give_client_auth(
+            req,
+            oauth,
+            &[
                 ("grant_type", "urn:ietf:params:oauth:grant-type:device_code"),
                 ("device_code", self.device_code.expose_secret()),
-            ])
-            .build()?;
-        tracing::debug!(req = ?req, "complete req");
-        let resp = crate::client::http().execute(req).await?;
+            ],
+        );
+        let resp = req.send().await?;
 
         let status = resp.status();
         if status != reqwest::StatusCode::OK {

src/oauth_refresh_token.rs

@@ -0,0 +1,85 @@
+pub(crate) async fn refresh_token(
+    token: &crate::token::ServerToken,
+) -> crate::Result<crate::token::ServerToken> {
+    let flow = OAuthRefreshTokenFlow::try_from(token)?;
+    flow.perform().await
+}
+
+#[derive(Debug)]
+pub struct OAuthRefreshTokenFlow {
+    server: crate::config::Server,
+    refresh_token: secrecy::SecretString,
+}
+
+impl TryFrom<&crate::token::ServerToken> for OAuthRefreshTokenFlow {
+    type Error = crate::Error;
+    fn try_from(value: &crate::token::ServerToken) -> Result<Self, Self::Error> {
+        let sid = value.server.id();
+        if !value.has_active_refresh_token() {
+            return Err(crate::Error::ConfigError(format!(
+                "{sid} token does not have refresh_token"
+            )));
+        }
+        Ok(Self {
+            refresh_token: value.refresh_token.clone().unwrap(),
+            server: value.server.clone(),
+        })
+    }
+}
+
+impl OAuthRefreshTokenFlow {
+    pub async fn perform(&self) -> crate::Result<crate::token::ServerToken> {
+        use secrecy::ExposeSecret;
+        tracing::info!(flow = ?self, "Refreshing access token");
+        let client = oauth2_client_from_server(&self.server)?;
+        let rt = oauth2::RefreshToken::new(self.refresh_token.expose_secret().to_string());
+        let req = client.exchange_refresh_token(&rt).add_scopes(
+            // scope is not mandatory on the standard, however some AS refuses to refresh access tokens when omit (e.g. Microsoft)
+            self.server
+                .oauth
+                .as_ref()
+                .unwrap() // XXX: unwrap(), ensured by oauth2_client_from_server
+                .scope
+                .iter()
+                .map(|x| oauth2::Scope::new(x.to_owned())),
+        );
+        let resp = req.request_async(&crate::client::http()).await?;
+        Ok(crate::token::ServerToken::from_token_response(
+            self.server.clone(),
+            resp,
+        ))
+    }
+}
+
+fn oauth2_client_from_server(
+    server: &crate::config::Server,
+) -> crate::Result<
+    crate::ext_oauth2::SecrecyClient<
+        oauth2::EndpointNotSet,
+        oauth2::EndpointNotSet,
+        oauth2::EndpointNotSet,
+        oauth2::EndpointNotSet,
+        oauth2::EndpointSet,
+    >,
+> {
+    let oauth = server.oauth.as_ref().ok_or_else(|| {
+        crate::Error::ConfigError(format!(
+            "Server '{}' is missing OAuth 2.0 client configuration",
+            server.id()
+        ))
+    })?;
+
+    let mut client =
+        crate::ext_oauth2::SecrecyClient::new(oauth2::ClientId::new(oauth.client_id.clone()))
+            .set_token_uri(oauth2::TokenUrl::from_url(
+                oauth
+                    .token_endpoint
+                    .clone()
+                    .map(Ok)
+                    .unwrap_or_else(|| server.url.join("oauth/token"))?,
+            ));
+    if let Some(ref secret) = oauth.client_secret {
+        client = client.set_client_secret(oauth2::ClientSecret::new(secret.to_owned()));
+    }
+    Ok(client)
+}