feat(auth): tower session

3 files changed, 179 insertions, 6 deletions

diff --git a/crates/auth/src/server/csrf_token_validation.rs b/crates/auth/src/server/csrf_token_validation.rs
new file mode 100644
index 0000000..c9a627c
--- /dev/null
+++ b/crates/auth/src/server/csrf_token_validation.rs
@@ -0,0 +1,49 @@
+use anyhow::{Context, anyhow};
+use axum_extra::headers;
+use oauth2::CsrfToken;
+use time::OffsetDateTime;
+use tower_sessions::{CachingSessionStore, SessionStore, session::Id};
+use tower_sessions_moka_store::MokaStore;
+use tower_sessions_sqlx_store::PostgresStore;
+
+use crate::{
+    error::AppError,
+    server::{COOKIE_NAME, CSRF_TOKEN, routes::authorised::AuthRequest},
+    state::AppHandle,
+};
+
+pub struct Session {
+    id: String,
+    expires_at: OffsetDateTime,
+    user_id: String,
+}
+
+pub async fn csrf_token_validation_workflow(
+    auth_request: &AuthRequest,
+    store: &CachingSessionStore<MokaStore, PostgresStore>,
+    oauth_session_id: Id,
+) -> Result<(), AppError> {
+    let oauth_session = store.load(&oauth_session_id).await.unwrap().unwrap();
+
+    // Extract the CSRF token from the session
+    let csrf_token_serialized = oauth_session
+        .data
+        .get(CSRF_TOKEN)
+        .context("failed to get value from session")?;
+    let csrf_token = serde_json::from_value::<CsrfToken>(csrf_token_serialized.clone())
+        .context("CSRF token not found in session")?
+        .to_owned();
+
+    // Cleanup the CSRF token session
+    store
+        .delete(&oauth_session_id)
+        .await
+        .context("Failed to destroy old session")?;
+
+    // Validate CSRF token is the same as the one in the auth request
+    if *csrf_token.secret() != auth_request.state {
+        return Err(anyhow!("CSRF token mismatch").into());
+    }
+
+    Ok(())
+}
diff --git a/crates/auth/src/server/routes/authorised.rs b/crates/auth/src/server/routes/authorised.rs
index ddf048d..42bbde2 100644
--- a/crates/auth/src/server/routes/authorised.rs
+++ b/crates/auth/src/server/routes/authorised.rs
@@ -1,23 +1,108 @@
+use std::{str::FromStr, time::Duration};
+
+use anyhow::Context;
 use axum::{
     extract::{Query, State},
-    response::IntoResponse,
+    http::HeaderMap,
+    response::{IntoResponse, Redirect},
 };
 use axum_extra::{TypedHeader, headers};
-use serde::Deserialize;
+use oauth2::{AuthorizationCode, TokenResponse};
+use reqwest::header::SET_COOKIE;
+use serde::{Deserialize, Serialize};
+use sqlx::types::uuid;
+use tower_sessions::{
+    SessionStore,
+    session::{Id, Record},
+};
 
-use crate::{error::AppError, server::routes::Provider, state::AppHandle};
+use crate::{
+    error::AppError,
+    server::{
+        OAUTH_CSRF_COOKIE, csrf_token_validation::csrf_token_validation_workflow, routes::Provider,
+    },
+    state::AppHandle,
+};
 
 #[derive(Debug, Deserialize)]
 pub struct AuthRequest {
     provider: Provider,
     code: String,
-    state: String,
+    pub state: String,
 }
 
+#[derive(Debug, Deserialize, Serialize)]
+struct User {
+    id: String,
+    avatar: Option<String>,
+    username: String,
+    discriminator: String,
+}
+
+/// The cookie to store the session id for user information.
+const SESSION_COOKIE: &str = "info";
+const SESSION_DATA_KEY: &str = "data";
+
 async fn login_authorized(
     Query(query): Query<AuthRequest>,
     State(state): State<AppHandle>,
     TypedHeader(cookies): TypedHeader<headers::Cookie>,
 ) -> Result<impl IntoResponse, AppError> {
-    Ok("")
+    let oauth_session_id = Id::from_str(
+        cookies
+            .get(OAUTH_CSRF_COOKIE)
+            .context("missing session cookie")?,
+    )
+    .unwrap();
+    csrf_token_validation_workflow(&query, &state.session_store, oauth_session_id).await?;
+
+    let client = state.http_client.clone();
+    let store = state.session_store.clone();
+
+    // Get an auth token
+    let token = state
+        .discord_client
+        .exchange_code(AuthorizationCode::new(query.code.clone()))
+        .request_async(&client)
+        .await
+        .context("failed in sending request request to authorization server")?;
+
+    let user_data: User = client
+        // https://discord.com/developers/docs/resources/user#get-current-user
+        .get("https://discordapp.com/api/users/@me")
+        .bearer_auth(token.access_token().secret())
+        .send()
+        .await
+        .context("failed in sending request to target Url")?
+        .json::<User>()
+        .await
+        .context("failed to deserialize response as JSON")?;
+
+    // Create a new session filled with user data
+    let session_id = Id(i128::from_le_bytes(uuid::Uuid::new_v4().to_bytes_le()));
+    store
+        .create(&mut Record {
+            id: session_id,
+            data: [(
+                SESSION_DATA_KEY.to_string(),
+                serde_json::to_value(user_data).unwrap(),
+            )]
+            .into(),
+            expiry_date: time::OffsetDateTime::now_utc()
+                + Duration::from_secs(state.local_config.oauth.session_lifespan),
+        })
+        .await
+        .context("failed in inserting serialized value into session")?;
+
+    // Store session and get corresponding cookie.
+    let cookie = format!("{SESSION_COOKIE}={session_id}; SameSite=Lax; HttpOnly; Secure; Path=/");
+
+    // Set cookie
+    let mut headers = HeaderMap::new();
+    headers.insert(
+        SET_COOKIE,
+        cookie.parse().context("failed to parse cookie")?,
+    );
+
+    Ok((headers, Redirect::to("/")))
 }
diff --git a/crates/auth/src/server/routes/discord/discord_auth.rs b/crates/auth/src/server/routes/discord/discord_auth.rs
index b07fa7a..5257a33 100644
--- a/crates/auth/src/server/routes/discord/discord_auth.rs
+++ b/crates/auth/src/server/routes/discord/discord_auth.rs
@@ -1,11 +1,24 @@
+use std::time::Duration;
+
+use anyhow::Context;
 use axum::{
     extract::State,
     http::HeaderMap,
     response::{IntoResponse, Redirect},
 };
 use oauth2::{CsrfToken, Scope};
+use reqwest::header::SET_COOKIE;
+use sqlx::types::uuid;
+use tower_sessions::{
+    SessionStore,
+    session::{Id, Record},
+};
 
-use crate::{error::AppError, state::AppHandle};
+use crate::{
+    error::AppError,
+    server::{CSRF_TOKEN, OAUTH_CSRF_COOKIE},
+    state::AppHandle,
+};
 
 pub async fn discord_auth(State(state): State<AppHandle>) -> Result<impl IntoResponse, AppError> {
     let (auth_url, csrf_token) = state
@@ -14,7 +27,33 @@ pub async fn discord_auth(State(state): State<AppHandle>) -> Result<impl IntoRes
         .add_scope(Scope::new("identify".to_string()))
         .url();
 
+    // Store the token in the session and retrieve the session cookie.
+    let session_id = Id(i128::from_le_bytes(uuid::Uuid::new_v4().to_bytes_le()));
+    let store = state.session_store.clone();
+
+    store
+        .create(&mut Record {
+            id: session_id,
+            data: [(
+                CSRF_TOKEN.to_string(),
+                serde_json::to_value(csrf_token).unwrap(),
+            )]
+            .into(),
+            expiry_date: time::OffsetDateTime::now_utc()
+                + Duration::from_secs(state.local_config.oauth.session_lifespan),
+        })
+        .await
+        .unwrap();
+    //        .context("failed in inserting CSRF token into session")?;
+
+    // Attach the session cookie to the response header
+    let cookie =
+        format!("{OAUTH_CSRF_COOKIE}={session_id}; SameSite=Lax; HttpOnly; Secure; Path=/");
     let mut headers = HeaderMap::new();
+    headers.insert(
+        SET_COOKIE,
+        cookie.parse().context("failed to parse cookie")?,
+    );
 
     Ok((headers, Redirect::to(auth_url.as_ref())))
 }