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) -> Result, 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) -> Result, 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::(0) .map_err(sql_reject)? .parse::() .map_err(|e| reject::custom(BrevoReject::UriParseError(e)))?, ))), None => Err(reject::not_found()), } } pub async fn submit(form_data: SubmitForm, brevo: Arc) -> Result, 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::(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, err: Rejection, ) -> Result, 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, err: Rejection, ) -> Result, Rejection> { if err.is_not_found() || err.find::().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::() { 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::().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) -> 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 {}