vtse/vtse-server/src/state.rs

use crate::user::{ApiKey, Password, Username};
use sodiumoxide::crypto::pwhash::argon2id13::{pwhash_verify, HashedPassword};
use sqlx::{query, PgPool};
use thiserror::Error;
use vtse_common::{
    net::{ServerResponse, UserError},
    stock::{Stock, StockName},
};

#[derive(Error, Debug)]
pub(crate) enum StateError {
    #[error("Was not in the correct state")]
    WrongState,
    #[error("Got SQLx error: {0}`")]
    Database(#[from] sqlx::Error),
    #[error("Failed to hash password")]
    PasswordHash,
}

type OperationResult = Result<ServerResponse, StateError>;

#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Debug)]
pub(crate) enum AppState {
    Unauthorized,
    Authorized { user_id: i32 },
}

/// Helper functions
impl AppState {
    pub(crate) fn new() -> Self {
        Self::Unauthorized
    }

    fn assert_state(&self, expected_state: AppState) -> Result<(), StateError> {
        if *self != expected_state {
            Err(StateError::WrongState)
        } else {
            Ok(())
        }
    }
}

/// Query operations impl
impl AppState {
    pub(crate) async fn stock_info(&self, stock_name: StockName, pool: &PgPool) -> OperationResult {
        let stock = query!(
            "SELECT name, symbol, description, price FROM stocks WHERE name = $1::text",
            stock_name as StockName
        )
        .fetch_one(pool)
        .await?;

        Ok(ServerResponse::StockInfo(Stock::new(
            stock.name,
            stock.symbol,
            stock.description,
            stock.price,
        )))
    }
}

/// User operation implementation
impl AppState {
    pub(crate) async fn login(&mut self, api_key: ApiKey, pool: &PgPool) -> OperationResult {
        self.assert_state(AppState::Unauthorized)?;
        let user_id = {
            let query = query!(
                "SELECT (user_id) FROM user_api_keys
                JOIN api_keys ON user_api_keys.key_id = api_keys.key_id
                WHERE key = $1",
                api_key.0
            )
            .fetch_optional(pool)
            .await?;
            match query {
                Some(id) => id.user_id,
                None => return Ok(ServerResponse::UserError(UserError::InvalidApiKey)),
            }
        };
        *self = AppState::Authorized { user_id };
        Ok(ServerResponse::Success)
    }

    pub(crate) async fn register(
        &mut self,
        username: Username,
        password: Password,
        pool: &PgPool,
    ) -> OperationResult {
        self.assert_state(AppState::Unauthorized)?;
        query!(
            "INSERT INTO users
            (username, pwhash_data)
            VALUES ($1::varchar, $2::bytea)",
            username as Username,
            password.salt().map_err(|_| StateError::PasswordHash)?.0,
        )
        .execute(pool)
        .await?;
        Ok(ServerResponse::Success)
    }

    pub(crate) async fn generate_api_key(
        &self,
        username: Username,
        password: Password,
        pool: &PgPool,
    ) -> OperationResult {
        self.assert_state(AppState::Unauthorized)?;
        match self.validate_password(&username, &password, pool).await {
            Ok(None) => Ok(ServerResponse::UserError(UserError::InvalidPassword)),
            Ok(Some(user_id)) => self.create_api_key(user_id, pool).await,
            Err(_) => todo!(),
        }
    }

    async fn create_api_key(&self, user_id: i32, pool: &PgPool) -> OperationResult {
        let mut transaction = pool.begin().await?;
        let api_key = uuid::Uuid::new_v4();
        let api_key_id = query!(
            "INSERT INTO api_keys (key)
            VALUES ($1)
            RETURNING api_keys.key_id;",
            &api_key,
        )
        .fetch_one(&mut transaction)
        .await?
        .key_id;

        query!(
            "INSERT INTO user_api_keys (user_id, key_id)
            VALUES ($1, $2)",
            user_id,
            api_key_id
        )
        .execute(&mut transaction)
        .await?;

        transaction.commit().await?;

        Ok(ServerResponse::NewApiKey(api_key))
    }

    async fn validate_password(
        &self,
        username: &Username,
        password: &Password,
        pool: &PgPool,
    ) -> Result<Option<i32>, StateError> {
        let result = query!(
            "SELECT user_id, pwhash_data FROM users
            WHERE username = $1::varchar",
            username as &Username
        )
        .fetch_one(pool)
        .await?;

        let login_attempt = HashedPassword::from_slice(&result.pwhash_data)
            .map(|pw| pwhash_verify(&pw, password.as_bytes()))
            .unwrap();

        if login_attempt {
            Ok(Some(result.user_id))
        } else {
            Ok(None)
        }
    }
}

/// Market operation implementation
impl AppState {
    pub(crate) fn buy(&self, stock_name: StockName, pool: &PgPool) -> OperationResult {
        if !matches!(self, Self::Authorized{..}) {
            return Err(StateError::WrongState);
        }
        todo!()
    }

    pub(crate) fn sell(&self, stock_name: StockName, pool: &PgPool) -> OperationResult {
        if !matches!(self, Self::Authorized{..}) {
            return Err(StateError::WrongState);
        }
        todo!()
    }
}

impl Default for AppState {
    fn default() -> Self {
        Self::new()
    }
}