From 9d9e2bef055e1209d0350cce919bdc0f6383ed8e687f69c58a0c881b290c50d9 Mon Sep 17 00:00:00 2001
From: raven <7156279+RavenX8@users.noreply.github.com>
Date: Fri, 20 Dec 2024 17:46:04 -0500
Subject: [PATCH] - add: session_id to the validate token response - add:
 session_id to the jwt generated token

---
 auth-service/src/grpc.rs  | 17 ++++++++++-------
 auth-service/src/jwt.rs   |  8 +++++---
 auth-service/src/users.rs |  5 +++--
 proto/auth.proto          |  1 +
 4 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/auth-service/src/grpc.rs b/auth-service/src/grpc.rs
index f25c986..f0229e6 100644
--- a/auth-service/src/grpc.rs
+++ b/auth-service/src/grpc.rs
@@ -26,16 +26,14 @@ impl AuthService for MyAuthService {
 
         info!("Login attempt for username: {}", req.username);
 
-        if let Some(user_id) = verify_user(self.db_client.as_ref().clone(), &req.username, &req.password).await {
-            let token = generate_token(&user_id, vec!["user".to_string()])
-                .map_err(|_| Status::internal("Token generation failed"))?;
-
+        if let Some(user) = verify_user(self.db_client.as_ref().clone(), &req.username, &req.password).await {
+            let user_id = user.user_id.to_string();
             let session_id = uuid::Uuid::new_v4().to_string();
             let response = self
                 .session_client.as_ref().clone()
                 .create_session(CreateSessionRequest {
                     session_id: session_id.clone(),
-                    user_id: user_id.parse().unwrap(),
+                    user_id: user.user_id,
                     username: req.username.to_string(),
                     character_id: 0,
                     ip_address: req.ip_address.to_string(),
@@ -48,6 +46,9 @@ impl AuthService for MyAuthService {
             };
             let session_id = session.into_inner().session_id;
 
+            let token = generate_token(&user_id, &&session_id.clone(), user.roles)
+                .map_err(|_| Status::internal("Token generation failed"))?;
+            
             info!("Login successful for username: {}", req.username);
             Ok(Response::new(LoginResponse { token, user_id, session_id }))
         } else {
@@ -77,13 +78,15 @@ impl AuthService for MyAuthService {
         let req = request.into_inner();
 
         match validate_token(&req.token) {
-            Ok(user_id) => Ok(Response::new(ValidateTokenResponse {
+            Ok(user_data) => Ok(Response::new(ValidateTokenResponse {
                 valid: true,
-                user_id,
+                user_id: user_data.0,
+                session_id: user_data.1,
             })),
             Err(_) => Ok(Response::new(ValidateTokenResponse {
                 valid: false,
                 user_id: "".to_string(),
+                session_id: "".to_string(),
             })),
         }
     }
diff --git a/auth-service/src/jwt.rs b/auth-service/src/jwt.rs
index 9117fa3..d217dde 100644
--- a/auth-service/src/jwt.rs
+++ b/auth-service/src/jwt.rs
@@ -5,11 +5,12 @@ use std::env;
 #[derive(Debug, Serialize, Deserialize)]
 struct Claims {
     sub: String, // Subject (user ID)
+    session_id: String, // Session ID
     roles: Vec<String>, // Roles/permissions
     exp: usize,  // Expiration time
 }
 
-pub fn generate_token(user_id: &str, roles: Vec<String>) -> Result<String, jsonwebtoken::errors::Error> {
+pub fn generate_token(user_id: &str, session_id: &str, roles: Vec<String>) -> Result<String, jsonwebtoken::errors::Error> {
     let secret = env::var("JWT_SECRET").expect("JWT_SECRET must be set");
     let expiration = chrono::Utc::now()
         .checked_add_signed(chrono::Duration::days(1))
@@ -18,6 +19,7 @@ pub fn generate_token(user_id: &str, roles: Vec<String>) -> Result<String, jsonw
 
     let claims = Claims {
         sub: user_id.to_owned(),
+        session_id: session_id.to_owned(),
         roles,
         exp: expiration,
     };
@@ -25,12 +27,12 @@ pub fn generate_token(user_id: &str, roles: Vec<String>) -> Result<String, jsonw
     encode(&Header::default(), &claims, &EncodingKey::from_secret(secret.as_ref()))
 }
 
-pub fn validate_token(token: &str) -> Result<String, jsonwebtoken::errors::Error> {
+pub fn validate_token(token: &str) -> Result<(String, String), jsonwebtoken::errors::Error> {
     let secret = env::var("JWT_SECRET").expect("JWT_SECRET must be set");
     let token_data = decode::<Claims>(
         token,
         &DecodingKey::from_secret(secret.as_ref()),
         &Validation::default(),
     )?;
-    Ok(token_data.claims.sub)
+    Ok((token_data.claims.sub, token_data.claims.session_id))
 }
diff --git a/auth-service/src/users.rs b/auth-service/src/users.rs
index 0375f54..309c334 100644
--- a/auth-service/src/users.rs
+++ b/auth-service/src/users.rs
@@ -1,4 +1,5 @@
 use crate::database_client::DatabaseClientTrait;
+use crate::database::GetUserResponse;
 
 use argon2::{
     password_hash::{
@@ -20,11 +21,11 @@ pub fn verify_password(password: &str, hash: &str) -> bool {
 }
 
 pub async fn verify_user<T: DatabaseClientTrait>(mut db_client: T,
-                         username: &str, password: &str) -> Option<String> {
+                         username: &str, password: &str) -> Option<GetUserResponse> {
     let user = db_client.get_user_by_username(username).await.ok()?;
 
     if verify_password(password, &user.hashed_password) {
-        Some(user.user_id.to_string())
+        Some(user)
     } else {
         None
     }
diff --git a/proto/auth.proto b/proto/auth.proto
index 0b11af7..bf54786 100644
--- a/proto/auth.proto
+++ b/proto/auth.proto
@@ -35,6 +35,7 @@ message ValidateTokenRequest {
 message ValidateTokenResponse {
     bool valid = 1;
     string user_id = 2;
+    string session_id = 3;
 }
 
 message ValidateSessionRequest {