FrostByte/server/src/db.rs

use crate::types::{Comment, Post, User};
use argon2::{
    password_hash::{rand_core::OsRng, SaltString},
    Argon2, PasswordHasher, PasswordVerifier,
};
use log::{info, warn};
use sqlx::PgPool;

pub async fn db_new_comment(
    pool: &PgPool,
    parent_post_id: i64,
    parent_comment_id: Option<i64>,
    user_id: i64,
    content: &str,
) -> bool {
    let insert_query = sqlx::query!(
        "INSERT INTO comments (parent_post_id, parent_comment_id, author_user_id, content) VALUES ($1, $2, $3, $4)",
        parent_post_id,
        parent_comment_id.unwrap_or(-1), // This is a bit of a hack
        user_id,
        content
    )
    .execute(pool)
    .await;

    if insert_query.is_err() {
        let s = insert_query.err().unwrap();
        warn!("Error inserting comment into database: {}", s);
        return false;
    }

    true
}

pub async fn db_get_comments(
    pool: &PgPool,
    parent_post_id: i64,
    limit: i64,
    offset: i64,
) -> Vec<Comment> {
    sqlx::query_as!(
        Comment,
        "SELECT * FROM comments WHERE parent_post_id = $1 ORDER BY created_at DESC LIMIT $2 OFFSET $3",
        parent_post_id,
        limit,
        offset
    )
    .fetch_all(pool)
    .await
    .unwrap()
}

/// Gets the latest posts from the database, ordered by created_at
pub async fn db_get_latest_posts(pool: &PgPool, limit: i64, offset: i64) -> Vec<Post> {
    sqlx::query_as!(
        Post,
        "SELECT * FROM posts ORDER BY created_at DESC LIMIT $1 OFFSET $2",
        limit,
        offset
    )
    .fetch_all(pool)
    .await
    .unwrap()
}

/// Gets the post with id from the database
pub async fn db_get_post(id: i64, pool: &PgPool) -> Option<Post> {
    sqlx::query_as!(Post, "SELECT * FROM posts WHERE id = $1", id)
        .fetch_one(pool)
        .await
        .ok()
}

/// Inserts a new post to the database
pub async fn db_new_post(userid: i64, content: &str, pool: &PgPool) -> Option<Post> {
    info!("User with id {} submitted a post", userid);

    let insert_query = sqlx::query!(
        "INSERT INTO posts (user_id, content) VALUES ($1, $2)",
        userid,
        content
    )
    .execute(pool)
    .await;

    if insert_query.is_err() {
        let s = insert_query.err().unwrap();
        warn!("Error inserting post into database: {}", s);
        return None;
    }

    // Dips into the database to get the post we just inserted
    let post = sqlx::query_as!(
        Post,
        "SELECT * FROM posts WHERE id = (SELECT MAX(id) FROM posts)"
    )
    .fetch_one(pool)
    .await
    .ok()?;

    Some(post)
}

/// Checks if the user exists in the database
pub async fn db_user_exists(username: String, pool: &PgPool) -> bool {
    let exists = sqlx::query!("SELECT username FROM users WHERE username = $1", username)
        .fetch_one(pool)
        .await
        .ok()
        .map(|row| row.username);

    exists.is_some()
}

/// Checks if the user exists and if the password is correct
pub async fn db_user_login(username: String, password: String, pool: &PgPool) -> Option<User> {
    let username = username.clone();
    let user = sqlx::query_as!(User, "SELECT * FROM users WHERE username = $1", username)
        .fetch_one(pool)
        .await
        .ok()?;

    let phc_password = user.password.clone();
    let phc_password = match argon2::PasswordHash::new(&phc_password) {
        Ok(phc_password) => phc_password,
        Err(_) => {
            warn!(
                "Invalid hash for user {} fetched from database (not a valid PHC string)",
                username
            );
            return None;
        }
    };

    let argon2 = Argon2::default();
    let password = password.as_bytes();

    match argon2.verify_password(password, &phc_password) {
        Ok(_) => Some(user),
        Err(_) => None,
    }
}

/// Creates a new user if the username is not already taken
pub async fn db_new_user(username: String, password: String, pool: &PgPool) -> Option<User> {
    // First check if the user already exists
    match db_user_exists(username.clone(), pool).await {
        true => {
            warn!("User \"{}\" already exists", username);
            return None;
        }
        false => {}
    }

    // Unwrapping here because if this fails, we have a serious problem
    let phc_hash = Argon2::default()
        .hash_password(password.as_bytes(), &SaltString::generate(&mut OsRng))
        .unwrap()
        .to_string();

    // Insert our new user into the database
    let insert_query = sqlx::query!(
        "INSERT INTO users (username, password) VALUES ($1, $2)",
        username,
        phc_hash
    )
    .execute(pool)
    .await;

    match insert_query {
        Ok(_) => {
            info!("User: {} registered", username);
            let user = sqlx::query_as!(User, "SELECT * FROM users WHERE username = $1", username)
                .fetch_one(pool)
                .await
                .ok()?;

            Some(user)
        }
        Err(e) => {
            warn!("Error inserting user into database: {}", e);
            return None;
        }
    }
}