From cf093c061b7c0daab052e103c3b2e852dafcfcdc Mon Sep 17 00:00:00 2001
From: mdecimus <mauro@stalw.art>
Date: Mon, 7 Oct 2024 10:53:18 +0200
Subject: [PATCH] Allow nonces in OAuth code requests

---
 crates/jmap/src/auth/oauth/auth.rs  |  8 +++++--
 crates/jmap/src/auth/oauth/mod.rs   |  3 +++
 crates/jmap/src/auth/oauth/token.rs | 34 +++++++++++++++++------------
 3 files changed, 29 insertions(+), 16 deletions(-)

diff --git a/crates/jmap/src/auth/oauth/auth.rs b/crates/jmap/src/auth/oauth/auth.rs
index 08acb948..db5cdd49 100644
--- a/crates/jmap/src/auth/oauth/auth.rs
+++ b/crates/jmap/src/auth/oauth/auth.rs
@@ -82,6 +82,7 @@ impl OAuthApiHandler for Server {
             OAuthCodeRequest::Code {
                 client_id,
                 redirect_uri,
+                nonce,
             } => {
                 // Validate clientId
                 if client_id.len() > CLIENT_ID_MAX_LEN {
@@ -109,6 +110,7 @@ impl OAuthApiHandler for Server {
                     status: OAuthStatus::Authorized,
                     account_id: access_token.primary_id(),
                     client_id,
+                    nonce,
                     params: redirect_uri.unwrap_or_default(),
                 })
                 .serialize();
@@ -189,8 +191,8 @@ impl OAuthApiHandler for Server {
         session: HttpSessionData,
     ) -> trc::Result<HttpResponse> {
         // Parse form
-        let client_id = FormData::from_request(req, MAX_POST_LEN, session.session_id)
-            .await?
+        let mut form_data = FormData::from_request(req, MAX_POST_LEN, session.session_id).await?;
+        let client_id = form_data
             .remove("client_id")
             .filter(|client_id| client_id.len() <= CLIENT_ID_MAX_LEN)
             .ok_or_else(|| {
@@ -198,6 +200,7 @@ impl OAuthApiHandler for Server {
                     .into_err()
                     .details("Client ID is missing.")
             })?;
+        let nonce = form_data.remove("nonce");
 
         // Generate device code
         let device_code = thread_rng()
@@ -225,6 +228,7 @@ impl OAuthApiHandler for Server {
             status: OAuthStatus::Pending,
             account_id: u32::MAX,
             client_id,
+            nonce,
             params: device_code.clone(),
         })
         .serialize();
diff --git a/crates/jmap/src/auth/oauth/mod.rs b/crates/jmap/src/auth/oauth/mod.rs
index 738ac989..7fddc8a1 100644
--- a/crates/jmap/src/auth/oauth/mod.rs
+++ b/crates/jmap/src/auth/oauth/mod.rs
@@ -40,6 +40,7 @@ pub struct OAuthCode {
     pub status: OAuthStatus,
     pub account_id: u32,
     pub client_id: String,
+    pub nonce: Option<String>,
     pub params: String,
 }
 
@@ -146,6 +147,8 @@ pub enum OAuthCodeRequest {
     Code {
         client_id: String,
         redirect_uri: Option<String>,
+        #[serde(default)]
+        nonce: Option<String>,
     },
     Device {
         code: String,
diff --git a/crates/jmap/src/auth/oauth/token.rs b/crates/jmap/src/auth/oauth/token.rs
index e629eab3..fcd2f36b 100644
--- a/crates/jmap/src/auth/oauth/token.rs
+++ b/crates/jmap/src/auth/oauth/token.rs
@@ -43,6 +43,7 @@ pub trait TokenHandler: Sync + Send {
         issuer: String,
         nonce: Option<String>,
         with_refresh_token: bool,
+        with_id_token: bool,
     ) -> impl Future<Output = trc::Result<OAuthResponse>> + Send;
 }
 
@@ -64,11 +65,10 @@ impl TokenHandler for Server {
             .await;
 
         if grant_type.eq_ignore_ascii_case("authorization_code") {
-            response = if let (Some(code), Some(client_id), Some(redirect_uri), nonce) = (
+            response = if let (Some(code), Some(client_id), Some(redirect_uri)) = (
                 params.get("code"),
                 params.get("client_id"),
                 params.get("redirect_uri"),
-                params.get("nonce"),
             ) {
                 // Obtain code
                 match self
@@ -106,7 +106,8 @@ impl TokenHandler for Server {
                                     oauth.account_id,
                                     &oauth.client_id,
                                     issuer,
-                                    nonce.map(Into::into),
+                                    oauth.nonce,
+                                    true,
                                     true,
                                 )
                                 .await
@@ -130,11 +131,9 @@ impl TokenHandler for Server {
         } else if grant_type.eq_ignore_ascii_case("urn:ietf:params:oauth:grant-type:device_code") {
             response = TokenResponse::error(ErrorType::ExpiredToken);
 
-            if let (Some(device_code), Some(client_id), nonce) = (
-                params.get("device_code"),
-                params.get("client_id"),
-                params.get("nonce"),
-            ) {
+            if let (Some(device_code), Some(client_id)) =
+                (params.get("device_code"), params.get("client_id"))
+            {
                 // Obtain code
                 if let Some(auth_code) = self
                     .core
@@ -167,7 +166,8 @@ impl TokenHandler for Server {
                                         oauth.account_id,
                                         &oauth.client_id,
                                         issuer,
-                                        nonce.map(Into::into),
+                                        oauth.nonce,
+                                        true,
                                         true,
                                     )
                                     .await
@@ -204,6 +204,7 @@ impl TokenHandler for Server {
                             None,
                             token_info.expires_in
                                 <= self.core.oauth.oauth_expiry_refresh_token_renew,
+                            false,
                         )
                         .await
                         .map(TokenResponse::Granted)
@@ -265,6 +266,7 @@ impl TokenHandler for Server {
         issuer: String,
         nonce: Option<String>,
         with_refresh_token: bool,
+        with_id_token: bool,
     ) -> trc::Result<OAuthResponse> {
         Ok(OAuthResponse {
             access_token: self
@@ -289,12 +291,16 @@ impl TokenHandler for Server {
             } else {
                 None
             },
-            id_token: match self.issue_id_token(account_id.to_string(), issuer, client_id, nonce) {
-                Ok(id_token) => Some(id_token),
-                Err(err) => {
-                    trc::error!(err);
-                    None
+            id_token: if with_id_token {
+                match self.issue_id_token(account_id.to_string(), issuer, client_id, nonce) {
+                    Ok(id_token) => Some(id_token),
+                    Err(err) => {
+                        trc::error!(err);
+                        None
+                    }
                 }
+            } else {
+                None
             },
             scope: None,
         })