From 29c9c4d76224773ea4dc8f48552da18cb332fa95 Mon Sep 17 00:00:00 2001 From: Jimmie Lovell Date: Fri, 7 Feb 2025 01:40:05 +0300 Subject: [PATCH 1/2] upgrade lib version --- Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index 9133ca1..8875f1d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "ruts" description = "A middleware for tower sessions" -version = "0.5.1" +version = "0.5.2" edition = "2021" rust-version = "1.75.0" authors = ["Jimmie Lovell "] From 1cbe953b67598415b445bae9f0458926627f1626 Mon Sep 17 00:00:00 2001 From: Jimmie Lovell Date: Sat, 8 Feb 2025 00:38:16 +0300 Subject: [PATCH 2/2] added prepare_regenerate for atomic session regeneration --- CHANGELOG.md | 6 +++ Cargo.toml | 2 +- README.md | 16 ++++-- src/extract/mod.rs | 6 +-- src/lib.rs | 5 +- src/session/mod.rs | 113 ++++++++++++++++++++++++++++++++------- src/store/memory.rs | 95 ++++++++++++++++++++++++++++++++ src/store/redis/lua.rs | 62 +++++++++++++++++++++ src/store/redis/mod.rs | 63 +++++++++++++++++++--- src/store/store_trait.rs | 35 ++++++++++++ tests/common/mod.rs | 6 +-- tests/core.rs | 98 +++++++++++++++++++++++++++++++++ tests/redis.rs | 72 ++++++++++++++++++++++++- 13 files changed, 538 insertions(+), 41 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 20fb1b9..d9c1cd4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,11 @@ # Changelog +## [0.5.3] - 2024-02-08 + +### Added +- New `prepare_regenerate()` method for atomic session ID regeneration with updates +- Support for atomic (insert/update with ID regeneration) operations in both Redis and Memory stores + ## [0.5.2] - 2024-02-07 ### Added diff --git a/Cargo.toml b/Cargo.toml index 8875f1d..34e6094 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "ruts" description = "A middleware for tower sessions" -version = "0.5.2" +version = "0.5.3" edition = "2021" rust-version = "1.75.0" authors = ["Jimmie Lovell "] diff --git a/README.md b/README.md index 35611cc..e2b57c4 100644 --- a/README.md +++ b/README.md @@ -23,7 +23,7 @@ Add the following to your `Cargo.toml`: ```toml [dependencies] -ruts = "0.4.2" +ruts = "0.5.3" ``` ## Quick Start @@ -87,13 +87,21 @@ async fn handler(session: Session>) -> String { ```rust // Get session data -let value = session.get::("key").await?; +let value: ValueType = session.get("key").await?; // Insert new data -session.insert::("key", &value).await?; +session.insert("key", &value, optional_field_expiration).await?; + +// Prepare a new session ID for the next insert +let new_id = session.prepare_regenerate(); +session.insert("key", &value, optional_field_expiration).await?; // Update existing data -session.update::("key", &new_value).await?; +session.update("key", &new_value, optional_field_expiration).await?; + +// Prepare a new session ID for the next update +let new_id = session.prepare_regenerate(); +session.update("key", &value, optional_field_expiration).await?; // Remove data session.remove("key").await?; diff --git a/src/extract/mod.rs b/src/extract/mod.rs index cf46e1c..2e34093 100644 --- a/src/extract/mod.rs +++ b/src/extract/mod.rs @@ -21,21 +21,21 @@ where tracing::error!("session layer not found in the request extensions"); ( StatusCode::INTERNAL_SERVER_ERROR, - "session not found in the request", + "Session not found in the request", ) })?; // Cookies are only used if the SessionLayer has a cookie_options set. let cookie_name = session_inner.cookie_name.ok_or_else(|| { tracing::error!("missing cookie options"); - (StatusCode::INTERNAL_SERVER_ERROR, "missing cookie options") + (StatusCode::INTERNAL_SERVER_ERROR, "Missing cookie options") })?; let cookies_ext = parts.extensions.get::().ok_or_else(|| { tracing::error!("cookies not found in the request extensions"); ( StatusCode::INTERNAL_SERVER_ERROR, - "cookies not found in the request", + "Cookies not found in the request", ) })?; diff --git a/src/lib.rs b/src/lib.rs index e1b6685..2fd30cc 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,5 +7,8 @@ pub use session::*; mod extract; mod service; - pub use service::*; + +// Reexport external crates +pub use cookie; +pub use tower_cookies; diff --git a/src/session/mod.rs b/src/session/mod.rs index df420fe..4b81244 100644 --- a/src/session/mod.rs +++ b/src/session/mod.rs @@ -165,16 +165,35 @@ where where T: Send + Sync + Serialize, { - let id = self.inner.get_or_set_id(); - let inserted = self - .inner - .store - .insert(&id, field, value, self.max_age(), field_expire) - .await - .map_err(|err| { - tracing::error!(err = %err, "failed to insert field-value to session store"); - err - })?; + let current_id = self.inner.get_or_set_id(); + let pending_id = self.inner.take_pending_id(); + + let inserted = match pending_id { + Some(new_id) => { + let inserted = self.inner + .store + .insert_with_rename(¤t_id, &new_id, field, value, self.max_age(), field_expire) + .await + .map_err(|err| { + tracing::error!(err = %err, "failed to insert field-value with rename to session store"); + err + })?; + if inserted { + *self.inner.id.write() = Some(new_id); + } + inserted + } + None => { + self.inner + .store + .insert(¤t_id, field, value, self.max_age(), field_expire) + .await + .map_err(|err| { + tracing::error!(err = %err, "failed to insert field-value to session store"); + err + })? + } + }; if inserted { self.inner.set_changed(); @@ -232,16 +251,37 @@ where where T: Send + Sync + Serialize, { - let id = self.inner.get_or_set_id(); - let updated = self - .inner - .store - .update(&id, field, value, self.max_age(), field_expire) - .await - .map_err(|err| { - tracing::error!(err = %err, "failed to update field in session store"); - err - })?; + let current_id = self.inner.get_or_set_id(); + let pending_id = self.inner.take_pending_id(); + + let updated = match pending_id { + Some(new_id) => { + let updated = self.inner + .store + .update_with_rename(¤t_id, &new_id, field, value, self.max_age(), field_expire) + .await + .map_err(|err| { + tracing::error!(err = %err, "failed to update field-value with rename in session store"); + err + })?; + + if updated { + *self.inner.id.write() = Some(new_id); + } + + updated + } + None => { + self.inner + .store + .update(¤t_id, field, value, self.max_age(), field_expire) + .await + .map_err(|err| { + tracing::error!(err = %err, "failed to update field in session store"); + err + })? + } + }; if updated { self.inner.set_changed(); @@ -415,6 +455,29 @@ where Ok(None) } + /// Prepares a new session ID to be used in the next store operation. + /// The new ID will be used to rename the current session when the next + /// insert or update operation is performed. + /// + /// # Example + /// + /// ```rust + /// use ruts::Session; + /// use fred::clients::Client; + /// use ruts::store::redis::RedisStore; + /// + /// async fn some_handler_could_be_axum(session: Session>) { + /// let new_id = session.prepare_regenerate(); + /// // The next update/insert operation will use this new ID + /// session.update("field", &"value", None).await.unwrap(); + /// } + /// ``` + pub fn prepare_regenerate(&self) -> Id { + let new_id = Id::default(); + self.inner.set_pending_id(Some(new_id)); + new_id + } + /// Returns the session ID, if it exists. pub fn id(&self) -> Option { self.inner.get_id() @@ -433,6 +496,7 @@ const DEFAULT_COOKIE_MAX_AGE: i64 = 10 * 60; pub struct Inner { pub state: AtomicU8, pub id: RwLock>, + pub pending_id: RwLock>, pub cookie_max_age: AtomicI64, pub cookie_name: Option<&'static str>, pub cookies: OnceLock, @@ -448,6 +512,7 @@ impl Inner { Self { state: AtomicU8::new(0), id: RwLock::new(None), + pending_id: RwLock::new(None), cookie_max_age: AtomicI64::new(cookie_max_age.unwrap_or(DEFAULT_COOKIE_MAX_AGE)), cookie_name, cookies: OnceLock::new(), @@ -475,6 +540,14 @@ impl Inner { *self.id.write() = id; } + pub fn set_pending_id(&self, id: Option) { + *self.pending_id.write() = id; + } + + pub fn take_pending_id(&self) -> Option { + self.pending_id.write().take() + } + pub fn set_changed(&self) { self.state .fetch_or(SESSION_STATE_CHANGED, Ordering::Relaxed); diff --git a/src/store/memory.rs b/src/store/memory.rs index a923363..fd636bc 100644 --- a/src/store/memory.rs +++ b/src/store/memory.rs @@ -166,6 +166,101 @@ impl SessionStore for MemoryStore { Ok(true) } + async fn insert_with_rename( + &self, + old_session_id: &Id, + new_session_id: &Id, + field: &str, + value: &T, + key_seconds: i64, + _field_seconds: Option, + ) -> Result + where + T: Send + Sync + Serialize, + { + self.cleanup_expired(); + + let mut data = self.data.write(); + + // Check if old session exists and new session doesn't + if !data.contains_key(&old_session_id.to_string()) || data.contains_key(&new_session_id.to_string()) { + return Ok(false); + } + + // Get the fields map, return false if field exists + let fields = data.get_mut(&old_session_id.to_string()).unwrap(); + if fields.contains_key(field) { + return Ok(false); + } + + // Calculate expiration + let expires_at = if key_seconds > 0 { + Some(Instant::now() + Duration::from_secs(key_seconds as u64)) + } else { + None + }; + + // Insert the new field + fields.insert( + field.to_string(), + StoredValue { + data: serialize_value(value)?, + expires_at, + }, + ); + + // Move the map to the new session ID + let fields = data.remove(&old_session_id.to_string()).unwrap(); + data.insert(new_session_id.to_string(), fields); + + Ok(true) + } + + async fn update_with_rename( + &self, + old_session_id: &Id, + new_session_id: &Id, + field: &str, + value: &T, + key_seconds: i64, + _field_seconds: Option, + ) -> Result + where + T: Send + Sync + Serialize, + { + self.cleanup_expired(); + + let mut data = self.data.write(); + + // Check if old session exists and new session doesn't + if !data.contains_key(&old_session_id.to_string()) || data.contains_key(&new_session_id.to_string()) { + return Ok(false); + } + + // Calculate expiration + let expires_at = if key_seconds > 0 { + Some(Instant::now() + Duration::from_secs(key_seconds as u64)) + } else { + None + }; + + // Update the field + let fields = data.get_mut(&old_session_id.to_string()).unwrap(); + fields.insert( + field.to_string(), + StoredValue { + data: serialize_value(value)?, + expires_at, + }, + ); + + // Move the map to the new session ID + let fields = data.remove(&old_session_id.to_string()).unwrap(); + data.insert(new_session_id.to_string(), fields); + + Ok(true) + } + async fn rename_session_id( &self, old_session_id: &Id, diff --git a/src/store/redis/lua.rs b/src/store/redis/lua.rs index 682e95b..c4e3a32 100644 --- a/src/store/redis/lua.rs +++ b/src/store/redis/lua.rs @@ -2,6 +2,9 @@ use tokio::sync::OnceCell; pub(crate) static INSERT_SCRIPT_HASH: OnceCell = OnceCell::const_new(); pub(crate) static UPDATE_SCRIPT_HASH: OnceCell = OnceCell::const_new(); +pub(crate) static INSERT_WITH_RENAME_SCRIPT_HASH: OnceCell = OnceCell::const_new(); +pub(crate) static UPDATE_WITH_RENAME_SCRIPT_HASH: OnceCell = OnceCell::const_new(); + pub(crate) static RENAME_SCRIPT_HASH: OnceCell = OnceCell::const_new(); pub(crate) static INSERT_SCRIPT: &str = r#" @@ -38,6 +41,65 @@ pub(crate) static UPDATE_SCRIPT: &str = r#" return updated "#; +pub(crate) static INSERT_WITH_RENAME_SCRIPT: &str = r#" + local old_key = KEYS[1] + local new_key = KEYS[2] + local field = ARGV[1] + local value = ARGV[2] + local key_seconds = tonumber(ARGV[3]) + local field_seconds = ARGV[4] + + local exists = redis.call('EXISTS', old_key) + if exists == 0 then + return 0 + end + + local new_exists = redis.call('EXISTS', new_key) + if new_exists == 1 then + return 0 + end + + local inserted = redis.call('HSETNX', old_key, field, value) + if inserted == 1 then + redis.call('RENAMENX', old_key, new_key) + redis.call('EXPIRE', new_key, key_seconds) + if field_seconds ~= '' then + redis.call('HEXPIRE', new_key, tonumber(field_seconds), 'FIELDS', 1, field) + end + end + return inserted +"#; + +pub(crate) static UPDATE_WITH_RENAME_SCRIPT: &str = r#" + local old_key = KEYS[1] + local new_key = KEYS[2] + local field = ARGV[1] + local value = ARGV[2] + local key_seconds = tonumber(ARGV[3]) + local field_seconds = ARGV[4] + + local exists = redis.call('EXISTS', old_key) + if exists == 0 then + return 0 + end + + local new_exists = redis.call('EXISTS', new_key) + if new_exists == 1 then + return 0 + end + + -- Update the field + local updated = redis.call('HSET', old_key, field, value) + if updated == 1 then + redis.call('RENAMENX', old_key, new_key) + redis.call('EXPIRE', new_key, key_seconds) + if field_seconds ~= '' then + redis.call('HEXPIRE', new_key, tonumber(field_seconds), 'FIELDS', 1, field) + end + end + return updated +"#; + pub(crate) const RENAME_SCRIPT: &str = r#" local old_key = KEYS[1] local new_key = KEYS[2] diff --git a/src/store/redis/mod.rs b/src/store/redis/mod.rs index fd4ac20..e301fb4 100644 --- a/src/store/redis/mod.rs +++ b/src/store/redis/mod.rs @@ -1,9 +1,6 @@ mod lua; -use crate::store::redis::lua::{ - INSERT_SCRIPT, INSERT_SCRIPT_HASH, RENAME_SCRIPT, RENAME_SCRIPT_HASH, UPDATE_SCRIPT, - UPDATE_SCRIPT_HASH, -}; +use crate::store::redis::lua::{INSERT_SCRIPT, INSERT_SCRIPT_HASH, INSERT_WITH_RENAME_SCRIPT, INSERT_WITH_RENAME_SCRIPT_HASH, RENAME_SCRIPT, RENAME_SCRIPT_HASH, UPDATE_SCRIPT, UPDATE_SCRIPT_HASH, UPDATE_WITH_RENAME_SCRIPT, UPDATE_WITH_RENAME_SCRIPT_HASH}; use crate::store::{deserialize_value, serialize_value, Error, SessionStore}; use crate::Id; use fred::clients::Pool; @@ -91,7 +88,7 @@ where { insert_update( Arc::clone(&self.client), - session_id, + vec![session_id], field, value, key_seconds, @@ -115,7 +112,7 @@ where { insert_update( Arc::clone(&self.client), - session_id, + vec![session_id], field, value, key_seconds, @@ -126,6 +123,56 @@ where .await } + async fn insert_with_rename( + &self, + old_session_id: &Id, + new_session_id: &Id, + field: &str, + value: &T, + key_seconds: i64, + field_seconds: Option, + ) -> Result + where + T: Send + Sync + Serialize, + { + insert_update( + Arc::clone(&self.client), + vec![old_session_id, new_session_id], + field, + value, + key_seconds, + field_seconds, + &INSERT_WITH_RENAME_SCRIPT_HASH, + INSERT_WITH_RENAME_SCRIPT, + ) + .await + } + + async fn update_with_rename( + &self, + old_session_id: &Id, + new_session_id: &Id, + field: &str, + value: &T, + key_seconds: i64, + field_seconds: Option, + ) -> Result + where + T: Send + Sync + Serialize, + { + insert_update( + Arc::clone(&self.client), + vec![old_session_id, new_session_id], + field, + value, + key_seconds, + field_seconds, + &UPDATE_WITH_RENAME_SCRIPT_HASH, + UPDATE_WITH_RENAME_SCRIPT, + ) + .await + } + async fn rename_session_id( &self, old_session_id: &Id, @@ -171,7 +218,7 @@ where async fn insert_update( client: Arc, - session_id: &Id, + session_ids: Vec<&Id>, field: &str, value: &T, key_seconds: i64, @@ -202,7 +249,7 @@ where let done: bool = client .evalsha( hash, - vec![session_id], + session_ids, vec![ field.as_bytes(), &serialized_value, diff --git a/src/store/store_trait.rs b/src/store/store_trait.rs index df9a915..b7ab24b 100644 --- a/src/store/store_trait.rs +++ b/src/store/store_trait.rs @@ -69,6 +69,41 @@ pub trait SessionStore: Clone + Send + Sync + 'static { where T: Send + Sync + Serialize; + /// Sets a `field` stored at `session_id` to its provided `value` and renames + /// the session ID from `old_session_id` to `new_session_id`, + /// only if the `field` does not exist. + /// + /// Returns `true` if the `field` was inserted, otherwise, `false` if the `field` already exists. + fn insert_with_rename( + &self, + old_session_id: &Id, + new_session_id: &Id, + field: &str, + value: &T, + key_seconds: i64, + field_seconds: Option, + ) -> impl Future> + Send + where + T: Send + Sync + Serialize; + + /// Updates a `field` stored at `session_id` to the new `value` and renames + /// the session ID from `old_session_id` to `new_session_id`. + /// + /// If the `field` does not exist, it is set to the corresponding `value`. + /// + /// Returns `true` if the `field` was updated, `false` if the `value` has not changed. + fn update_with_rename( + &self, + old_session_id: &Id, + new_session_id: &Id, + field: &str, + value: &T, + key_seconds: i64, + field_seconds: Option, + ) -> impl Future> + Send + where + T: Send + Sync + Serialize; + /// Renames the `old_session_id` to `new_session_id` if the `old_session_id` exists. /// /// Returns an error when `old_session_id` does not exist. diff --git a/tests/common/mod.rs b/tests/common/mod.rs index 758eb84..b4564b8 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -2,7 +2,7 @@ use ruts::CookieOptions; use serde::{Deserialize, Serialize}; #[derive(Debug, Serialize, Deserialize, PartialEq, Clone)] -pub struct TestUser { +pub(crate) struct TestUser { pub id: i64, pub name: String, } @@ -14,7 +14,7 @@ pub struct TestSession { } #[derive(Debug, Serialize, Deserialize, PartialEq, Clone)] -pub struct TestPreferences { +pub(crate) struct TestPreferences { pub theme: String, pub language: String, } @@ -38,6 +38,6 @@ pub fn build_cookie_options() -> CookieOptions { .http_only(true) .same_site(cookie::SameSite::Lax) .secure(true) - .max_age(3600) + .max_age(15) .path("/") } diff --git a/tests/core.rs b/tests/core.rs index 0ff61d2..63b4141 100644 --- a/tests/core.rs +++ b/tests/core.rs @@ -56,6 +56,104 @@ mod tests { assert!(retrieved.is_none()); } + #[tokio::test] + async fn test_prepare_regenerate_with_update() { + let store = Arc::new(MemoryStore::new()); + let inner = create_inner(store.clone(), Some("test_sess"), Some(3600)); + let session = Session::new(inner); + let test_data = create_test_session(); + + // Insert initial data + session.insert("test", &test_data, None).await.unwrap(); + let original_id = session.id().unwrap(); + + // Prepare new ID and update + let prepared_id = session.prepare_regenerate(); + let mut updated_data = test_data.clone(); + updated_data.user.name = "Updated User".to_string(); + let updated = session.update("test", &updated_data, None).await.unwrap(); + assert!(updated); + + // Verify ID changed and data updated + let current_id = session.id().unwrap(); + assert_eq!(current_id, prepared_id); + assert_ne!(current_id, original_id); + + let retrieved: Option = session.get("test").await.unwrap(); + assert_eq!(retrieved.unwrap(), updated_data); + + // Verify old session is gone by directly checking the store + let result: Result, _> = store.get(&original_id, "test").await; + assert!(result.unwrap().is_none()); + } + + #[tokio::test] + async fn test_prepare_regenerate_with_insert() { + let store = Arc::new(MemoryStore::new()); + let inner = create_inner(store.clone(), Some("test_sess"), Some(3600)); + let session = Session::new(inner); + let test_data = create_test_session(); + + // Insert initial data + session.insert("test1", &test_data, None).await.unwrap(); + let original_id = session.id().unwrap(); + + // Prepare new ID and insert new field + let prepared_id = session.prepare_regenerate(); + let mut new_data = test_data.clone(); + new_data.user.name = "New User".to_string(); + let inserted = session.insert("test2", &new_data, None).await.unwrap(); + assert!(inserted); + + // Verify ID changed and both fields exist + let current_id = session.id().unwrap(); + assert_eq!(current_id, prepared_id); + assert_ne!(current_id, original_id); + + let retrieved1: Option = session.get("test1").await.unwrap(); + let retrieved2: Option = session.get("test2").await.unwrap(); + assert_eq!(retrieved1.unwrap(), test_data); + assert_eq!(retrieved2.unwrap(), new_data); + + // Verify old session is gone by directly checking the store + let result: Result, _> = store.get(&original_id, "test1").await; + assert!(result.unwrap().is_none()); + } + + #[tokio::test] + async fn test_multiple_prepare_regenerate() { + let store = Arc::new(MemoryStore::new()); + let inner = create_inner(store.clone(), Some("test_sess"), Some(3600)); + let session = Session::new(inner); + let test_data = create_test_session(); + + // Insert initial data + session.insert("test", &test_data, None).await.unwrap(); + let original_id = session.id().unwrap(); + + // First prepare_regenerate + let first_prepared_id = session.prepare_regenerate(); + + // Second prepare_regenerate before any operation + let second_prepared_id = session.prepare_regenerate(); + assert_ne!(first_prepared_id, second_prepared_id); + + // Update - should use the last prepared ID + let mut updated_data = test_data.clone(); + updated_data.user.name = "Updated User".to_string(); + session.update("test", &updated_data, None).await.unwrap(); + + // Verify the last prepared ID was used + let current_id = session.id().unwrap(); + assert_eq!(current_id, second_prepared_id); + assert_ne!(current_id, first_prepared_id); + assert_ne!(current_id, original_id); + + // Verify original session is gone by directly checking the store + let result: Result, _> = store.get(&original_id, "test").await; + assert!(result.unwrap().is_none()); + } + #[tokio::test] async fn test_session_regeneration() { let store = Arc::new(MemoryStore::new()); diff --git a/tests/redis.rs b/tests/redis.rs index 50d6ed1..5d444b3 100644 --- a/tests/redis.rs +++ b/tests/redis.rs @@ -29,7 +29,7 @@ mod tests { async fn insert_handler(session: Session>) -> Result { let test_data = create_test_session(); session - .insert("user", &test_data, Some(30)) + .insert("user", &test_data, Some(5)) .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; Ok("Success".to_string()) @@ -242,4 +242,74 @@ mod tests { assert_eq!(body_str, "Test User"); } } + + async fn prepare_and_update_handler( + session: Session>, + ) -> Result { + let test_data = create_test_session(); + + // Insert initial data + session + .insert("user", &test_data, Some(30)) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + + // Prepare new ID and update + session.prepare_regenerate(); + + let mut updated_data = test_data; + updated_data.user.name = "Updated User".to_string(); + session + .update("user", &updated_data, Some(30)) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + + Ok("Success".to_string()) + } + + #[tokio::test] + async fn test_prepare_regenerate_flow() { + let store = setup_redis().await; + let app = Router::new() + .route("/prepare_update", get(prepare_and_update_handler)) + .route("/get", get(get_handler)) + .layer(SessionLayer::new(store).with_cookie_options(build_cookie_options())) + .layer(CookieManagerLayer::new()); + + // Create and update session with new ID + let response = app + .clone() + .oneshot(Request::builder().uri("/prepare_update").body(Body::empty()).unwrap()) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + + let new_cookie = response + .headers() + .get(SET_COOKIE) + .expect("Set-Cookie header should be present") + .to_str() + .unwrap() + .to_string(); + + // Verify updated data with new session ID + let response = app + .clone() + .oneshot( + Request::builder() + .uri("/get") + .header(COOKIE, &new_cookie) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + let body = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap(); + let body_str = String::from_utf8(body.to_vec()).unwrap(); + assert_eq!(body_str, "Updated User"); + } }