rvault/src/auth/token.rs

use crate::storage::DbPool;
use axum::extract::State;
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use axum::routing::post;
use axum::{Json, Router};
use log::error;
use rand::{Rng, distributions::Alphanumeric};
use serde::{Deserialize, Serialize};
use sqlx::Error;
use uuid::Uuid;

#[derive(Debug, Serialize)]
pub struct IdentityDTO {
    id: String,
    name: String,
}

#[derive(Debug)]
pub struct TokenDTO {
    key: String,
    id: String,
    identity_id: Option<String>,
    parent_id: Option<String>,
    expiry: Option<i64>,
}

#[derive(Debug)]
pub struct TokenRoleMembershipDTO {
    role_name: String,
    token_id: String,
}

/// Represents a request body for the `/auth/token/lookup` endpoint.
#[derive(Deserialize)]
struct RequestBodyPostLookup {
    token: String,
}

/// Represents the response body for the `/auth/token/lookup` endpoint.
#[derive(Serialize)]
struct TokenLookupResponse {
    id: String,
    type_name: String,
    roles: Vec<String>,
}

/// Represents an error response for the API.
#[derive(Serialize)]
struct ErrorResponse {
    error: String,
}

/// Generates a random string of the specified length using alphanumeric characters.
// TODO: Make string generation secure
fn get_random_string(len: usize) -> String {
    rand::thread_rng()
        .sample_iter(&Alphanumeric)
        .take(len)
        .map(char::from)
        .collect()
}

/// Creates a root token if none exists in the database.
/// Returns true if a new root token was created, false if one already exists.
pub async fn create_root_token_if_none_exist(pool: &DbPool) -> bool {
    // Check if a root token already exists
    let exists = sqlx::query!(
        r#"SELECT service_token.* FROM service_token, service_token_role_membership
                          WHERE service_token.id = service_token_role_membership.token_id AND
                                service_token_role_membership.role_name = 'root'
                          LIMIT 1"#
    )
    .fetch_one(pool)
    .await
    .is_ok();
    if exists {
        return false;
    }
    // If no root token exists, create one
    let result = create_root_token(pool).await;
    if result.is_err() {
        let error = result.err().unwrap();
        // Log the error and panic
        error!("create_root_token failed: {error:?}");
        panic!("create_root_token failed: {error:?}");
    }
    // If successful, print the root token. This will only happen once.
    println!("\n\nYour root token is: {}", result.unwrap());
    println!("It will only be displayed once!\n\n");
    true
}

/// Creates a root token in the database.
async fn create_root_token(pool: &DbPool) -> Result<String, Error> {
    let id = Uuid::new_v4().to_string();
    let key = "s.".to_string() + &get_random_string(24);
    // Insert the root token into the database
    let result = sqlx::query!(r#"
                    INSERT INTO service_token (id, key) VALUES ($1, $2);
                    INSERT INTO service_token_role_membership (token_id, role_name) VALUES ($3, 'root');
            "#, id, key, id).execute(pool).await;
    // If the insert was successful, return the key
    if result.is_ok() {
        return Ok(key);
    }
    // Else, return the error
    Err(result.unwrap_err())
}

/// Gets the current time in seconds since unix epoch
fn get_time_as_int() -> i64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_secs() as i64
}

/// Gets the type of token. (The first character of the key always specifies the type)
fn get_token_type(token: &TokenDTO) -> Result<String, &str> {
    Ok(match token.key.clone().chars().next().unwrap_or('?') {
        's' => "service",
        'b' => "batch",
        'r' => "recovery",
        _ => {
            error!("Unsupported token type");
            return Err("Unsupported token type");
        }
    }
    .to_string())
}

/// Retrieves a token from the database using its key.
/// If the token is found and not expired, it returns the token.
/// Else, it returns an error.
pub async fn get_token_from_key(token_key: &str, pool: &DbPool) -> Result<TokenDTO, Error> {
    let time = get_time_as_int();
    sqlx::query_as!(
            TokenDTO,
            r#"SELECT * FROM 'service_token' WHERE key = $1 AND (expiry IS NULL OR expiry > $2) LIMIT 1"#,
            token_key, time).fetch_one(pool).await
}

/// Retrieves the roles associated with a given token from the database.
/// If the token does not exist, it returns an empty vector.
pub async fn get_roles_from_token(token: &TokenDTO, pool: &DbPool) -> Vec<String> {
    let result = sqlx::query_as!(
        TokenRoleMembershipDTO,
        r#"SELECT * FROM 'service_token_role_membership' WHERE token_id = $1"#,
        token.id
    )
    .fetch_all(pool)
    .await;
    result
        .unwrap_or(Vec::new())
        .iter()
        .map(|r| r.role_name.to_string())
        .collect()
}

/// Return a router, that may be used to route traffic to the corresponding handlers
pub fn token_auth_router(pool: DbPool) -> Router<DbPool> {
    Router::new()
        .route("/lookup", post(post_lookup))
        .with_state(pool)
}

/// Handles the `/auth/token/lookup` endpoint.
/// Retrieves the token and its associated roles from the database using the provided token key.
/// The output format does not yet match the openBao specification and is for testing only!
async fn post_lookup(
    State(pool): State<DbPool>,
    Json(body): Json<RequestBodyPostLookup>,
) -> Response {
    let token_str = body.token;
    // Validate the token string
    match get_token_from_key(&token_str, &pool).await {
        // If the token is found, retrieve its type and roles
        Ok(token) => {
            let type_name = get_token_type(&token).unwrap_or_else(|_| String::from("Unknown"));
            let roles = get_roles_from_token(&token, &pool).await;
            let resp = TokenLookupResponse {
                id: token.id,
                type_name,
                roles,
            };
            // Return the token information as a JSON response
            (StatusCode::OK, axum::Json(resp)).into_response()
        }
        // If the token is not found, return a 404 Not Found error
        Err(e) => {
            error!("Failed to retrieve token: {e:?}");
            let err = ErrorResponse {
                error: "Failed to retrieve token".to_string(),
            };
            (StatusCode::NOT_FOUND, axum::Json(err)).into_response()
        }
    }
}

//
// The following functions are placeholders for the various token-related operations.
//


async fn get_accessors() -> &'static str {
    todo!("not implemented")
}

async fn post_create() -> &'static str {
    todo!("not implemented")
}

async fn post_create_orphan() -> &'static str {
    todo!("not implemented")
}

async fn post_create_role() -> &'static str {
    todo!("not implemented")
}

async fn get_lookup() -> &'static str {
    todo!("not implemented")
}

async fn get_lookup_self() -> &'static str {
    todo!("not implemented")
}

async fn post_lookup_self() -> &'static str {
    todo!("not implemented")
}

async fn post_renew() -> &'static str {
    todo!("not implemented")
}

async fn post_renew_accessor() -> &'static str {
    todo!("not implemented")
}

async fn post_renew_self() -> &'static str {
    todo!("not implemented")
}

async fn post_revoke() -> &'static str {
    todo!("not implemented")
}

async fn post_revoke_accessor() -> &'static str {
    todo!("not implemented")
}

async fn post_revoke_orphan() -> &'static str {
    todo!("not implemented")
}

async fn post_revoke_self() -> &'static str {
    todo!("not implemented")
}

async fn get_roles() -> &'static str {
    todo!("not implemented")
}

async fn get_role_by_name() -> &'static str {
    todo!("not implemented")
}

async fn post_role_by_name() -> &'static str {
    todo!("not implemented")
}

async fn delete_role_by_name() -> &'static str {
    todo!("not implemented")
}

async fn post_tidy() -> &'static str {
    todo!("not implemented")
}