From d2b0950e28953d31d79ab1c82b9bfd38c483a2e3 Mon Sep 17 00:00:00 2001
From: Luc <luc@lucemans.nl>
Date: Sat, 27 Jul 2024 21:02:14 +0000
Subject: [PATCH] Update Sessions

---
 engine/migrations/0002_sessions.sql |  1 +
 engine/src/database/mod.rs          |  5 ++--
 engine/src/routes/auth.rs           | 40 +++++++++++++++++++++++------
 3 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/engine/migrations/0002_sessions.sql b/engine/migrations/0002_sessions.sql
index a69f676..178ae15 100644
--- a/engine/migrations/0002_sessions.sql
+++ b/engine/migrations/0002_sessions.sql
@@ -3,6 +3,7 @@ CREATE TABLE IF NOT EXISTS sessions
     id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
     user_id INT NOT NULL,
     user_agent VARCHAR(255) NOT NULL,
+    user_ip VARCHAR(255) NOT NULL,
     last_access TIMESTAMPTZ NOT NULL DEFAULT NOW(),
     valid BOOLEAN NOT NULL DEFAULT TRUE
 );
diff --git a/engine/src/database/mod.rs b/engine/src/database/mod.rs
index 8f0af1f..994c5e7 100644
--- a/engine/src/database/mod.rs
+++ b/engine/src/database/mod.rs
@@ -63,10 +63,11 @@ impl Database {
         Ok(user)
     }
 
-    pub async fn create_session(&self, user_id: i32, user_agent: &str) -> Result<SessionState, sqlx::Error> {
-        let session = sqlx::query_as::<_, SessionState>("INSERT INTO sessions (user_id, user_agent) VALUES ($1, $2) RETURNING *")
+    pub async fn create_session(&self, user_id: i32, user_agent: &str, user_ip: &str) -> Result<SessionState, sqlx::Error> {
+        let session = sqlx::query_as::<_, SessionState>("INSERT INTO sessions (user_id, user_agent, user_ip) VALUES ($1, $2, $3) RETURNING *")
             .bind(user_id)
             .bind(user_agent)
+            .bind(user_ip)
             .fetch_one(&self.pool)
             .await?;
         Ok(session)
diff --git a/engine/src/routes/auth.rs b/engine/src/routes/auth.rs
index f01ee68..e269450 100644
--- a/engine/src/routes/auth.rs
+++ b/engine/src/routes/auth.rs
@@ -1,10 +1,18 @@
 use crate::state::AppState;
 use openid::{Options, Token};
-use poem::{handler, web::{cookie::{Cookie, CookieJar}, Data, Json, Query, Redirect}, IntoResponse};
+use poem::{
+    handler,
+    http::HeaderMap,
+    web::{
+        cookie::{Cookie, CookieJar},
+        Data, Json, Query, RealIp, Redirect, RemoteAddr,
+    },
+    IntoResponse,
+};
 use serde::Deserialize;
+use std::sync::Arc;
 use tracing::info;
 use uuid::Uuid;
-use std::sync::Arc;
 
 #[handler]
 pub async fn login(state: Data<&Arc<AppState>>) -> impl IntoResponse {
@@ -44,7 +52,12 @@ pub struct MyQuery {
 }
 
 #[handler]
-pub async fn callback(query: Query<MyQuery>, state: Data<&Arc<AppState>>) -> impl IntoResponse {
+pub async fn callback(
+    query: Query<MyQuery>,
+    state: Data<&Arc<AppState>>,
+    ip: RealIp,
+    headers: &HeaderMap,
+) -> impl IntoResponse {
     let mut token = state.openid.request_token(&query.code).await.unwrap();
 
     let mut token = Token::from(token);
@@ -65,7 +78,14 @@ pub async fn callback(query: Query<MyQuery>, state: Data<&Arc<AppState>>) -> imp
         .await
         .unwrap();
 
-    let session = state.database.create_session(user.id, "test").await.unwrap();
+    let user_agent = headers.get("user-agent").unwrap().to_str().unwrap();
+    let user_ip = ip.0.unwrap().to_string();
+
+    let session = state
+        .database
+        .create_session(user.id, user_agent, &user_ip)
+        .await
+        .unwrap();
 
     // let session = state.database.get_session_by_id(&user.id).await.unwrap();
 
@@ -73,18 +93,22 @@ pub async fn callback(query: Query<MyQuery>, state: Data<&Arc<AppState>>) -> imp
 
     let token = session.id;
 
-    Redirect::temporary("http://localhost:3000/hello").with_header("Set-Cookie", format!("property.v3x.token={}", token))
+    Redirect::temporary("http://localhost:3000/me")
+        .with_header("Set-Cookie", format!("property.v3x.token={}", token))
 }
 
 #[handler]
-pub async fn me(state: Data<&Arc<AppState>>, cookies: &CookieJar) -> impl IntoResponse  {
-
+pub async fn me(state: Data<&Arc<AppState>>, cookies: &CookieJar) -> impl IntoResponse {
     let token = cookies.get("property.v3x.token").unwrap();
     let token = Uuid::parse_str(token.value_str()).unwrap();
 
     let session = state.database.get_session_by_id(token).await.unwrap();
 
-    let user = state.database.get_user_from_id(session.user_id).await.unwrap();
+    let user = state
+        .database
+        .get_user_from_id(session.user_id)
+        .await
+        .unwrap();
 
     Json(user)
 }