Updated user authentication system

src/main.rs

@@ -1,24 +1,34 @@
 mod users;
 
-use std::sync::Mutex;
 use users::*;
 
-use actix_web::{App, HttpServer};
 use actix_web::web::Data;
+use actix_web::{App, HttpServer};
 use sqlx::sqlite::SqlitePool;
 
 #[derive(Clone)]
 struct AppState {
     database: SqlitePool,
+    token_expiration: u64,
 }
 
 #[actix_web::main]
 async fn main() -> std::io::Result<()> {
-    let pool = SqlitePool::connect("sqlite:database.db").await.expect("Could not connect to db");
-    let app_state = Data::new(AppState { database: pool });
+    let pool = SqlitePool::connect("sqlite:database.db")
+        .await
+        .expect("Could not connect to db");
+    let app_state = Data::new(AppState {
+        database: pool,
+        token_expiration: 86400,
+    });
     HttpServer::new(move || {
         App::new()
             .service(login)
+            .service(register)
+            .service(logout)
             .app_data(app_state.clone())
-    }).bind(("127.0.0.1", 8000))?.run().await
+    })
+        .bind(("127.0.0.1", 8000))?
+        .run()
+        .await
 }

src/users.rs

@@ -1,15 +1,34 @@
-use actix_web::web::Data;
-use actix_web::{post, HttpResponse, Responder};
-use argon2::password_hash::{SaltString, rand_core::OsRng, PasswordHash, PasswordVerifier, PasswordHasher};
-use argon2::Argon2;
-use sqlx::{query, query_as};
+use std::fs::File;
+use std::io::Read;
 use crate::AppState;
+use actix_web::cookie::Cookie;
+use actix_web::web::{Data, Json};
+use actix_web::{HttpRequest, HttpResponse, Responder, post};
+use argon2::Argon2;
+use argon2::password_hash::{
+    PasswordHash, PasswordHasher, PasswordVerifier, SaltString, rand_core::OsRng,
+};
+use jsonwebtoken::{
+    Algorithm, DecodingKey, EncodingKey, Header, Validation, decode, encode, get_current_timestamp,
+};
+use rand::Rng;
+use serde::{Deserialize, Serialize};
+use sqlx::{query, query_as};
+use uuid::Uuid;
 
+#[derive(Serialize, Deserialize)]
 struct UserLogin {
     username: String,
     password: String,
 }
 
+#[derive(Serialize, Deserialize)]
+struct UserRegister {
+    username: String,
+    password: String,
+    email: String,
+}
+
 struct User {
     id: i64,
     uuid: String,
@@ -18,29 +37,123 @@ struct User {
     email: String,
 }
 
+#[derive(Serialize, Deserialize)]
+struct UserTokenClaims {
+    exp: usize, // Required (validate_exp defaults to true in validation). Expiration time (as UTC timestamp)
+    kid: i64,
+    uid: String,
+}
+
 #[post("/login")]
-async fn login(user_login: Data<UserLogin>, app_state: Data<AppState>) -> impl Responder {
+async fn login(user_login: Json<UserLogin>, app_state: Data<AppState>) -> impl Responder {
     // Verify that the password is correct
     let argon2 = Argon2::default();
-    let user = query_as!(User, "SELECT * FROM users WHERE username = $1", user_login.username).fetch_one(&app_state.database).await;
+    let user = query_as!(
+        User,
+        "SELECT * FROM users WHERE username = $1",
+        user_login.username
+    )
+        .fetch_one(&app_state.database)
+        .await;
     if user.is_err() {
-        return HttpResponse::BadRequest();
+        return HttpResponse::BadRequest().finish();
     }
     let user = user.unwrap();
     let hash = PasswordHash::new(&user.hash);
     if hash.is_err() {
-        return HttpResponse::BadRequest();
+        return HttpResponse::BadRequest().finish();
     }
-    if argon2.verify_password(user_login.password.as_bytes(), &hash.unwrap()).is_err() {
-        return HttpResponse::BadRequest();
+    if argon2
+        .verify_password(user_login.password.as_bytes(), &hash.unwrap())
+        .is_err()
+    {
+        return HttpResponse::BadRequest().finish();
     }
     // Create the JWT
-
+    let header = Header::new(Algorithm::ES256);
+    // Put a random KeyId
+    let mut rng = rand::rng();
+    let key_id: i64 = rng.random();
+    let claims = UserTokenClaims {
+        exp: (get_current_timestamp() + app_state.token_expiration) as usize,
+        kid: key_id,
+        uid: user.uuid.clone(),
+    };
+    let mut key = File::open("priv.pem").unwrap();
+    let mut buf = vec![];
+    key.read_to_end(&mut buf).unwrap();
+    let token = encode(&header, &claims, &EncodingKey::from_ec_pem(&buf).unwrap()).unwrap();
     // Send the JWT as cookie
     HttpResponse::Ok()
+        .cookie(Cookie::new("token", token))
+        .finish()
 }
 
 #[post("/logout")]
-async fn logout(_token: Data<String>) -> impl Responder {
-    HttpResponse::Ok()
-}
+async fn logout(req: HttpRequest, app_state: Data<AppState>) -> impl Responder {
+    // Put the (KeyId, User) pair in the revoked table
+    // And remove data from client
+    let token = req.cookie("token");
+    println!("token: {:?}", token);
+    if token.is_none() {
+        return HttpResponse::BadRequest().finish();
+    }
+    let token = token.unwrap();
+    let token = token.value();
+    let mut key = File::open("pub.pem").unwrap();
+    let mut buf = vec![];
+    key.read_to_end(&mut buf).unwrap();
+    let token = decode::<UserTokenClaims>(
+        token,
+        &DecodingKey::from_ec_pem(&buf).unwrap(),
+        &Validation::new(Algorithm::ES256),
+    );
+    match token {
+        Ok(token) => {
+            if query!(
+                "INSERT INTO revoked ( token_id, user_id ) VALUES ( $1, $2 )",
+                token.claims.kid,
+                token.claims.uid
+            )
+                .execute(&app_state.database)
+                .await
+                .is_err()
+            {
+                return HttpResponse::InternalServerError().reason("Query fail").finish();
+            }
+        }
+        Err(e) => {
+            let message = format!("Error: {e}");
+            println!("{}", message);
+            return HttpResponse::InternalServerError().reason("Token caca").finish();
+        }
+    }
+    let mut cookie = Cookie::new("token", "");
+    cookie.make_removal();
+    HttpResponse::Ok().cookie(cookie).finish()
+}
+
+#[post("/register")]
+async fn register(user_register: Json<UserRegister>, app_state: Data<AppState>) -> impl Responder {
+    let uuid = Uuid::new_v4().to_string();
+    let argon2 = Argon2::default();
+    let salt = SaltString::generate(&mut OsRng);
+    let hash = argon2
+        .hash_password(user_register.password.as_bytes(), &salt)
+        .unwrap()
+        .to_string();
+    if query!(
+        "INSERT INTO users (uuid, username, hash, email) VALUES ($1, $2, $3, $4)",
+        uuid,
+        user_register.username,
+        hash,
+        user_register.email
+    )
+        .execute(&app_state.database)
+        .await
+        .is_err()
+    {
+        return HttpResponse::InternalServerError().finish();
+    };
+    HttpResponse::Ok().finish()
+}