char.lt-old/src/redirects.rs

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<String>,
    target: String,
    status_code: Option<u16>,
}

#[derive(serde::Deserialize)]
struct DeleteRedirectForm {
    path: String,
}

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<Response, Rejection> {
    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(())
    }

    if !auth_valid(&authorization) {
        return Ok(
            json_error_response("Invalid authorization", StatusCode::UNAUTHORIZED).into_response(),
        );
    }

    let path = add_redirect_form.path.unwrap_or_else(|| {
        use rand::{thread_rng, Rng};

        const LEN: usize = 16;
        const CHARSET: &str = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";

        let mut rng = thread_rng();
        let mut random_string = String::with_capacity(LEN);
        for _ in 0..LEN {
            let idx = rng.gen_range(0..CHARSET.len());
            let c: &str = &CHARSET[idx..=idx + 1]; // we know the charset is ASCII so this is okay
            random_string.push_str(c);
        }

        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),
    };

    if let Err(_e) = create_redirect(db, redir).await {
        return Ok(
            json_error_response("Internal error", StatusCode::INTERNAL_SERVER_ERROR)
                .into_response(),
        );
    }

    Ok(warp::reply::json(&serde_json::json!({"success": true})).into_response())
}

async fn delete_redirect(
    db: SqlitePool,
    form: DeleteRedirectForm,
    authorization: String,
) -> Result<Response, Rejection> {
    async fn delete_redirect(db: SqlitePool, redir_path: &str) -> sqlx::Result<()> {
        let mut conn = db.acquire().await?;
        let query = sqlx::query!("DELETE FROM redirects WHERE path = ?", redir_path);
        let _ = conn.execute(query).await?;
        Ok(())
    }

    if !auth_valid(&authorization) {
        return Ok(
            json_error_response("Invalid authorization", StatusCode::UNAUTHORIZED).into_response(),
        );
    }

    if let Err(_e) = delete_redirect(db, &form.path).await {
        return Ok(
            json_error_response("Internal error", StatusCode::INTERNAL_SERVER_ERROR)
                .into_response(),
        );
    }

    Ok(warp::reply::json(&serde_json::json!({"success": true})).into_response())
}

async fn list_redirects(db: SqlitePool, authorization: String) -> Result<impl Reply, Rejection> {
    if !auth_valid(&authorization) {
        return Err(warp::reject());
    }

    async fn get_redirects(db: SqlitePool) -> sqlx::Result<Vec<Redirect>> {
        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<impl Reply, Rejection> {
    async fn find_redirect(db: SqlitePool, path: &str) -> sqlx::Result<Redirect> {
        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::<String>("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::<String>("authorization"))
                .and_then(add_redirect)
        })
        .or({
            let db = db.clone();
            warp::delete()
                .map(move || db.clone())
                .and(warp::body::form())
                .and(warp::header::<String>("authorization"))
                .and_then(delete_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()
}