Update store

sugyan · Sep 11, 2024 · 8399132 · 8399132
1 parent 8fe20a4
commit 8399132
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 11 deletions.
diff --git a/atrium-oauth/oauth-client/src/oauth_client.rs b/atrium-oauth/oauth-client/src/oauth_client.rs
@@ -223,19 +223,25 @@ where
         }
     }
     pub async fn callback(&self, params: CallbackParams) -> Result<TokenSet> {
-        let Some(state) = params.state else {
+        let Some(state_key) = params.state else {
             return Err(Error::Callback("missing `state` parameter".into()));
         };
         let Some(state) = self
             .state_store
-            .get(&state)
+            .get(&state_key)
             .await
             .map_err(|e| Error::StateStore(Box::new(e)))?
         else {
             return Err(Error::Callback(format!(
-                "unknown authorization state: {state}"
+                "unknown authorization state: {state_key}"
             )));
         };
+        // Prevent any kind of replay
+        self.state_store
+            .del(&state_key)
+            .await
+            .map_err(|e| Error::StateStore(Box::new(e)))?;
+
         let metadata = self
             .resolver
             .get_authorization_server_metadata(&state.iss)

diff --git a/atrium-oauth/oauth-client/src/store.rs b/atrium-oauth/oauth-client/src/store.rs
@@ -5,7 +5,8 @@ use async_trait::async_trait;
 use std::error::Error;
 use std::hash::Hash;
 
-#[async_trait]
+#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
+#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
 pub trait SimpleStore<K, V>
 where
     K: Eq + Hash,
@@ -14,7 +15,7 @@ where
     type Error: Error + Send + Sync + 'static;
 
     async fn get(&self, key: &K) -> Result<Option<V>, Self::Error>;
-    async fn set(&self, key: K, value: V) -> Result<Option<V>, Self::Error>;
-    async fn del(&self, key: &K) -> Result<Option<V>, Self::Error>;
+    async fn set(&self, key: K, value: V) -> Result<(), Self::Error>;
+    async fn del(&self, key: &K) -> Result<(), Self::Error>;
     async fn clear(&self) -> Result<(), Self::Error>;
 }
diff --git a/atrium-oauth/oauth-client/src/store/memory.rs b/atrium-oauth/oauth-client/src/store/memory.rs
@@ -23,7 +23,8 @@ impl<K, V> Default for MemorySimpleStore<K, V> {
     }
 }
 
-#[async_trait]
+#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
+#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
 impl<K, V> SimpleStore<K, V> for MemorySimpleStore<K, V>
 where
     K: Debug + Eq + Hash + Send + Sync + 'static,
@@ -34,11 +35,13 @@ where
     async fn get(&self, key: &K) -> Result<Option<V>, Self::Error> {
         Ok(self.store.lock().unwrap().get(key).cloned())
     }
-    async fn set(&self, key: K, value: V) -> Result<Option<V>, Self::Error> {
-        Ok(self.store.lock().unwrap().insert(key, value))
+    async fn set(&self, key: K, value: V) -> Result<(), Self::Error> {
+        self.store.lock().unwrap().insert(key, value);
+        Ok(())
     }
-    async fn del(&self, key: &K) -> Result<Option<V>, Self::Error> {
-        Ok(self.store.lock().unwrap().remove(key))
+    async fn del(&self, key: &K) -> Result<(), Self::Error> {
+        self.store.lock().unwrap().remove(key);
+        Ok(())
     }
     async fn clear(&self) -> Result<(), Self::Error> {
         self.store.lock().unwrap().clear();