use std::path::PathBuf; use serde::Serialize; use sqlx::{Executor, SqlitePool}; use warp::{ filters::BoxedFilter, hyper::{self, StatusCode}, reply::Response, Filter, Rejection, Reply, }; // if we're making a debug build, we don't want to expand the migrate! macro at all #[inline] #[cfg(debug_assertions)] async fn migrate(_db: &SqlitePool) {} #[inline] #[cfg(not(debug_assertions))] async fn migrate(db: &SqlitePool) { let _ = sqlx::migrate!().run(db).await; } #[derive(Serialize)] struct Redirect { path: String, target_location: String, redirect_code: i64, } impl Redirect { fn status_code(&self) -> hyper::StatusCode { hyper::StatusCode::from_u16(self.redirect_code as u16) .expect("Status codes in database should always be valid") } } #[derive(serde::Deserialize)] struct AddRedirectForm { path: Option, target: String, status_code: Option, } fn auth_valid(authorization: &str) -> bool { if let Some(token) = authorization.strip_prefix("Bearer ") { if std::env::var("REDIRECTS_AUTH_KEY").ok().as_deref() == Some(token) { return true; } } false } fn json_error_response(message: &str, code: StatusCode) -> impl Reply { warp::reply::with_status( warp::reply::json(&serde_json::json!({ "error": message })), code, ) } async fn add_redirect( db: SqlitePool, add_redirect_form: AddRedirectForm, authorization: String, ) -> Result { if !auth_valid(&authorization) { return Ok( json_error_response("Invalid authorization", StatusCode::UNAUTHORIZED).into_response(), ); } async fn create_redirect(db: SqlitePool, redir: Redirect) -> sqlx::Result<()> { let mut conn = db.acquire().await?; let query = sqlx::query!( "INSERT OR REPLACE INTO redirects VALUES (?, ?, ?)", redir.path, redir.target_location, redir.redirect_code ); let _ = conn.execute(query).await?; Ok(()) } let path = add_redirect_form .path .unwrap_or_else(|| todo!("generate random string")); let redir = Redirect { path, target_location: add_redirect_form.target, redirect_code: add_redirect_form.status_code.map(i64::from).unwrap_or(302), }; Ok(match create_redirect(db, redir).await { Ok(_) => warp::reply::json(&serde_json::json!({"success": true})).into_response(), Err(_) => { json_error_response("Internal error", StatusCode::INTERNAL_SERVER_ERROR).into_response() } }) } async fn list_redirects(db: SqlitePool, authorization: String) -> Result { if !auth_valid(&authorization) { return Err(warp::reject()); } async fn get_redirects(db: SqlitePool) -> sqlx::Result> { let mut conn = db.acquire().await?; let query = sqlx::query_as!(Redirect, "SELECT * FROM redirects"); query.fetch_all(&mut conn).await } match get_redirects(db).await { Ok(redirects) => Ok(warp::reply::json(&redirects)), Err(_) => Err(warp::reject()), } } async fn lookup_redirect(path: String, db: SqlitePool) -> Result { async fn find_redirect(db: SqlitePool, path: &str) -> sqlx::Result { let mut conn = db.acquire().await?; let query = sqlx::query_as!(Redirect, "SELECT * FROM redirects WHERE path = ?", path); query.fetch_one(&mut conn).await } match find_redirect(db, &path).await { Ok(redir) => Ok(warp::reply::with_header( redir.status_code(), hyper::header::LOCATION, redir.target_location, )), Err(_) => Err(warp::reject()), } } pub async fn routes(static_dir: &str) -> BoxedFilter<(Response,)> { let database_uri = std::env!("DATABASE_URL"); let db = SqlitePool::connect(database_uri) .await .unwrap_or_else(|_| panic!("Failed to open SQLite DB: {}", database_uri)); migrate(&db).await; let shorten_html = { let mut path = PathBuf::from(static_dir); path.push("shorten.html"); path }; let redirects_route = warp::path!("redirects.json").and({ let db = db.clone(); warp::get() .map(move || db.clone()) .and(warp::header::("authorization")) .and_then(list_redirects) }); let shorten_route = warp::path!("shorten") .and(warp::get().and(warp::fs::file(shorten_html))) // TODO: Templating to list existing shortlinks? .or({ let db = db.clone(); warp::post() .map(move || db.clone()) .and(warp::body::form()) .and(warp::header::("authorization")) .and_then(add_redirect) }); let lookup_route = warp::path::param() .and(warp::path::end()) .and(warp::any().map(move || db.clone())) .and_then(lookup_redirect); shorten_route .or(redirects_route) .or(lookup_route) .map(Reply::into_response) .boxed() }