shoutdotdev/src/v0_router.rs

use anyhow::Context;
use axum::{
    extract::Query,
    response::{IntoResponse, Json},
    routing::get,
    Router,
};
use chrono::TimeDelta;
use diesel::{dsl::insert_into, prelude::*, update};
use serde::Deserialize;
use serde_json::json;
use uuid::Uuid;

use crate::{
    api_keys::ApiKey,
    app_error::AppError,
    app_state::{AppState, DbConn},
    channels::Channel,
    governors::Governor,
    projects::Project,
    schema::{api_keys, governors, messages, projects},
};

const TEAM_GOVERNOR_DEFAULT_WINDOW_SIZE_SEC: i64 = 300;
const TEAM_GOVERNOR_DEFAULT_MAX_COUNT: i32 = 50;

pub fn new_router(state: AppState) -> Router<AppState> {
    Router::new().route("/say", get(say_get)).with_state(state)
}

#[derive(Deserialize)]
struct SayQuery {
    key: Uuid,
    project: String,
    message: String,
}

async fn say_get(
    DbConn(db_conn): DbConn,
    Query(query): Query<SayQuery>,
) -> Result<impl IntoResponse, AppError> {
    // TODO: do some validation of message contents
    let api_key = {
        let query_key = query.key.clone();
        db_conn
            .interact::<_, Result<ApiKey, AppError>>(move |conn| {
                update(api_keys::table.filter(ApiKey::with_id(query_key)))
                    .set(api_keys::last_used_at.eq(diesel::dsl::now))
                    .returning(ApiKey::as_returning())
                    .get_result(conn)
                    .optional()
                    .context("failed to get API key")?
                    .ok_or(AppError::ForbiddenError("Key not accepted.".to_string()))
            })
            .await
            .unwrap()?
    };

    let project = {
        let project_name = query.project.to_lowercase();
        db_conn
            .interact::<_, Result<Project, AppError>>(move |conn| {
                conn.transaction(move |conn| {
                    Ok(
                        match Project::all()
                            .filter(Project::with_team(api_key.team_id))
                            .filter(Project::with_name(project_name.clone()))
                            .first(conn)
                            .optional()
                            .context("failed to load project")?
                        {
                            Some(project) => project,
                            None => insert_into(projects::table)
                                .values((
                                    projects::id.eq(Uuid::now_v7()),
                                    projects::team_id.eq(api_key.team_id),
                                    projects::name.eq(project_name),
                                ))
                                .get_result(conn)
                                .context("failed to insert project")?,
                        },
                    )
                })
            })
            .await
            .unwrap()?
    };

    let team_governor = {
        let team_id = project.team_id.clone();
        db_conn
            .interact::<_, Result<Governor, AppError>>(move |conn| {
                // TODO: extract this logic to a method in crate::governors,
                // and create governor proactively on team creation
                match Governor::all()
                    .filter(Governor::with_team(team_id.clone()))
                    .filter(Governor::with_project(None))
                    .first(conn)
                {
                    diesel::QueryResult::Ok(governor) => Ok(governor),
                    diesel::QueryResult::Err(diesel::result::Error::NotFound) => {
                        // Lazily initialize governor
                        Ok(diesel::insert_into(governors::table)
                            .values((
                                governors::team_id.eq(team_id),
                                governors::id.eq(Uuid::now_v7()),
                                governors::project_id.eq(None as Option<Uuid>),
                                governors::window_size
                                    .eq(TimeDelta::seconds(TEAM_GOVERNOR_DEFAULT_WINDOW_SIZE_SEC)),
                                governors::max_count.eq(TEAM_GOVERNOR_DEFAULT_MAX_COUNT),
                            ))
                            .get_result(conn)?)
                    }
                    diesel::QueryResult::Err(err) => Err(err.into()),
                }
            })
            .await
            .unwrap()?
    };

    if db_conn
        .interact::<_, Result<Option<_>, anyhow::Error>>(move |conn| {
            team_governor.create_entry(conn)
        })
        .await
        .unwrap()?
        .is_none()
    {
        return Err(AppError::TooManyRequestsError(
            "team rate limit exceeded".to_string(),
        ));
    }

    let selected_channels = {
        let project = project.clone();
        db_conn
            .interact::<_, Result<Vec<Channel>, AppError>>(move |conn| {
                Ok(project
                    .selected_channels()
                    .load(conn)
                    .context("failed to load selected channels")?)
            })
            .await
            .unwrap()?
    };

    {
        let selected_channels = selected_channels.clone();
        db_conn
            .interact::<_, Result<_, AppError>>(move |conn| {
                for channel in selected_channels {
                    insert_into(messages::table)
                        .values((
                            messages::id.eq(Uuid::now_v7()),
                            messages::channel_id.eq(&channel.id),
                            messages::project_id.eq(&project.id),
                            messages::message.eq(&query.message),
                        ))
                        .execute(conn)?;
                }
                Ok(())
            })
            .await
            .unwrap()?;
    }
    tracing::debug!("queued {} messages", selected_channels.len());

    Ok(Json(json!({ "ok": true })))
}