pastebin/src/main.rs

// pastebin.run
// Copyright (C) 2020-2021 Konrad Borowski
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

#[macro_use]
extern crate diesel;
#[macro_use]
extern crate rocket;

mod models;
mod routes;
mod schema;

use crate::routes::{
    api_insert_paste, config, display_paste, index, insert_paste, metrics, raw_paste,
};
use chrono::{Duration, Utc};
use diesel::prelude::*;
use rocket::fairing::AdHoc;
use rocket::http::Header;
use rocket::shield::{Policy, Referrer, Shield};
use rocket::{Build, Rocket};
use rocket_dyn_templates::tera::{self, Value};
use rocket_dyn_templates::Template;
use rocket_sync_db_pools::database;
use std::collections::HashMap;
use std::net::SocketAddr;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::runtime::Runtime;
use tokio::signal;
use tokio::sync::oneshot::Sender;

#[database("main")]
pub struct Db(PgConnection);

fn js_path(_: &HashMap<String, Value>) -> Result<Value, tera::Error> {
    #[cfg(not(debug_assertions))]
    let path = concat!("/", env!("ENTRY_FILE_PATH"));
    #[cfg(debug_assertions)]
    let path = "http://localhost:5173/js/index.ts";
    Ok(path.into())
}

fn css_stylesheet(_: &HashMap<String, Value>) -> Result<Value, tera::Error> {
    #[cfg(not(debug_assertions))]
    let path = concat!("<link rel=stylesheet href='/", env!("CSS_PATH"), "'>");
    #[cfg(debug_assertions)]
    let path = "";
    Ok(path.into())
}

#[derive(Default)]
struct ContentSecurityPolicy;

impl Policy for ContentSecurityPolicy {
    const NAME: &'static str = "Content-Security-Policy";
    fn header(&self) -> Header<'static> {
        const CONTENT_SECURITY_POLICY: &str = if cfg!(debug_assertions) {
            concat!(
                "default-src 'none';",
                "script-src 'self' localhost:5173;",
                "style-src 'unsafe-inline';",
                "img-src data: https:;",
                "connect-src 'self' ws://localhost:5173;",
                "sandbox allow-forms allow-scripts allow-same-origin;",
                "form-action 'self';",
                "frame-ancestors 'none';",
                "base-uri 'none';",
                "worker-src 'none';",
                "manifest-src 'none'",
            )
        } else {
            concat!(
                "default-src 'none';",
                "script-src 'self';",
                "style-src 'self' 'unsafe-inline';",
                "img-src data: https:;",
                "connect-src 'self';",
                "sandbox allow-forms allow-scripts allow-same-origin;",
                "form-action 'self';",
                "frame-ancestors 'none';",
                "base-uri 'none';",
                "worker-src 'none';",
                "manifest-src 'none'",
            )
        };
        Header::new(Self::NAME, CONTENT_SECURITY_POLICY)
    }
}

fn main() -> Result<(), rocket::Error> {
    let webapp = rocket();

    let (s, r) = tokio::sync::oneshot::channel();

    let runtime = Runtime::new().unwrap();
    runtime.spawn(webapp.launch());
    runtime.spawn(cli());
    runtime.spawn(shutdown_monitor(s));

    let handle = runtime.handle().clone();
    handle.block_on(async move {
        r.await.unwrap();
        runtime.shutdown_background();
    });

    Ok(())
}

async fn shutdown_monitor(s: Sender<i32>) {
    tokio::select! {
        _ = signal::ctrl_c() => {
            s.send(0).unwrap();
        },
    }
}

async fn cli() {
    println!("starting cli");
    let addr = "0.0.0.0:9000".parse::<SocketAddr>().unwrap();
    let listener = TcpListener::bind(addr).await.unwrap();

    let mut db = PgConnection::establish(&std::env::var("POSTGRES_URL").unwrap())
        .expect("error in establishing connection with postgresql");

    let host = std::env::var("HOST").unwrap();

    loop {
        let (socket, details) = listener.accept().await.unwrap();

        println!("received connection from {:?}", details);
        save_paste(&mut db, &host, socket).await.unwrap();
    }
}

const BUFFER_SIZE: usize = 128;

async fn save_paste(
    db: &mut PgConnection,
    host: &str,
    mut socket: TcpStream,
) -> Result<(), Box<dyn std::error::Error>> {
    let mut paste = vec![];
    let mut buffer = [0; BUFFER_SIZE];

    loop {
        let read = socket.read(&mut buffer).await?;
        if read < BUFFER_SIZE {
            paste.extend(&buffer[..read]);
            break;
        }

        paste.extend(buffer);
        buffer.fill(0);
    }

    use models::paste;
    let v = paste::insert(
        db,
        // Terminal paste should stay for 3 days
        Some(Utc::now() + Duration::days(3)),
        &String::from_utf8_lossy(&paste),
    )
    .unwrap();

    let output = format!("{}/{}\n", host, v);
    let output: Vec<u8> = output.into_bytes();

    socket.write(&output).await?;
    socket.flush().await.unwrap();
    Ok(())
}

fn rocket() -> Rocket<Build> {
    println!("starting rocket");
    let mut rocket = rocket::build()
        .attach(Template::custom(|engines| {
            engines.tera.register_function("js_path", js_path);
            engines
                .tera
                .register_function("css_stylesheet", css_stylesheet);
        }))
        .attach(Db::fairing())
        .attach(AdHoc::on_ignite("Migrations", |rocket| async {
            Db::get_one(&rocket)
                .await
                .expect("a database")
                .run(|conn| diesel_migrations::run_pending_migrations(conn))
                .await
                .expect("database to be migrated");
            rocket
        }))
        .attach(
            Shield::default()
                .enable(ContentSecurityPolicy)
                .enable(Referrer::NoReferrer),
        )
        .mount(
            "/",
            routes![
                api_insert_paste,
                config,
                index,
                insert_paste,
                display_paste,
                raw_paste,
                metrics,
            ],
        );
    if cfg!(not(debug_assertions)) {
        rocket = rocket.mount("/assets", rocket::fs::FileServer::from("dist/assets"));
    }
    rocket
}