asklyphe/authservice/src/db/session.rs

/*
 * authservice db/session.rs
 * - pgsql session functions
 *
 * Copyright (C) 2025 Real Microsoft, LLC
 *
 * This program is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, version 3.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License along with this program. If not, see <https://www.gnu.org/licenses/>.
*/

use chrono::{DateTime, Duration, TimeDelta, Utc};
use log::error;
use rand::Rng;
use sea_orm::ActiveValue::Set;
use sea_orm::{ActiveModelTrait, ColumnTrait, DatabaseConnection, EntityTrait, IntoActiveModel, QueryFilter};
use serde::{Deserialize, Serialize};
use entity::user_session;
use crate::db::{genid, user};
use crate::db::user::FetchUserError;

// todo: make this configurable
// 10 days
pub const SESSION_TTL_SECONDS: u64 = 10 * 24 * 60 * 60;

fn generate_session_token() -> String {
    rand::thread_rng()
        .sample_iter(&rand::distributions::Alphanumeric)
        .take(64)
        .map(char::from)
        .collect()
}

#[derive(Debug)]
pub enum CreateSessionError {
    UserDoesntExist,
    DatabaseError,
}

pub async fn create_session(db: &DatabaseConnection, user_id: String, ip_address: String) -> Result<String, CreateSessionError> {
    user::check_if_user_exists(db, user_id.to_string()).await.map_err(|e| match e {
        FetchUserError::UserNotFound => { CreateSessionError::UserDoesntExist }
        FetchUserError::DatabaseError => { CreateSessionError::DatabaseError }
    })?;

    let display_id = genid();
    let token = generate_session_token();
    let now = Utc::now();
    let death = now + Duration::try_seconds(SESSION_TTL_SECONDS as i64).expect("BAD SESSION TTL SECONDS");

    let record = user_session::ActiveModel {
        id: Set(display_id.clone()),
        user: Set(user_id),
        display_id: Set(display_id),
        token: Set(token.clone()),
        flags: Set(0),
        ip_address: Set(Some(ip_address)),
        created_on: Set(now.naive_utc()),
        death_date: Set(death.naive_utc()),
    };
    
    user_session::Entity::insert(record).exec(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE CREATESESSION: {e}");
        CreateSessionError::DatabaseError
    })?;

    Ok(token)
}

#[derive(Debug)]
pub enum DeleteSessionError {
    /// session doesn't exist
    SessionDoesntExist,
    DatabaseError,
}

/// deletes the given session from the session table, logging out the user
/// returns an error if the user or session doesn't exist, will delete the session regardless of
/// if it is expired
pub async fn delete_session(db: &DatabaseConnection, token: String) -> Result<(), DeleteSessionError> {
    match check_if_session_died(db, token.clone(), false, false).await {
        Ok(_) => {}
        Err(e) => match e {
            FetchSessionError::SessionInvalid => {}
            FetchSessionError::SessionDoesntExist => { return Err(DeleteSessionError::SessionDoesntExist); }
            FetchSessionError::DatabaseError => { return Err(DeleteSessionError::DatabaseError); }
        }
    }
    
    user_session::Entity::delete(user_session::Entity::find().filter(user_session::Column::Token.eq(&token)).one(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE DELETESESSION (query): {e}");
        DeleteSessionError::DatabaseError
    })?.ok_or(DeleteSessionError::SessionDoesntExist)?.into_active_model()).exec(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE DELETESESSION (delete): {e}");
        DeleteSessionError::DatabaseError
    })?;

    Ok(())
}

#[derive(Debug)]
pub enum FetchSessionError {
    /// death date reached
    SessionInvalid,
    /// session doesn't exist
    SessionDoesntExist,
    DatabaseError,
}

pub async fn check_if_session_died(db: &DatabaseConnection, token: String, extend_life: bool, rm_session: bool) -> Result<(), FetchSessionError> {
    let record = user_session::Entity::find().filter(user_session::Column::Token.eq(&token)).one(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE CHECKSESSIONDIED (query): {e}");
        FetchSessionError::DatabaseError
    })?.ok_or(FetchSessionError::SessionDoesntExist)?;
    
    let now = Utc::now();
    // is now bigger than death date?
    if record.death_date.signed_duration_since(now.naive_utc()) <= TimeDelta::nanoseconds(0) {
        // died
        if rm_session {
            user_session::Entity::delete(record.into_active_model()).exec(db).await.map_err(|e| {
                error!("DATABASE ERROR WHILE CHECKSESSIONDIED (delete): {e}");
                FetchSessionError::DatabaseError
            })?;
        }
        return Err(FetchSessionError::SessionInvalid);
    }

    if extend_life {
        let death = now + Duration::try_seconds(SESSION_TTL_SECONDS as i64).expect("BAD SESSION TTL SECONDS");
        let mut record = record.into_active_model();
        record.death_date = Set(death.naive_utc());
        record.update(db).await.map_err(|e| {
            error!("DATABASE ERROR WHILE CHECKSESSIONDIED (update): {e}");
            FetchSessionError::DatabaseError
        })?;
    }

    Ok(())
}

pub async fn get_session_death_date(db: &DatabaseConnection, token: String) -> Result<DateTime<Utc>, FetchSessionError> {
    check_if_session_died(db, token.clone(), false, true).await?;
    Ok(user_session::Entity::find().filter(user_session::Column::Token.eq(token)).one(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE GETSESSIONDEATHDATE: {e}");
        FetchSessionError::DatabaseError
    })?.ok_or(FetchSessionError::SessionDoesntExist)?.death_date.and_utc())
}

pub async fn get_userid_from_session(db: &DatabaseConnection, token: String, extend_life: bool) -> Result<String, FetchSessionError> {
    check_if_session_died(db, token.clone(), extend_life, true).await?;
    Ok(user_session::Entity::find().filter(user_session::Column::Token.eq(&token)).one(db).await.map_err(|e| {
        error!("DATABASE ERROR WHILE SESSION2USERID: {e}");
        FetchSessionError::DatabaseError
    })?.ok_or(FetchSessionError::SessionDoesntExist)?.user)
}