refresh and cache certs in account creation flow

crates/service/web/storyteller_web/src/http_server/endpoints/users/create_account_from_google_sign_in_handler.rs

@@ -14,6 +14,7 @@ use log::{info, warn};
 use sqlx::MySqlPool;
 use utoipa::ToSchema;
 use actix_helpers::extractors::get_request_origin_uri::get_request_origin_uri;
+use errors::AnyhowResult;
 use google_sign_in::decode_and_verify_token_claims::decode_and_verify_token_claims;
 use google_sign_in::download_certs::download_cert_key_set;
 use google_sign_in::VerificationOptions;
@@ -30,6 +31,7 @@ use crate::http_server::session::http::http_user_session_manager::HttpUserSessio
 use crate::http_server::validations::is_reserved_username::is_reserved_username;
 use crate::http_server::validations::validate_passwords::validate_passwords;
 use crate::http_server::validations::validate_username::validate_username;
+use crate::state::certs::google_sign_in_cert::GoogleSignInCert;
 use crate::util::email_to_gravatar::email_to_gravatar;
 use crate::util::enroll_in_studio::enroll_in_studio;
 use crate::util::generate_random_username::generate_random_username;
@@ -133,21 +135,37 @@ pub async fn create_account_from_google_sign_in_handler(
   request: web::Json<GoogleCreateAccountRequest>,
   mysql_pool: web::Data<MySqlPool>,
   session_cookie_manager: web::Data<HttpUserSessionManager>,
+  google_sign_in_cert: web::Data<GoogleSignInCert>,
 ) -> Result<HttpResponse, GoogleCreateAccountErrorResponse>
 {
-  // TODO(bt,2024-09-20): Cache certs in shared server state
-  let keys = download_cert_key_set()
+  let keys = google_sign_in_cert.fetch_key_map(false)
       .await
       .map_err(|e| {
         warn!("error downloading google certs: {:?}", e);
         GoogleCreateAccountErrorResponse::server_error()
       })?;
 
-  let claims = decode_and_verify_token_claims(&keys, &request.google_credential, None)
-    .map_err(|e| {
-      warn!("error decoding google token claims: {:?}", e);
-      GoogleCreateAccountErrorResponse::bad_request()
-    })?;
+  let claims = match decode_and_verify_token_claims(&keys, &request.google_credential, None) {
+    Ok(claims) => claims,
+    Err(err) => {
+      warn!("error decoding google token claims (will retry certs): {:?}", err);
+
+      let keys = google_sign_in_cert.fetch_key_map(true) // NB: REFRESH
+          .await
+          .map_err(|e| {
+            warn!("error refreshing google certs: {:?}", e);
+            GoogleCreateAccountErrorResponse::server_error()
+          })?;
+
+      let claims = decode_and_verify_token_claims(&keys, &request.google_credential, None)
+          .map_err(|e| {
+            warn!("error decoding google token claims: {:?}", e);
+            GoogleCreateAccountErrorResponse::bad_request()
+          })?;
+
+      claims
+    },
+  };
 
   info!("Google JWT credential claims: email {:?}, verified: {}",
     claims.email(),

crates/service/web/storyteller_web/src/main.rs

@@ -84,6 +84,7 @@ use crate::http_server::session::http::http_user_session_manager::HttpUserSessio
 use crate::http_server::session::session_checker::SessionChecker;
 use crate::http_server::web_utils::handle_multipart_error::handle_multipart_error;
 use crate::http_server::web_utils::scoped_temp_dir_creator::ScopedTempDirCreator;
+use crate::state::certs::google_sign_in_cert::GoogleSignInCert;
 use crate::state::memory_cache::model_token_to_info_cache::ModelTokenToInfoCache;
 use crate::state::server_state::{DurableInMemoryCaches, EnvConfig, EphemeralInMemoryCaches, InMemoryCaches, ServerInfo, ServerState, StaticFeatureFlags, StripeSettings, TrollBans};
 use crate::threads::db_health_checker_thread::db_health_check_status::HealthCheckStatus;
@@ -457,6 +458,7 @@ async fn main() -> AnyhowResult<()> {
       ip_addresses: ip_address_troll_bans,
     },
     temp_dir_creator: ScopedTempDirCreator::for_directory(easyenv::get_env_pathbuf_or_default("TEMP_DIR", "/tmp")),
+    google_sign_in_cert: GoogleSignInCert::new(),
   };
 
   serve(server_state)
@@ -600,6 +602,7 @@ pub async fn serve(server_state: ServerState) -> AnyhowResult<()>
       .app_data(web::Data::new(server_state_arc.stripe.clone().config.clone()))
       .app_data(web::Data::new(server_state_arc.stripe.clone().client.clone()))
       .app_data(web::Data::new(server_state_arc.third_party_url_redirector))
+      .app_data(web::Data::new(server_state_arc.google_sign_in_cert.clone()))
       .app_data(web::Data::new(server_state_arc.email_sender.clone()))
       .app_data(web::Data::new(old_server_environment))
       .app_data(web::Data::new(new_server_environment))

crates/service/web/storyteller_web/src/state/certs/google_sign_in_cert.rs

@@ -0,0 +1,57 @@
+use crate::http_server::endpoints::users::create_account_from_google_sign_in_handler::GoogleCreateAccountErrorResponse;
+use errors::AnyhowResult;
+use google_sign_in::certs::key_map::KeyMap;
+use google_sign_in::download_certs::download_cert_key_set;
+use log::{error, info, warn};
+use std::sync::{Arc, LockResult, RwLock};
+use anyhow::anyhow;
+
+/// The Google Sign In cert needs to be refreshed periodically
+#[derive(Clone)]
+pub struct GoogleSignInCert {
+  /// Multiple recent JWK public keys from Google
+  key_map: Arc<RwLock<Option<KeyMap>>>
+}
+
+impl GoogleSignInCert {
+  pub fn new() -> Self {
+    Self {
+      key_map: Arc::new(RwLock::new(None))
+    }
+  }
+
+  /// Get the current key map
+  /// If available, this returns a copy so that we don't spend time in the critical section.
+  pub async fn fetch_key_map(&self, refresh: bool) -> AnyhowResult<KeyMap> {
+    if !refresh {
+      match self.key_map.read() {
+        Err(err) => {
+          error!("Error reading lock: {:?}", err);
+          return Err(anyhow!("Error reading lock: {:?}", err));
+        }
+        Ok(read) => {
+          if let Some(key_map) = &*read {
+            return Ok(key_map.clone());
+          }
+        }
+      }
+    }
+    info!("Fetching Google Sign In certs...");
+    let new_certs = download_cert_key_set().await?;
+    let _r = self.set_key_map(new_certs.clone());
+    Ok(new_certs)
+  }
+
+  fn set_key_map(&self, key_map: KeyMap) -> AnyhowResult<()> {
+    match self.key_map.write() {
+      Err(err) => {
+        error!("Error writing lock: {:?}", err);
+        Err(anyhow!("Error writing lock: {:?}", err))
+      }
+      Ok(mut write) => {
+        *write = Some(key_map);
+        Ok(())
+      }
+    }
+  }
+}

crates/service/web/storyteller_web/src/state/certs/mod.rs

@@ -0,0 +1 @@
+pub mod google_sign_in_cert;

crates/service/web/storyteller_web/src/state/mod.rs

@@ -1,3 +1,4 @@
 pub mod cached_queries;
+pub mod certs;
 pub mod memory_cache;
 pub mod server_state;

crates/service/web/storyteller_web/src/state/server_state.rs

@@ -33,6 +33,7 @@ use crate::http_server::session::http::http_user_session_manager::HttpUserSessio
 use crate::http_server::session::session_checker::SessionChecker;
 use crate::http_server::web_utils::redis_rate_limiter::RedisRateLimiter;
 use crate::http_server::web_utils::scoped_temp_dir_creator::ScopedTempDirCreator;
+use crate::state::certs::google_sign_in_cert::GoogleSignInCert;
 use crate::state::memory_cache::model_token_to_info_cache::ModelTokenToInfoCache;
 use crate::threads::db_health_checker_thread::db_health_check_status::HealthCheckStatus;
 use crate::util::encrypted_sort_id::SortKeyCrypto;
@@ -97,6 +98,8 @@ pub struct ServerState {
 
   pub caches: InMemoryCaches,
 
+  pub google_sign_in_cert: GoogleSignInCert,
+
   pub temp_dir_creator: ScopedTempDirCreator,
 }