fix(homeserver): always get sessions from Cookies instead of manually from headers

pubky-homeserver/src/core/layers/authz.rs

@@ -1,4 +1,4 @@
-use axum::http::{header, HeaderMap, Method};
+use axum::http::Method;
 use axum::response::IntoResponse;
 use axum::{
     body::Body,
@@ -79,14 +79,10 @@ where
                 }
             };
 
+            let cookies = req.extensions().get::<Cookies>();
+
             // Authorize the request
-            if let Err(e) = authorize(
-                &state,
-                req.method(),
-                req.headers(),
-                pubky.public_key(),
-                path,
-            ) {
+            if let Err(e) = authorize(&state, req.method(), cookies, pubky.public_key(), path) {
                 return Ok(e.into_response());
             }
 
@@ -100,7 +96,7 @@ where
 fn authorize(
     state: &AppState,
     method: &Method,
-    headers: &HeaderMap,
+    cookies: Option<&Cookies>,
     public_key: &PublicKey,
     path: &str,
 ) -> Result<()> {
@@ -118,45 +114,34 @@ fn authorize(
         ));
     }
 
-    let session_secret = session_secret_from_headers(headers, public_key)
-        .ok_or(Error::with_status(StatusCode::UNAUTHORIZED))?;
+    if let Some(cookies) = cookies {
+        let session_secret = session_secret_from_cookies(cookies, public_key)
+            .ok_or(Error::with_status(StatusCode::UNAUTHORIZED))?;
 
-    let session = state
-        .db
-        .get_session(&session_secret)?
-        .ok_or(Error::with_status(StatusCode::UNAUTHORIZED))?;
+        let session = state
+            .db
+            .get_session(&session_secret)?
+            .ok_or(Error::with_status(StatusCode::UNAUTHORIZED))?;
 
-    if session.pubky() == public_key
-        && session.capabilities().iter().any(|cap| {
-            path.starts_with(&cap.scope)
-                && cap
-                    .actions
-                    .contains(&pubky_common::capabilities::Action::Write)
-        })
-    {
-        return Ok(());
+        if session.pubky() == public_key
+            && session.capabilities().iter().any(|cap| {
+                path.starts_with(&cap.scope)
+                    && cap
+                        .actions
+                        .contains(&pubky_common::capabilities::Action::Write)
+            })
+        {
+            return Ok(());
+        }
+
+        return Err(Error::with_status(StatusCode::FORBIDDEN));
     }
 
-    Err(Error::with_status(StatusCode::FORBIDDEN))
+    Err(Error::with_status(StatusCode::UNAUTHORIZED))
 }
 
-pub fn session_secret_from_cookies(cookies: Cookies, public_key: &PublicKey) -> Option<String> {
+pub fn session_secret_from_cookies(cookies: &Cookies, public_key: &PublicKey) -> Option<String> {
     cookies
         .get(&public_key.to_string())
         .map(|c| c.value().to_string())
 }
-
-// TODO: unit test this
-fn session_secret_from_headers(headers: &HeaderMap, public_key: &PublicKey) -> Option<String> {
-    headers
-        .get_all(header::COOKIE)
-        .iter()
-        .filter_map(|h| h.to_str().ok())
-        .find(|h| h.starts_with(&public_key.to_string()))
-        .and_then(|h| {
-            h.split(';')
-                .next()
-                .and_then(|key_value| key_value.split('=').last())
-        })
-        .map(|s| s.to_string())
-}

pubky-homeserver/src/core/routes/auth.rs

@@ -1,6 +1,5 @@
 use axum::{
     extract::{Host, State},
-    http::StatusCode,
     response::IntoResponse,
 };
 use axum_extra::{headers::UserAgent, TypedHeader};
@@ -9,13 +8,7 @@ use tower_cookies::{cookie::SameSite, Cookie, Cookies};
 
 use pubky_common::{crypto::random_bytes, session::Session, timestamp::Timestamp};
 
-use crate::core::{
-    database::tables::users::User,
-    error::{Error, Result},
-    extractors::PubkyHost,
-    layers::authz::session_secret_from_cookies,
-    AppState,
-};
+use crate::core::{database::tables::users::User, error::Result, AppState};
 
 pub async fn signup(
     State(state): State<AppState>,
@@ -29,36 +22,6 @@ pub async fn signup(
     signin(State(state), user_agent, cookies, host, body).await
 }
 
-pub async fn session(
-    State(state): State<AppState>,
-    cookies: Cookies,
-    pubky: PubkyHost,
-) -> Result<impl IntoResponse> {
-    if let Some(secret) = session_secret_from_cookies(cookies, pubky.public_key()) {
-        if let Some(session) = state.db.get_session(&secret)? {
-            // TODO: add content-type
-            return Ok(session.serialize());
-        };
-    }
-
-    Err(Error::with_status(StatusCode::NOT_FOUND))
-}
-
-pub async fn signout(
-    State(mut state): State<AppState>,
-    cookies: Cookies,
-    pubky: PubkyHost,
-) -> Result<impl IntoResponse> {
-    // TODO: Set expired cookie to delete the cookie on client side.
-
-    if let Some(secret) = session_secret_from_cookies(cookies, pubky.public_key()) {
-        state.db.delete_session(&secret)?;
-    }
-
-    // Idempotent Success Response (200 OK)
-    Ok(())
-}
-
 pub async fn signin(
     State(state): State<AppState>,
     user_agent: Option<TypedHeader<UserAgent>>,

pubky-homeserver/src/core/routes/tenants/mod.rs

@@ -14,9 +14,8 @@ use crate::core::{
     AppState,
 };
 
-use super::auth;
-
 pub mod read;
+pub mod session;
 pub mod write;
 
 pub fn router(state: AppState) -> Router<AppState> {
@@ -28,8 +27,8 @@ pub fn router(state: AppState) -> Router<AppState> {
         .route("/pub/*path", put(write::put))
         .route("/pub/*path", delete(write::delete))
         // - Session routes
-        .route("/session", get(auth::session))
-        .route("/session", delete(auth::signout))
+        .route("/session", get(session::session))
+        .route("/session", delete(session::signout))
         // Layers
         // TODO: different max size for sessions and other routes?
         .layer(DefaultBodyLimit::max(100 * 1024 * 1024))

pubky-homeserver/src/core/routes/tenants/session.rs

@@ -0,0 +1,38 @@
+use axum::{extract::State, http::StatusCode, response::IntoResponse};
+use tower_cookies::Cookies;
+
+use crate::core::{
+    error::{Error, Result},
+    extractors::PubkyHost,
+    layers::authz::session_secret_from_cookies,
+    AppState,
+};
+
+pub async fn session(
+    State(state): State<AppState>,
+    cookies: Cookies,
+    pubky: PubkyHost,
+) -> Result<impl IntoResponse> {
+    if let Some(secret) = session_secret_from_cookies(&cookies, pubky.public_key()) {
+        if let Some(session) = state.db.get_session(&secret)? {
+            // TODO: add content-type
+            return Ok(session.serialize());
+        };
+    }
+
+    Err(Error::with_status(StatusCode::NOT_FOUND))
+}
+pub async fn signout(
+    State(mut state): State<AppState>,
+    cookies: Cookies,
+    pubky: PubkyHost,
+) -> Result<impl IntoResponse> {
+    // TODO: Set expired cookie to delete the cookie on client side.
+
+    if let Some(secret) = session_secret_from_cookies(&cookies, pubky.public_key()) {
+        state.db.delete_session(&secret)?;
+    }
+
+    // Idempotent Success Response (200 OK)
+    Ok(())
+}

pubky/src/shared/auth.rs

@@ -346,4 +346,43 @@ mod tests {
             StatusCode::FORBIDDEN
         );
     }
+
+    #[tokio::test]
+    async fn multiple_users() {
+        let testnet = Testnet::new(10).unwrap();
+        let server = Homeserver::start_test(&testnet).await.unwrap();
+
+        let client = Client::test(&testnet);
+
+        let first_keypair = Keypair::random();
+        let second_keypair = Keypair::random();
+
+        client
+            .signup(&first_keypair, &server.public_key())
+            .await
+            .unwrap();
+
+        client
+            .signup(&second_keypair, &server.public_key())
+            .await
+            .unwrap();
+
+        let session = client
+            .session(&first_keypair.public_key())
+            .await
+            .unwrap()
+            .unwrap();
+
+        assert_eq!(session.pubky(), &first_keypair.public_key());
+        assert!(session.capabilities().contains(&Capability::root()));
+
+        let session = client
+            .session(&second_keypair.public_key())
+            .await
+            .unwrap()
+            .unwrap();
+
+        assert_eq!(session.pubky(), &second_keypair.public_key());
+        assert!(session.capabilities().contains(&Capability::root()));
+    }
 }