From 3cfe6567f040f38b508a5b69a0a3c0701b18396a Mon Sep 17 00:00:00 2001
From: nazeh <ar.nazeh@gmail.com>
Date: Thu, 19 Dec 2024 16:10:31 +0300
Subject: [PATCH] feat(homeserver): replace session_id with the first 8 bytes
 of a key

---
 .../src/core/database/tables/sessions.rs      | 16 ++++++++-----
 pubky-homeserver/src/core/mod.rs              |  9 ++++---
 pubky-homeserver/src/core/routes/auth.rs      |  9 ++++---
 pubky/src/native/cookies.rs                   | 24 ++++++++++++-------
 4 files changed, 34 insertions(+), 24 deletions(-)

diff --git a/pubky-homeserver/src/core/database/tables/sessions.rs b/pubky-homeserver/src/core/database/tables/sessions.rs
index e2ff58d..3ffe916 100644
--- a/pubky-homeserver/src/core/database/tables/sessions.rs
+++ b/pubky-homeserver/src/core/database/tables/sessions.rs
@@ -29,10 +29,11 @@ impl DB {
     pub fn get_session_bytes(
         &mut self,
         cookies: Cookies,
-        _public_key: &PublicKey,
+        public_key: &PublicKey,
     ) -> anyhow::Result<Option<Vec<u8>>> {
-        // TODO: support coookie for key in the path
-        if let Some(cookie) = cookies.get("session_id") {
+        if let Some(cookie) =
+            cookies.get(&public_key.to_string().chars().take(8).collect::<String>())
+        {
             let rtxn = self.env.read_txn()?;
 
             let session = self
@@ -52,10 +53,13 @@ impl DB {
     pub fn delete_session(
         &mut self,
         cookies: Cookies,
-        _public_key: &PublicKey,
+        public_key: &PublicKey,
     ) -> anyhow::Result<bool> {
-        // TODO: support coookie for key in the path
-        if let Some(cookie) = cookies.get("session_id") {
+        // TODO: Set expired cookie to delete the cookie on client side.
+
+        if let Some(cookie) =
+            cookies.get(&public_key.to_string().chars().take(8).collect::<String>())
+        {
             let mut wtxn = self.env.write_txn()?;
 
             let deleted = self.tables.sessions.delete(&mut wtxn, cookie.value())?;
diff --git a/pubky-homeserver/src/core/mod.rs b/pubky-homeserver/src/core/mod.rs
index c4cf6de..ce33985 100644
--- a/pubky-homeserver/src/core/mod.rs
+++ b/pubky-homeserver/src/core/mod.rs
@@ -103,11 +103,10 @@ impl HomeserverCore {
 
         wtxn.commit()?;
 
-        let mut cookie = if true {
-            Cookie::new("session_id", session_secret)
-        } else {
-            Cookie::new(public_key.to_string(), session_secret)
-        };
+        let mut cookie = Cookie::new(
+            public_key.to_string().chars().take(8).collect::<String>(),
+            session_secret,
+        );
 
         cookie.set_path("/");
 
diff --git a/pubky-homeserver/src/core/routes/auth.rs b/pubky-homeserver/src/core/routes/auth.rs
index 9b143ab..f1efdb2 100644
--- a/pubky-homeserver/src/core/routes/auth.rs
+++ b/pubky-homeserver/src/core/routes/auth.rs
@@ -96,11 +96,10 @@ pub async fn signin(
 
     wtxn.commit()?;
 
-    let mut cookie = if true {
-        Cookie::new("session_id", session_secret)
-    } else {
-        Cookie::new(public_key.to_string(), session_secret)
-    };
+    let mut cookie = Cookie::new(
+        public_key.to_string().chars().take(8).collect::<String>(),
+        session_secret,
+    );
 
     cookie.set_path("/");
 
diff --git a/pubky/src/native/cookies.rs b/pubky/src/native/cookies.rs
index 5a030cb..d4a2adb 100644
--- a/pubky/src/native/cookies.rs
+++ b/pubky/src/native/cookies.rs
@@ -12,11 +12,15 @@ pub struct CookieJar {
 impl CookieJar {
     pub(crate) fn store_session_after_signup(&self, response: &Response, pubky: &PublicKey) {
         for (header_name, header_value) in response.headers() {
-            if header_name == "set-cookie" && header_value.as_ref().starts_with(b"session_id=") {
+            let cookie_name = &pubky.to_string().chars().take(8).collect::<String>();
+
+            if header_name == "set-cookie"
+                && header_value.as_ref().starts_with(cookie_name.as_bytes())
+            {
                 if let Ok(Ok(cookie)) =
                     std::str::from_utf8(header_value.as_bytes()).map(cookie::Cookie::parse)
                 {
-                    if cookie.name() == "session_id" {
+                    if cookie.name() == cookie_name {
                         let domain = format!("_pubky.{pubky}");
                         tracing::debug!(?cookie, "Storing coookie after signup");
 
@@ -62,13 +66,17 @@ impl CookieStore for CookieJar {
             .collect::<Vec<_>>()
             .join("; ");
 
+        // TODO: should we return if empty or just push?
         if s.is_empty() {
-            return self
-                .pubky_sessions
-                .read()
-                .unwrap()
-                .get(url.host_str().unwrap())
-                .map(|secret| HeaderValue::try_from(format!("session_id={secret}")).unwrap());
+            let host = url.host_str().unwrap_or("");
+
+            if let Ok(public_key) = PublicKey::try_from(host) {
+                let cookie_name = public_key.to_string().chars().take(8).collect::<String>();
+
+                return self.pubky_sessions.read().unwrap().get(host).map(|secret| {
+                    HeaderValue::try_from(format!("{cookie_name}={secret}")).unwrap()
+                });
+            }
         }
 
         HeaderValue::from_maybe_shared(bytes::Bytes::from(s)).ok()