Refresh tokens and endpoint

Author: yungcomputerchair

Cargo.lock

@@ -14,6 +14,7 @@ log = "0.4.22"
 ring = "0.17.8"
 rustls = { version = "0.23.12", features = ["ring"], default-features = false, optional = true }
 serde = { version = "1.0.209", features = ["derive"] }
+serde_repr = "0.1.19"
 simplelog = "0.12.2"
 sqlite = "0.36.1"
 tokio = { version = "1.40.0", features = ["full"] }

Cargo.toml

@@ -25,8 +25,12 @@ placeholders = true
 
 [auth]
 route = "/auth"
+refresh_subroute = "/session"
 secret_path = "secret"
-valid_secs = 604_800 # one week
+# Session tokens are used to authenticate requests. They should be short-lived.
+# Refresh tokens are used to generate new session tokens. They should be long-lived and cached by clients.
+valid_secs_refresh = 604_800 # one week
+valid_secs_session = 900 # 15 minutes
 
 [cookie]
 route = "/cookie"

config.toml

@@ -1,18 +1,26 @@
 use std::sync::{Arc, OnceLock};
 
-use axum::{extract::State, http::StatusCode, routing::post, Json, Router};
+use axum::{
+    extract::State,
+    http::{HeaderMap, StatusCode},
+    routing::post,
+    Json, Router,
+};
 use jsonwebtoken::{get_current_timestamp, DecodingKey, EncodingKey, Validation};
 use log::{info, warn};
 use ring::rand::{SecureRandom, SystemRandom};
 use serde::{Deserialize, Serialize};
+use serde_repr::{Deserialize_repr, Serialize_repr};
 
 use crate::{util, AppState};
 
 #[derive(Deserialize, Clone)]
 pub struct AuthConfig {
     route: String,
+    refresh_subroute: String,
     secret_path: String,
-    valid_secs: u64,
+    valid_secs_refresh: u64,
+    valid_secs_session: u64,
 }
 
 #[derive(Deserialize)]
@@ -21,11 +29,19 @@ pub struct AuthRequest {
     password: String,
 }
 
+#[repr(u8)]
+#[derive(Deserialize_repr, Serialize_repr, PartialEq, Eq)]
+pub enum TokenKind {
+    Refresh = 0,
+    Session = 1,
+}
+
 #[derive(Deserialize, Serialize)]
 pub struct Claims {
-    sub: String, // account id as a string
-    crt: u64,    // creation timestamp in UTC
-    exp: u64,    // expiration timestamp in UTC
+    sub: String,     // account id as a string
+    crt: u64,        // creation timestamp in UTC
+    exp: u64,        // expiration timestamp in UTC
+    kind: TokenKind, // kind of token
 }
 
 static SECRET_KEY: OnceLock<Vec<u8>> = OnceLock::new();
@@ -55,21 +71,33 @@ pub fn register(
     rng: &SystemRandom,
 ) -> Router<Arc<AppState>> {
     let route = &config.route;
+    let refresh_route = util::get_subroute(route, &config.refresh_subroute);
     info!("Registering auth route @ {}", route);
+    info!("\tRefresh route @ {}", refresh_route);
     check_secret(&config.secret_path, rng);
-    routes.route(route, post(do_auth))
+    routes
+        .route(route, post(do_auth))
+        .route(&refresh_route, post(do_refresh))
 }
 
-fn gen_jwt(account_id: i64, valid_secs: u64) -> Result<String, String> {
+fn gen_jwt(auth_config: &AuthConfig, account_id: i64, kind: TokenKind) -> Result<String, String> {
     let secret = SECRET_KEY.get().unwrap();
     let key = EncodingKey::from_secret(secret);
+
+    let valid_secs = match kind {
+        TokenKind::Refresh => auth_config.valid_secs_refresh,
+        TokenKind::Session => auth_config.valid_secs_session,
+    };
+
     let crt = get_current_timestamp();
     let exp = crt + valid_secs;
     let claims = Claims {
         sub: account_id.to_string(),
         crt,
         exp,
+        kind,
     };
+
     jsonwebtoken::encode(&jsonwebtoken::Header::default(), &claims, &key)
         .map_err(|e| format!("JWT error: {}", e))
 }
@@ -85,7 +113,7 @@ fn get_validator(account_id: Option<i64>) -> Validation {
     validation
 }
 
-pub fn validate_jwt(jwt: &str) -> Result<i64, String> {
+pub fn validate_jwt(jwt: &str, kind: TokenKind) -> Result<i64, String> {
     let Some(secret) = SECRET_KEY.get() else {
         return Err("Auth module not initialized".to_string());
     };
@@ -102,6 +130,10 @@ pub fn validate_jwt(jwt: &str) -> Result<i64, String> {
         return Err("Expired JWT".to_string());
     }
 
+    if token.claims.kind != kind {
+        return Err("Bad token kind".to_string());
+    }
+
     match token.claims.sub.parse() {
         Ok(id) => Ok(id),
         Err(e) => Err(format!("Bad account ID: {}", e)),
@@ -118,8 +150,11 @@ async fn do_auth(
         warn!("Auth error: {}", e);
         (StatusCode::UNAUTHORIZED, "Invalid credentials".to_string())
     })?;
-    let valid_secs = app.config.auth.as_ref().unwrap().valid_secs;
-    match gen_jwt(account_id, valid_secs) {
+    match gen_jwt(
+        app.config.auth.as_ref().unwrap(),
+        account_id,
+        TokenKind::Refresh,
+    ) {
         Ok(jwt) => Ok(jwt),
         Err(e) => {
             warn!("Auth error: {}", e);
@@ -130,3 +165,30 @@ async fn do_auth(
         }
     }
 }
+
+async fn do_refresh(
+    State(app): State<Arc<AppState>>,
+    headers: HeaderMap,
+) -> Result<String, (StatusCode, String)> {
+    assert!(app.is_tls);
+    let db = app.db.lock().await;
+    // TODO validate the refresh token against the last password reset timestamp
+    let account_id = match util::validate_authed_request(&headers, TokenKind::Refresh) {
+        Ok(id) => id,
+        Err(e) => return Err((StatusCode::UNAUTHORIZED, e)),
+    };
+    match gen_jwt(
+        app.config.auth.as_ref().unwrap(),
+        account_id,
+        TokenKind::Session,
+    ) {
+        Ok(jwt) => Ok(jwt),
+        Err(e) => {
+            warn!("Refresh error: {}", e);
+            Err((
+                StatusCode::INTERNAL_SERVER_ERROR,
+                "Server error".to_string(),
+            ))
+        }
+    }
+}

src/auth.rs

@@ -13,7 +13,7 @@ use ring::rand::{SecureRandom, SystemRandom};
 use serde::{Deserialize, Serialize};
 use sqlite::Connection;
 
-use crate::{util, AppState};
+use crate::{auth::TokenKind, util, AppState};
 
 #[derive(Deserialize, Clone)]
 pub struct CookieConfig {
@@ -69,7 +69,7 @@ async fn get_cookie(
     assert!(app.is_tls);
 
     let db = app.db.lock().await;
-    let account_id = match util::validate_authed_request(&headers) {
+    let account_id = match util::validate_authed_request(&headers, TokenKind::Session) {
         Ok(id) => id,
         Err(e) => return Err((StatusCode::UNAUTHORIZED, e)),
     };

src/cookie.rs

@@ -2,7 +2,7 @@ use axum::http::HeaderMap;
 use log::info;
 use sqlite::{Connection, State};
 
-use crate::auth;
+use crate::auth::{self, TokenKind};
 
 const MIN_DATABASE_VERSION: i64 = 6;
 
@@ -23,6 +23,12 @@ pub fn version_to_string(version: usize) -> String {
     format!("{}.{}.{}", major, minor, patch)
 }
 
+pub fn get_subroute(route: &str, subroute: &str) -> String {
+    let route_noslash = route.trim_end_matches('/');
+    let subroute_noslash = subroute.trim_start_matches('/');
+    format!("{}/{}", route_noslash, subroute_noslash)
+}
+
 pub fn connect_to_db(path: &str) -> Connection {
     const QUERY: &str = "SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='Meta';";
     const VERSION_QUERY: &str = "SELECT Value FROM Meta WHERE Key = 'DatabaseVersion';";
@@ -84,7 +90,7 @@ pub fn parse_csv(data: &str) -> Vec<Vec<String>> {
         .collect()
 }
 
-pub fn validate_authed_request(headers: &HeaderMap) -> Result<i64, String> {
+pub fn validate_authed_request(headers: &HeaderMap, kind: TokenKind) -> Result<i64, String> {
     let auth_header = headers.get("authorization").ok_or("No auth header")?;
     // auth header uses the Bearer scheme
     let parts: Vec<&str> = auth_header
@@ -96,7 +102,7 @@ pub fn validate_authed_request(headers: &HeaderMap) -> Result<i64, String> {
         return Err("Invalid auth header".to_string());
     }
     let token = parts[1];
-    auth::validate_jwt(token)
+    auth::validate_jwt(token, kind)
 }
 
 pub fn find_account(db: &Connection, username: &str) -> Option<Account> {