Pass the rng and clock around

crates/handlers/src/oauth2/authorization/complete.rs

@@ -190,6 +190,8 @@ pub(crate) async fn complete(
     policy_factory: &PolicyFactory,
     mut txn: Transaction<'_, Postgres>,
 ) -> Result<AuthorizationResponse<Option<AccessTokenResponse>>, GrantCompletionError> {
+    let (clock, mut rng) = crate::rng_and_clock()?;
+
     // Verify that the grant is in a pending stage
     if !grant.stage.is_pending() {
         return Err(GrantCompletionError::NotPending);
@@ -226,7 +228,7 @@ pub(crate) async fn complete(
     }
 
     // All good, let's start the session
-    let session = derive_session(&mut txn, &grant, browser_session).await?;
+    let session = derive_session(&mut txn, &mut rng, &clock, &grant, browser_session).await?;
 
     let grant = fulfill_grant(&mut txn, grant, session.clone()).await?;
 

crates/handlers/src/oauth2/authorization/mod.rs

@@ -37,7 +37,7 @@ use oauth2_types::{
     requests::{AuthorizationRequest, GrantType, Prompt, ResponseMode},
     response_type::ResponseType,
 };
-use rand::{distributions::Alphanumeric, thread_rng, Rng};
+use rand::{distributions::Alphanumeric, Rng};
 use serde::Deserialize;
 use sqlx::PgPool;
 use thiserror::Error;
@@ -159,6 +159,7 @@ pub(crate) async fn get(
     cookie_jar: PrivateCookieJar<Encrypter>,
     Form(params): Form<Params>,
 ) -> Result<Response, RouteError> {
+    let (clock, mut rng) = crate::rng_and_clock()?;
     let mut txn = pool.begin().await?;
 
     // First, figure out what client it is
@@ -265,7 +266,7 @@ pub(crate) async fn get(
                 }
 
                 // 32 random alphanumeric characters, about 190bit of entropy
-                let code: String = thread_rng()
+                let code: String = (&mut rng)
                     .sample_iter(&Alphanumeric)
                     .take(32)
                     .map(char::from)
@@ -296,6 +297,8 @@ pub(crate) async fn get(
 
             let grant = new_authorization_grant(
                 &mut txn,
+                &mut rng,
+                &clock,
                 client,
                 redirect_uri.clone(),
                 params.auth.scope,

crates/handlers/src/oauth2/consent.rs

@@ -119,6 +119,7 @@ pub(crate) async fn post(
     Path(grant_id): Path<Ulid>,
     Form(form): Form<ProtectedForm<()>>,
 ) -> Result<Response, RouteError> {
+    let (clock, mut rng) = crate::rng_and_clock()?;
     let mut txn = pool
         .begin()
         .await
@@ -163,6 +164,8 @@ pub(crate) async fn post(
         .collect();
     insert_client_consent(
         &mut txn,
+        &mut rng,
+        &clock,
         &session.user,
         &grant.client,
         &scope_without_device,

crates/handlers/src/oauth2/introspection.rs

@@ -28,6 +28,7 @@ use mas_storage::{
         client::ClientFetchError,
         refresh_token::{lookup_active_refresh_token, RefreshTokenLookupError},
     },
+    Clock,
 };
 use oauth2_types::requests::{IntrospectionRequest, IntrospectionResponse};
 use sqlx::PgPool;
@@ -158,6 +159,7 @@ pub(crate) async fn post(
     State(encrypter): State<Encrypter>,
     client_authorization: ClientAuthorization<IntrospectionRequest>,
 ) -> Result<impl IntoResponse, RouteError> {
+    let clock = Clock::default();
     let mut conn = pool.acquire().await?;
 
     let client = client_authorization.credentials.fetch(&mut conn).await?;
@@ -227,7 +229,8 @@ pub(crate) async fn post(
             }
         }
         TokenType::CompatAccessToken => {
-            let (token, session) = lookup_active_compat_access_token(&mut conn, token).await?;
+            let (token, session) =
+                lookup_active_compat_access_token(&mut conn, &clock, token).await?;
 
             let device_scope = session.device.to_scope_token();
             let scope = [device_scope].into_iter().collect();

crates/handlers/src/oauth2/token.rs

@@ -50,7 +50,6 @@ use oauth2_types::{
     },
     scope,
 };
-use rand::thread_rng;
 use serde::Serialize;
 use serde_with::{serde_as, skip_serializing_none};
 use sqlx::{PgPool, Postgres, Transaction};
@@ -235,12 +234,13 @@ async fn authorization_code_grant(
     url_builder: &UrlBuilder,
     mut txn: Transaction<'_, Postgres>,
 ) -> Result<AccessTokenResponse, RouteError> {
+    let (clock, mut rng) = crate::rng_and_clock()?;
+
     // TODO: there is a bunch of unnecessary cloning here
     // TODO: handle "not found" cases
     let authz_grant = lookup_grant_by_code(&mut txn, &grant.code).await?;
 
-    // TODO: that's not a timestamp from the DB. Let's assume they are in sync
-    let now = Utc::now();
+    let now = clock.now();
 
     let session = match authz_grant.stage {
         AuthorizationGrantStage::Cancelled { cancelled_at } => {
@@ -257,7 +257,7 @@ async fn authorization_code_grant(
             // Ending the session if the token was already exchanged more than 20s ago
             if now - exchanged_at > Duration::seconds(20) {
                 debug!("Ending potentially compromised session");
-                end_oauth_session(&mut txn, session).await?;
+                end_oauth_session(&mut txn, &clock, session).await?;
                 txn.commit().await?;
             }
 
@@ -303,22 +303,32 @@ async fn authorization_code_grant(
     let browser_session = &session.browser_session;
 
     let ttl = Duration::minutes(5);
-    let (access_token_str, refresh_token_str) = {
-        let mut rng = thread_rng();
-        (
-            TokenType::AccessToken.generate(&mut rng),
-            TokenType::RefreshToken.generate(&mut rng),
-        )
-    };
+    let access_token_str = TokenType::AccessToken.generate(&mut rng);
+    let refresh_token_str = TokenType::RefreshToken.generate(&mut rng);
 
-    let access_token = add_access_token(&mut txn, session, access_token_str.clone(), ttl).await?;
+    let access_token = add_access_token(
+        &mut txn,
+        &mut rng,
+        &clock,
+        session,
+        access_token_str.clone(),
+        ttl,
+    )
+    .await?;
 
-    let _refresh_token =
-        add_refresh_token(&mut txn, session, access_token, refresh_token_str.clone()).await?;
+    let _refresh_token = add_refresh_token(
+        &mut txn,
+        &mut rng,
+        &clock,
+        session,
+        access_token,
+        refresh_token_str.clone(),
+    )
+    .await?;
 
     let id_token = if session.scope.contains(&scope::OPENID) {
         let mut claims = HashMap::new();
-        let now = Utc::now();
+        let now = clock.now();
         claims::ISS.insert(&mut claims, url_builder.oidc_issuer().to_string())?;
         claims::SUB.insert(&mut claims, &browser_session.user.sub)?;
         claims::AUD.insert(&mut claims, client.client_id.clone())?;
@@ -346,7 +356,7 @@ async fn authorization_code_grant(
         let signer = key.params().signing_key_for_alg(&alg)?;
         let header = JsonWebSignatureHeader::new(alg)
             .with_kid(key.kid().context("key has no `kid` for some reason")?);
-        let id_token = Jwt::sign(header, claims, &signer)?;
+        let id_token = Jwt::sign_with_rng(&mut rng, header, claims, &signer)?;
 
         Some(id_token.as_str().to_owned())
     } else {
@@ -362,7 +372,7 @@ async fn authorization_code_grant(
         params = params.with_id_token(id_token);
     }
 
-    exchange_grant(&mut txn, authz_grant).await?;
+    exchange_grant(&mut txn, &clock, authz_grant).await?;
 
     txn.commit().await?;
 
@@ -374,6 +384,8 @@ async fn refresh_token_grant(
     client: &Client<PostgresqlBackend>,
     mut txn: Transaction<'_, Postgres>,
 ) -> Result<AccessTokenResponse, RouteError> {
+    let (clock, mut rng) = crate::rng_and_clock()?;
+
     let (refresh_token, session) =
         lookup_active_refresh_token(&mut txn, &grant.refresh_token).await?;
 
@@ -383,24 +395,33 @@ async fn refresh_token_grant(
     }
 
     let ttl = Duration::minutes(5);
-    let (access_token_str, refresh_token_str) = {
-        let mut rng = thread_rng();
-        (
-            TokenType::AccessToken.generate(&mut rng),
-            TokenType::RefreshToken.generate(&mut rng),
-        )
-    };
+    let access_token_str = TokenType::AccessToken.generate(&mut rng);
+    let refresh_token_str = TokenType::RefreshToken.generate(&mut rng);
 
-    let new_access_token =
-        add_access_token(&mut txn, &session, access_token_str.clone(), ttl).await?;
+    let new_access_token = add_access_token(
+        &mut txn,
+        &mut rng,
+        &clock,
+        &session,
+        access_token_str.clone(),
+        ttl,
+    )
+    .await?;
 
-    let new_refresh_token =
-        add_refresh_token(&mut txn, &session, new_access_token, refresh_token_str).await?;
+    let new_refresh_token = add_refresh_token(
+        &mut txn,
+        &mut rng,
+        &clock,
+        &session,
+        new_access_token,
+        refresh_token_str,
+    )
+    .await?;
 
-    consume_refresh_token(&mut txn, &refresh_token).await?;
+    consume_refresh_token(&mut txn, &clock, &refresh_token).await?;
 
     if let Some(access_token) = refresh_token.access_token {
-        revoke_access_token(&mut txn, access_token).await?;
+        revoke_access_token(&mut txn, &clock, access_token).await?;
     }
 
     let params = AccessTokenResponse::new(access_token_str)

crates/handlers/src/oauth2/userinfo.rs

@@ -54,6 +54,7 @@ pub async fn get(
     user_authorization: UserAuthorization,
 ) -> Result<Response, FancyError> {
     // TODO: error handling
+    let (_clock, mut rng) = crate::rng_and_clock()?;
     let mut conn = pool.acquire().await?;
 
     let session = user_authorization.protected(&mut conn).await?;
@@ -88,7 +89,7 @@ pub async fn get(
             user_info,
         };
 
-        let token = Jwt::sign(header, user_info, &signer)?;
+        let token = Jwt::sign_with_rng(&mut rng, header, user_info, &signer)?;
         Ok(JwtResponse(token).into_response())
     } else {
         Ok(Json(user_info).into_response())