pastabble/src/main.rs

use std::{env, fs::{self, File}, thread};

use chrono::Utc;
use kv::{Config, Store, Json, Bucket, Value};
use paste::Paste;
use rand::{Rng, distributions::Alphanumeric};
use rouille::{Response, router, try_or_400, post_input, Request, Server};
use signal_hook::{iterator::Signals, consts::{SIGINT, SIGTERM}};

mod paste;

// Generate random key function
fn random_string() -> String {
    rand::thread_rng()
        .sample_iter(&Alphanumeric)
        .take(5)
        .map(char::from)
        .collect()
}

// Generate key using provided key or generate new one
fn generate_key<T : Value>(key: Option<String>, store: &Bucket<String, T>) -> String {
    let mut new_key = key.unwrap_or(random_string());

    loop {
        if !store.contains(&new_key).unwrap() {
            break new_key;
        }
        new_key = random_string();
    }
}

fn get_query(req: &Request, name: &str) -> Option<String> {
    let params: Vec<&str> = req.raw_query_string().split('&').collect();
    let param = params.iter().find(|p|p.starts_with(&format!("{}=", name)));
    match param {
        Some(q) => q.split('=').last(),
        None => None,
    }.map(|d|d.to_string())
}

fn main() {
    // Get settings
    let port: u16 = env::var("PORT").unwrap_or(String::from("8080")).parse().expect("Failed to parse PORT variable");
    let data_dir: String = env::var("DATA_DIR").unwrap_or(String::from("./data"));
    let source_dir: String = env::var("WWW_DIR").unwrap_or(String::from("./wwwroot"));
    let prefix: Option<String> = env::var("PREFIX").ok();

    let config = Config::new(data_dir);
    let store = Store::new(config).expect("Failed to initialize database store");
    let about = fs::read_to_string("./about.html").expect("about.html not found!!");

    // Get pastes bucket
    let pastes = store.bucket::<String, Json<Paste>>(Some("pastes"))
        .expect("Failed to open pastes bucket");

    // Get links bucket
    let links = store.bucket::<String, String>(Some("links"))
        .expect("Failed to open links bucket");

    
    // Create server
    let server = Server::new(format!("0.0.0.0:{}", port), move |req| {
        if req.url().ends_with(".html") || req.url().ends_with(".css") || req.url().ends_with(".js") {
            let asset_response = rouille::match_assets(req, &source_dir);
            if asset_response.is_success() {
                return asset_response;
            }
        }

        router!(req,
            (GET) (/) => {
                match File::open(format!("{}/index.html", &source_dir)) {
                    Ok(file) => {
                        Response::from_file("text/html", file)
                    },
                    Err(_) => {
                        Response::html(&about)
                    }
                }
            },
            (GET) (/about) => {
                Response::html(&about)
            },
            (GET) (/to/{id: String}) => {
                match links.get(&id).expect("Failed to access links DB") {
                    Some(lnk) => {
                        Response::redirect_301(lnk)
                    },
                    None => {
                        Response::text(format!("Redirect with ID '{}' not found", &id))
                            .with_status_code(404)
                    }
                }
            },
            (POST) (/to/{id: String}) => { register_link(&links, req, Some(id), prefix.clone()) },
            (POST) (/to) => { register_link(&links, req, None, prefix.clone()) },
            (GET) (/{id: String}) => {
                // Get note from database, if exists
                match pastes.get(&id).expect("Failed to access pastes DB") {
                    Some(paste) => {
                        // If details, return JSON. Else, return text
                        if get_query(req, "details").is_some() {
                            Response::json(&paste.0)
                        } else {
                            Response::text(paste.0.content)
                                .with_additional_header("Access-Control-Allow-Origin", "*")
                        }
                    },
                    None => {
                        Response::text(format!("Note with ID '{}' not found", &id))
                            .with_status_code(404)
                    }
                }
            },
            (POST) (/{id: String}) => { 
                let id = if id.is_empty() { None } else { Some(id) };
                register_paste(&pastes, req, id, prefix.clone()) 
            },
            _ => Response::empty_404()
        )
    }).expect("Failed to start server");

    // Start server (run() or ::start_server is not able to handle graceful shutdown, hence stoppable() or poll() in loop)
    println!("Started PASTABBLE server on {:?}", server.server_addr());
    let (handle, stop) = server.stoppable();

    // Listen for signals
    let mut signals = Signals::new(&[SIGINT, SIGTERM]).unwrap();
    thread::spawn(move || {
        for sig in signals.forever() {
            match sig {
                _ => {
                    println!("Received stop signal {:?}, stopping server...", sig);
                    stop.send(()).expect("Failed to stop server");
                }
            }
        }
    });

    // Join server thread
    handle.join().unwrap();
}

// Register paste handler
fn register_paste(pastes: &Bucket<String, Json<Paste>>, req: &Request, id: Option<String>, prefix: Option<String>) -> Response {

    // Try read body from form data, and if not present from request body
    let body = match post_input!(req, {
        content: String
    }) {
        Ok(f) => {
            f.content
        },
        Err(_) => {
            try_or_400!(rouille::input::plain_text_body(req))
        }
    };

    // If lang param is present in query, extract it
    let language = get_query(req, "lang");

    // Create and save new note
    let paste = Json(Paste { 
        content: body,
        expires: None, 
        language, 
        created: Utc::now()
    });

    let key = generate_key(id, &pastes);
    pastes.set(&key, &paste).expect("Failed to save note");
    pastes.flush().expect("Failed to save paste to database");

    // Return url with key
    let url = if prefix.is_some() {
        format!("{}/{}", prefix.unwrap(), key)
    } else {
        key
    };

    Response::text(url)
}

// Register link route handler
fn register_link(links: &Bucket<String,String>, req: &Request, id: Option<String>, prefix: Option<String>) -> Response {

    // Try read body from form data, and if not present from request body
    let link = match post_input!(req, {
        link: String
    }) {
        Ok(f) => {
            f.link
        },
        Err(_) => {
            try_or_400!(rouille::input::plain_text_body(req))
        }
    }.trim().to_string();

    let key = generate_key(id, &links);
    links.set(&key, &link).expect("Failed to save link");
    links.flush().expect("Failed to save link to database");

    // Return url with key
    let url = if prefix.is_some() {
        format!("{}/to/{}", prefix.unwrap(), key)
    } else {
        key
    };

    Response::text(url)
}