pastabble/src/main.rs

use std::{env, fs, 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 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 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| {
        router!(req,
            (GET) (/) => {
                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)) },
            (POST) (/to) => { register_link(&links, req, None) },
            (GET) (/{id: String}) => {
                // Get note from database, if exists
                match pastes.get(&id).expect("Failed to access pastes DB") {
                    Some(paste) => {
                        Response::text(paste.0.content)
                    },
                    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) 
            },
            _ => 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>) -> 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))
        }
    };

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

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

    // Return key
    Response::text(key)
}

// Register link route handler
fn register_link(links: &Bucket<String,String>, req: &Request, id: 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 key
    Response::text(key)
}