From 59668387d1f482c4a7a4c5b203c3029bd097b188 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABlle=20Huisman?= <danielle@huisman.me>
Date: Wed, 15 Jan 2025 10:00:09 +0100
Subject: [PATCH] Shield: Move user logic to core

---
 packages/core/shield/src/shield.rs            | 51 +++++++++++++++++++
 .../integrations/shield-tower/src/service.rs  | 23 +--------
 2 files changed, 53 insertions(+), 21 deletions(-)

diff --git a/packages/core/shield/src/shield.rs b/packages/core/shield/src/shield.rs
index 5604aa9..db15b57 100644
--- a/packages/core/shield/src/shield.rs
+++ b/packages/core/shield/src/shield.rs
@@ -77,6 +77,21 @@ impl<U: User> Shield<U> {
         })
     }
 
+    pub async fn subprovider_by_id(
+        &self,
+        provider_id: &str,
+        subprovider_id: Option<&str>,
+    ) -> Result<Option<Box<dyn Subprovider>>, ShieldError> {
+        match self.provider_by_id(provider_id) {
+            Some(provider) => {
+                provider
+                    .subprovider_by_id(subprovider_id.expect("TODO"))
+                    .await
+            }
+            None => Ok(None),
+        }
+    }
+
     pub async fn sign_in(
         &self,
         request: SignInRequest,
@@ -145,6 +160,42 @@ impl<U: User> Shield<U> {
 
         Ok(response)
     }
+
+    pub async fn user(&self, session: &Session) -> Result<Option<U>, ShieldError> {
+        let authentication = {
+            let session_data = session.data();
+            let session_data = session_data
+                .lock()
+                .map_err(|err| SessionError::Lock(err.to_string()))?;
+
+            session_data.authentication.clone()
+        };
+
+        match authentication {
+            Some(authentication) => {
+                if self
+                    .subprovider_by_id(
+                        &authentication.provider_id,
+                        authentication.subprovider_id.as_deref(),
+                    )
+                    .await?
+                    .is_none()
+                {
+                    session.purge().await?;
+                    return Ok(None);
+                }
+
+                let user = self.storage().user_by_id(&authentication.user_id).await?;
+
+                if user.is_none() {
+                    session.purge().await?;
+                }
+
+                Ok(user)
+            }
+            None => Ok(None),
+        }
+    }
 }
 
 #[cfg(test)]
diff --git a/packages/integrations/shield-tower/src/service.rs b/packages/integrations/shield-tower/src/service.rs
index 7b8a684..1b9a9c9 100644
--- a/packages/integrations/shield-tower/src/service.rs
+++ b/packages/integrations/shield-tower/src/service.rs
@@ -75,30 +75,11 @@ where
                 };
             let shield_session = Session::new(session_storage);
 
-            let authenticated = match shield_session.data().lock() {
-                Ok(session) => session.authentication.clone(),
+            let user = match shield.user(&shield_session).await {
+                Ok(user) => user,
                 Err(_err) => return Ok(Self::internal_server_error()),
             };
 
-            let user = if let Some(authenticated) = authenticated {
-                // TODO: Verify provider and subprovider still exist.
-
-                match shield.storage().user_by_id(&authenticated.user_id).await {
-                    Ok(user) => {
-                        if user.is_none() {
-                            if let Err(_err) = shield_session.purge().await {
-                                return Ok(Self::internal_server_error());
-                            }
-                        }
-
-                        user
-                    }
-                    Err(_err) => return Ok(Self::internal_server_error()),
-                }
-            } else {
-                None
-            };
-
             debug!("{:?}", user.as_ref().map(|user| user.id()));
 
             req.extensions_mut().insert(shield);