refresh tokens

Cargo.lock

@@ -72,6 +72,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90"
 dependencies = [
  "axum-core",
+ "axum-macros",
  "bytes",
  "form_urlencoded",
  "futures-util",
@@ -98,6 +99,19 @@ dependencies = [
  "tracing",
 ]
 
+[[package]]
+name = "axum-cookie"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c687fa82965c81556c57eb0f7fa71f11058b4bc7cc18adee10b7f7c19f4c59b9"
+dependencies = [
+ "axum-core",
+ "cookie-rs",
+ "http",
+ "tower-layer",
+ "tower-service",
+]
+
 [[package]]
 name = "axum-core"
 version = "0.5.6"
@@ -117,6 +131,17 @@ dependencies = [
  "tracing",
 ]
 
+[[package]]
+name = "axum-macros"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7aa268c23bfbbd2c4363b9cd302a4f504fb2a9dfe7e3451d66f35dd392e20aca"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "base64"
 version = "0.22.1"
@@ -245,6 +270,12 @@ version = "0.10.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c"
 
+[[package]]
+name = "cookie-rs"
+version = "0.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5933ded34847986a1348361a52d526fd2fc9ac976cfc84cfaf9ffb9fda73132a"
+
 [[package]]
 name = "core-foundation-sys"
 version = "0.8.7"
@@ -1349,6 +1380,7 @@ version = "0.1.0"
 dependencies = [
  "argon2",
  "axum",
+ "axum-cookie",
  "chrono",
  "dotenvy",
  "hex",
@@ -1360,6 +1392,8 @@ dependencies = [
  "sqlx",
  "thiserror",
  "tokio",
+ "tower",
+ "tower-http",
  "tracing",
  "tracing-appender",
  "tracing-subscriber",
@@ -1993,6 +2027,22 @@ dependencies = [
  "tracing",
 ]
 
+[[package]]
+name = "tower-http"
+version = "0.6.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d4e6559d53cc268e5031cd8429d05415bc4cb4aefc4aa5d6cc35fbf5b924a1f8"
+dependencies = [
+ "bitflags",
+ "bytes",
+ "http",
+ "http-body",
+ "pin-project-lite",
+ "tower-layer",
+ "tower-service",
+ "tracing",
+]
+
 [[package]]
 name = "tower-layer"
 version = "0.3.3"

Cargo.toml

@@ -11,7 +11,7 @@ tracing-subscriber = {version="0.3.23", features = ["env-filter", "json"]}
 tracing-tree = "0.4.1"
 tokio = { version = "1.52.1", features = ["rt-multi-thread", "macros", "signal"] }
 sqlx = { version = "0.8", features = [ "runtime-tokio", "postgres", "chrono", "uuid" ] }
-axum = "0.8.9"
+axum = { version = "0.8.9", features = ["macros"] }
 thiserror = "2"
 serde = { version = "1.0.228", features = ["derive"] }
 serde_json = "1.0.149"
@@ -22,3 +22,6 @@ uuid = { version = "1.23.1", features = ["serde", "v4"] }
 rand = "0.10.1"
 sha2 = "0.11.0"
 hex = "0.4.3"
+axum-cookie = "0.2.4"
+tower = "0.5.3"
+tower-http = { version = "0.6.8", features = ["trace", "tracing"] }

src/controller/mod.rs

@@ -1,6 +1,8 @@
 use axum::{Router, routing::get};
 
 use crate::state::AppState;
+use tower::ServiceBuilder;
+use tower_http::trace::TraceLayer;
 
 pub mod model;
 mod v1;
@@ -9,4 +11,5 @@ pub fn router() -> Router<AppState> {
     Router::new()
         .route("/", get("Server is going brr 🚀"))
         .nest("/api/v1", v1::router_v1())
+        .layer(ServiceBuilder::new().layer(TraceLayer::new_for_http()))
 }

src/controller/model/auth_model.rs

@@ -14,5 +14,4 @@ pub struct RegisterRequest {
 #[derive(Serialize)]
 pub struct AuthResponse {
     pub access_token: String,
-    pub refresh_token: String,
 }

src/controller/v1/auth_controller.rs

@@ -1,10 +1,11 @@
-use axum::extract::State;
+use axum::{routing::post, Json, Router};
-use axum::{Json, Router, routing::post};
+use axum::{extract::State, http::header::SET_COOKIE, response::IntoResponse};
+use axum::http::HeaderMap;
 
 use crate::{
-    controller::model::auth_model::{AuthResponse, LoginRequest, RegisterRequest},
+    controller::model::auth_model::AuthResponse,
     errors::AppError,
-    service::auth_service::{login, register},
+    service::auth_service::{login, refresh as refresh_service, register},
     state::AppState,
 };
 
@@ -12,17 +13,65 @@ pub fn auth_router() -> Router<AppState> {
     Router::new()
         .route("/login", post(login_handler))
         .route("/register", post(register_handler))
+        .route("/refresh", post(refresh_handler))
 }
 
 async fn login_handler(
     State(s): State<AppState>,
-    Json(payload): Json<LoginRequest>,
+    Json(payload): Json<crate::controller::model::auth_model::LoginRequest>,
-) -> Result<Json<AuthResponse>, AppError> {
+) -> Result<impl IntoResponse, AppError> {
-    login(&s, payload).await
+    let (response, refresh_token) = login(&s, payload).await?;
+    let cookie = build_refresh_cookie(&refresh_token, s.app_env == crate::config::AppEnv::Production);
+    let mut headers = HeaderMap::new();
+    headers.insert(SET_COOKIE, cookie.parse().unwrap());
+    Ok((headers, Json(response)))
 }
+
 async fn register_handler(
     State(s): State<AppState>,
-    Json(payload): Json<RegisterRequest>,
+    Json(payload): Json<crate::controller::model::auth_model::RegisterRequest>,
-) -> Result<Json<AuthResponse>, AppError> {
+) -> Result<impl IntoResponse, AppError> {
-    register(&s, payload).await
+    let (response, refresh_token) = register(&s, payload).await?;
+    let cookie = build_refresh_cookie(&refresh_token, s.app_env == crate::config::AppEnv::Production);
+    let mut headers = HeaderMap::new();
+    headers.insert(SET_COOKIE, cookie.parse().unwrap());
+    Ok((headers, Json(response)))
+}
+
+async fn refresh_handler(
+    State(s): State<AppState>,
+    headers: axum::http::request::Parts,
+) -> Result<impl IntoResponse, AppError> {
+    let cookie_header = headers.headers.get("cookie");
+    let refresh_token = parse_cookie(cookie_header, "refresh_token").ok_or(AppError::InvalidCredentials)?;
+    let (response, new_token) = refresh_service(&s, refresh_token).await?;
+    let cookie = build_refresh_cookie(&new_token, s.app_env == crate::config::AppEnv::Production);
+    let mut response_headers = HeaderMap::new();
+    response_headers.insert(SET_COOKIE, cookie.parse().unwrap());
+    Ok((response_headers, Json(response)))
+}
+
+fn build_refresh_cookie(token: &str, is_prod: bool) -> String {
+    let mut cookie = format!(
+        "refresh_token={}; HttpOnly; SameSite=Strict; Path=/auth/refresh; Max-Age=604800",
+        token
+    );
+    if is_prod {
+        cookie.push_str("; Secure");
+    }
+    cookie
+}
+
+fn parse_cookie(cookie_header: Option<&axum::http::HeaderValue>, name: &str) -> Option<String> {
+    let cookie_header = cookie_header?.to_str().ok()?;
+    cookie_header
+        .split(';')
+        .filter_map(|pair| {
+            let mut parts = pair.trim().splitn(2, '=');
+            match (parts.next(), parts.next()) {
+                (Some(k), Some(v)) if k == name => Some(v.to_string()),
+                _ => None,
+            }
+        })
+        .next()
 }

src/server.rs

@@ -7,6 +7,7 @@ pub async fn init(cfg: &config::Config, db: PgPool) -> Result<(), AppError> {
     let state = AppState {
         db,
         jwt_secret: cfg.jwt_secret.clone(),
+        app_env: cfg.app_env,
     };
     let app = Router::new().merge(controller::router()).with_state(state);
 

src/service/auth_service.rs

@@ -1,33 +1,36 @@
 use std::time::Instant;
 
-use axum::Json;
+use chrono::Duration;
-use chrono::{Duration, Utc};
 
 use crate::controller::model::auth_model::*;
-use crate::db::repository::refresh_token_repository::create_refresh_token;
+use crate::db::repository::refresh_token_repository::{self, create_refresh_token};
 use crate::db::repository::user_repository;
 use crate::errors::AppError;
 use crate::state::AppState;
 use crate::utils::anti_enumeration::anti_enumeration_delay;
-use crate::utils::hash;
 use crate::utils::jwt::generate_access_token;
 use crate::utils::refresh_token::generate_refresh_token;
+use crate::utils::{hash, refresh_token};
 
-pub async fn login(state: &AppState, req: LoginRequest) -> Result<Json<AuthResponse>, AppError> {
+pub async fn login(
+    _state: &AppState,
+    _req: LoginRequest,
+) -> Result<(AuthResponse, String), AppError> {
+    let start = Instant::now();
     todo!()
 }
+
 pub async fn register(
     state: &AppState,
     req: RegisterRequest,
-) -> Result<Json<AuthResponse>, AppError> {
+) -> Result<(AuthResponse, String), AppError> {
     let start = Instant::now();
     let mut tx = state.db.begin().await?;
     {
         let user = user_repository::get_user_by_email(&mut *tx, &req.email).await?;
         if user.is_some() {
-            // user already registered
             anti_enumeration_delay(start, 150, 300).await;
-            return Err(AppError::Internal);
+            return Err(AppError::Validation("Email already registered".to_string()));
         }
     }
     let h = hash::hash(&req.password)?;
@@ -40,9 +43,49 @@ pub async fn register(
 
     tx.commit().await?;
     anti_enumeration_delay(start, 150, 300).await;
-    // TODO: put refresh token in cookie
+
-    Ok(Json(AuthResponse {
+    Ok((AuthResponse { access_token }, refresh_plain))
-        access_token: access_token,
+}
-        refresh_token: refresh_plain,
+
-    }))
+pub async fn refresh(
+    state: &AppState,
+    refresh_token: String,
+) -> Result<(AuthResponse, String), AppError> {
+    let start = Instant::now();
+
+    let token_hash = refresh_token::hash_token(&refresh_token);
+    let stored_token = refresh_token_repository::find_by_hash(&state.db, &token_hash)
+        .await?
+        .ok_or(AppError::InvalidCredentials)?;
+
+    if stored_token.revoked_at.is_some() {
+        anti_enumeration_delay(start, 150, 300).await;
+        return Err(AppError::InvalidCredentials);
+    }
+    if stored_token.expires_at < chrono::Utc::now() {
+        anti_enumeration_delay(start, 150, 300).await;
+        return Err(AppError::InvalidCredentials);
+    }
+
+    let user = user_repository::get_user_by_id(&state.db, stored_token.user_id)
+        .await?
+        .ok_or(AppError::InvalidCredentials)?;
+
+    let new_access_token = generate_access_token(user.id, &state.jwt_secret)?;
+    let (new_refresh_plain, new_refresh_hash) = generate_refresh_token();
+    let new_expires_at = chrono::Utc::now() + Duration::days(7);
+
+    let mut tx = state.db.begin().await?;
+    crate::db::repository::refresh_token_repository::revoke(&mut *tx, stored_token.id).await?;
+    create_refresh_token(&mut *tx, user.id, new_refresh_hash, new_expires_at).await?;
+    tx.commit().await?;
+
+    anti_enumeration_delay(start, 150, 300).await;
+
+    Ok((
+        AuthResponse {
+            access_token: new_access_token,
+        },
+        new_refresh_plain,
+    ))
 }

src/state.rs

@@ -1,7 +1,9 @@
+use crate::config::AppEnv;
 use sqlx::PgPool;
 
 #[derive(Clone)]
 pub struct AppState {
     pub db: PgPool,
     pub jwt_secret: String,
+    pub app_env: AppEnv,
 }

src/utils/refresh_token.rs

@@ -16,3 +16,9 @@ pub fn generate_refresh_token() -> (String, String) {
 
     (plain, hash)
 }
+
+pub fn hash_token(token: &str) -> String {
+    let mut hasher = Sha256::new();
+    hasher.update(token.as_bytes());
+    hex::encode(hasher.finalize())
+}