From 235e2789aa3d618cddfb61c297ccb73bca9dd9d1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABlle=20Huisman?= <danielle@huisman.me>
Date: Wed, 25 Dec 2024 15:32:40 +0100
Subject: [PATCH] Shield Leptos/Tower: Fix session

---
 packages/integrations/shield-leptos/src/context.rs  | 13 +++++++++++++
 packages/integrations/shield-leptos/src/lib.rs      |  1 +
 .../shield-leptos/src/routes/sign_in.rs             | 13 +++++--------
 packages/integrations/shield-tower/src/service.rs   |  7 ++++---
 4 files changed, 23 insertions(+), 11 deletions(-)
 create mode 100644 packages/integrations/shield-leptos/src/context.rs

diff --git a/packages/integrations/shield-leptos/src/context.rs b/packages/integrations/shield-leptos/src/context.rs
new file mode 100644
index 0000000..bcf9671
--- /dev/null
+++ b/packages/integrations/shield-leptos/src/context.rs
@@ -0,0 +1,13 @@
+use std::sync::Arc;
+
+use leptos::prelude::expect_context;
+use shield::{ServerIntegration, Shield};
+
+pub fn expect_server_integration() -> Arc<dyn ServerIntegration> {
+    expect_context::<Arc<dyn ServerIntegration>>()
+}
+
+pub async fn expect_shield() -> Shield {
+    let server_integration = expect_server_integration();
+    server_integration.extract_shield().await
+}
diff --git a/packages/integrations/shield-leptos/src/lib.rs b/packages/integrations/shield-leptos/src/lib.rs
index 6a664ab..e206dc0 100644
--- a/packages/integrations/shield-leptos/src/lib.rs
+++ b/packages/integrations/shield-leptos/src/lib.rs
@@ -1 +1,2 @@
+pub mod context;
 pub mod routes;
diff --git a/packages/integrations/shield-leptos/src/routes/sign_in.rs b/packages/integrations/shield-leptos/src/routes/sign_in.rs
index bb1fed1..8d75b87 100644
--- a/packages/integrations/shield-leptos/src/routes/sign_in.rs
+++ b/packages/integrations/shield-leptos/src/routes/sign_in.rs
@@ -3,12 +3,9 @@ use shield::SubproviderVisualisation;
 
 #[server]
 pub async fn subproviders() -> Result<Vec<SubproviderVisualisation>, ServerFnError> {
-    use std::sync::Arc;
+    use crate::context::expect_shield;
 
-    use shield::ServerIntegration;
-
-    let server_integration = expect_context::<Arc<dyn ServerIntegration>>();
-    let shield = server_integration.extract_shield().await;
+    let shield = expect_shield().await;
 
     shield
         .subprovider_visualisations()
@@ -21,11 +18,11 @@ pub async fn sign_in(
     provider_id: String,
     subprovider_id: Option<String>,
 ) -> Result<(), ServerFnError> {
-    use std::sync::Arc;
+    use shield::{Response, ShieldError, SignInRequest};
 
-    use shield::{Response, ServerIntegration, ShieldError, SignInRequest};
+    use crate::context::expect_server_integration;
 
-    let server_integration = expect_context::<Arc<dyn ServerIntegration>>();
+    let server_integration = expect_server_integration();
     let shield = server_integration.extract_shield().await;
     let session = server_integration.extract_session().await;
 
diff --git a/packages/integrations/shield-tower/src/service.rs b/packages/integrations/shield-tower/src/service.rs
index 3643706..3b448bf 100644
--- a/packages/integrations/shield-tower/src/service.rs
+++ b/packages/integrations/shield-tower/src/service.rs
@@ -5,7 +5,7 @@ use std::{
 };
 
 use http::{Request, Response};
-use shield::Shield;
+use shield::{Session, Shield};
 use tower_service::Service;
 
 use crate::session::TowerSession;
@@ -66,10 +66,11 @@ where
                 }
             };
 
-            let shield_session = match TowerSession::load(session.clone(), session_key).await {
-                Ok(shield_session) => shield_session,
+            let session_storage = match TowerSession::load(session.clone(), session_key).await {
+                Ok(session_storage) => session_storage,
                 Err(_err) => return Ok(Self::internal_server_error()),
             };
+            let shield_session = Session::new(session_storage);
 
             req.extensions_mut().insert(shield);
             req.extensions_mut().insert(shield_session);