brevo/src/handlers.rs

use crate::util::{gen_id, make_link, render_reject};
use serde::{Deserialize, Serialize};
use sqlx::Row;
use std::{convert::Infallible, sync::Arc};
use tera::Context;
use url::Url;
use warp::{
    body::BodyDeserializeError,
    http::{uri::InvalidUri, Uri},
    redirect,
    reject::{self, MethodNotAllowed, Reject},
    reply,
    Rejection,
    Reply,
};

use crate::{sqlargs, util::sql_reject, Brevo};

pub async fn index(brevo: Arc<Brevo>) -> Result<Box<dyn Reply>, Rejection> {
    let rendered = brevo.tera.render("index.html", &Context::new());
    match rendered {
        Err(e) => Err(reject::custom(BrevoReject::RenderFail(e))),
        Ok(r) => Ok(Box::new(reply::html(r))),
    }
}

pub async fn shortened(id: String, brevo: Arc<Brevo>) -> Result<Box<dyn Reply>, Rejection> {
    make_table(Arc::clone(&brevo)).await?;
    let url = sqlx::query_with(
        "
        SELECT url FROM urls
        WHERE id = ?
        ",
        sqlargs![&id],
    )
    .fetch_optional(&brevo.pool)
    .await
    .map_err(sql_reject)?;

    match url {
        Some(u) => Ok(Box::new(redirect::permanent(
            u.try_get::<String, _>(0)
                .map_err(sql_reject)?
                .parse::<Uri>()
                .map_err(|e| reject::custom(BrevoReject::UriParseError(e)))?,
        ))),
        None => Err(reject::not_found()),
    }
}

pub async fn submit(form_data: SubmitForm, brevo: Arc<Brevo>) -> Result<Box<dyn Reply>, Rejection> {
    make_table(Arc::clone(&brevo)).await?;
    let url = form_data.url.as_str();
    let id = sqlx::query_with(
        "
        SELECT id FROM urls
        WHERE url = ?
        ",
        sqlargs![url],
    )
    .fetch_optional(&brevo.pool)
    .await
    .map_err(sql_reject)?;

    if let Some(id) = id {
        let id = id.try_get::<String, _>(0).map_err(sql_reject)?;
        let rendered = brevo
            .tera
            .render(
                "success.html",
                &Context::from_serialize(SuccessContent {
                    link: make_link(Arc::clone(&brevo), &id)?,
                })
                .map_err(render_reject)?,
            )
            .map_err(render_reject)?;

        Ok(Box::new(reply::html(rendered)))
    } else {
        let id = gen_id(Arc::clone(&brevo)).await;
        sqlx::query_with(
            "
            INSERT INTO urls ( id, url ) VALUES
            ( ?, ? )
            ",
            sqlargs![&id, url],
        )
        .execute(&brevo.pool)
        .await
        .map_err(sql_reject)?;

        let rendered = brevo
            .tera
            .render(
                "success.html",
                &Context::from_serialize(SuccessContent {
                    link: make_link(Arc::clone(&brevo), &id)?,
                })
                .map_err(render_reject)?,
            )
            .map_err(render_reject)?;
        Ok(Box::new(reply::html(rendered)))
    }
}

pub async fn handle_reject(
    brevo: Arc<Brevo>,
    err: Rejection,
) -> Result<Box<dyn Reply>, Infallible> {
    match try_handle_reject(brevo, err).await {
        Ok(x) => Ok(x),
        Err(e) => {
            log::error!("Unrecoverable reply error occured! {:?}", e);
            Ok(Box::new(reply::html("ERROR!! see logs")))
        },
    }
}

pub async fn try_handle_reject(
    brevo: Arc<Brevo>,
    err: Rejection,
) -> Result<Box<dyn Reply>, Rejection> {
    if err.is_not_found() || err.find::<MethodNotAllowed>().is_some() {
        let rendered = brevo
            .tera
            .render("404.html", &Context::new())
            .map_err(render_reject)?;

        return Ok(Box::new(reply::html(rendered)));
    }

    if let Some(rej) = err.find::<BrevoReject>() {
        return match rej {
            BrevoReject::SqlError(err) => {
                log::error!("SQL error: {:?}", err);
                Ok(Box::new(reply::html("SQL error! see logs.")))
            },

            BrevoReject::RenderFail(err) => {
                log::error!("Template error: {:?}", err);
                Ok(Box::new(reply::html(
                    "Failed to render template! see logs.",
                )))
            },

            BrevoReject::UrlParseError(err) => {
                log::error!("URL parse error: {:?}", err);
                Ok(Box::new(reply::html("Not a valid URL!")))
            },

            BrevoReject::UriParseError(err) => {
                log::error!("URI parse error: {:?}", err);
                Ok(Box::new(reply::html("Not a valid URI!")))
            },
        };
    }

    if err.find::<BodyDeserializeError>().is_some() {
        return Ok(Box::new(reply::html(
            "Invalid data! Did you enter a valid URI?",
        )));
    }

    log::error!("unknown rejection: {:?}", err);
    Ok(Box::new(reply::html("Found unknown rejection! see logs")))
}

async fn make_table(brevo: Arc<Brevo>) -> Result<(), Rejection> {
    // can't parameterize VARCHAR len, but this should be safe
    sqlx::query(&format!(
        "
        CREATE TABLE IF NOT EXISTS urls (
            id VARCHAR({}) PRIMARY KEY,
            url TEXT NOT NULL,
            UNIQUE (
                url(512)
            )
        )                                    
        ",
        brevo.config.link_len
    ))
    .execute(&brevo.pool)
    .await
    .map_err(sql_reject)?;
    Ok(())
}

#[derive(Serialize)]
struct SuccessContent {
    link: String,
}

#[derive(Deserialize)]
pub struct SubmitForm {
    url: Url,
}

// TODO: implement reject handerls
#[derive(Debug)]
pub enum BrevoReject {
    RenderFail(tera::Error),
    SqlError(sqlx::Error),
    UrlParseError(url::ParseError),
    UriParseError(InvalidUri),
}

impl Reject for BrevoReject {}