upend/src/main.rs

#[macro_use]
extern crate diesel;
#[macro_use]
extern crate diesel_migrations;
#[macro_use]
extern crate lazy_static;

use std::net::SocketAddr;
use std::path::PathBuf;

use actix_cors::Cors;
use actix_web::{middleware, App, HttpServer};
use anyhow::Result;
use clap::{App as ClapApp, Arg};
use log::{debug, info, warn};
use rand::{thread_rng, Rng};
use std::sync::Arc;
use tracing_subscriber::filter::{EnvFilter, LevelFilter};

use crate::{
    common::{build, get_static_dir},
    config::UpEndConfig,
    database::{
        stores::{fs::FsStore, UpStore},
        UpEndDatabase,
    },
    util::{exec::block_background, jobs::JobContainer},
};

mod addressing;
mod common;
mod config;
mod database;
mod extractors;
mod previews;
mod routes;
mod util;

fn main() -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            EnvFilter::builder()
                .with_default_directive(LevelFilter::INFO.into())
                .from_env_lossy(),
        )
        .init();

    let app = ClapApp::new("upend")
        .version(build::PKG_VERSION)
        .author("Tomáš Mládek <t@mldk.cz>")
        .arg(Arg::with_name("DIRECTORY").required(true).index(1))
        .arg(
            Arg::with_name("BIND")
                .long("bind")
                .default_value("127.0.0.1:8093")
                .help("address and port to bind the Web interface on")
                .required(true),
        )
        .arg(
            Arg::with_name("STORE_PATH")
                .long("store")
                .takes_value(true)
                .help(r#"path to store ($VAULT_PATH by default)"#),
        )
        .arg(
            Arg::with_name("NO_BROWSER")
                .long("no-browser")
                .help("Do not open web browser with the UI."),
        )
        .arg(
            Arg::with_name("NO_DESKTOP")
                .long("no-desktop")
                .help("Disable desktop features (webbrowser, native file opening)"),
        )
        .arg(
            Arg::with_name("TRUST_EXECUTABLES")
                .long("trust-executables")
                .help("Trust the vault and open local executable files."),
        )
        .arg(
            Arg::with_name("NO_UI")
                .long("no-ui")
                .help("Do not serve the web UI."),
        )
        .arg(
            Arg::with_name("NO_INITIAL_UPDATE")
                .long("no-initial-update")
                .help("Don't run a database update on start."),
        )
        .arg(
            Arg::with_name("CLEAN")
                .long("clean")
                .help("Clean up temporary files (e.g. previews) on start."),
        )
        .arg(
            Arg::with_name("REINITIALIZE")
                .long("reinitialize")
                .help("Delete and initialize database, if it exists already."),
        )
        .arg(
            Arg::with_name("VAULT_NAME")
                .takes_value(true)
                .long("name")
                .help("Name of the vault."),
        )
        .arg(
            Arg::with_name("SECRET")
                .takes_value(true)
                .long("secret")
                .env("UPEND_SECRET")
                .help("Secret to use for authentication."),
        )
        .arg(
            Arg::with_name("KEY")
                .takes_value(true)
                .long("key")
                .env("UPEND_KEY")
                .help("Authentication key users must supply."),
        );

    let matches = app.get_matches();

    info!("Starting UpEnd {}...", build::PKG_VERSION);
    let sys = actix::System::new("upend");

    let job_container = JobContainer::new();

    let vault_path = PathBuf::from(matches.value_of("DIRECTORY").unwrap());

    let open_result = UpEndDatabase::open(&vault_path, matches.is_present("REINITIALIZE"))
        .expect("failed to open database!");

    let upend = Arc::new(open_result.db);
    let store = Arc::new(Box::new(
        FsStore::from_path(
            matches
                .value_of("STORE_PATH")
                .map(PathBuf::from)
                .unwrap_or_else(|| vault_path.clone()),
        )
        .unwrap(),
    ) as Box<dyn UpStore + Send + Sync>);

    let ui_path = get_static_dir("webui");
    if ui_path.is_err() {
        warn!(
            "Couldn't locate Web UI directory ({:?}), disabling...",
            ui_path
        );
    }
    let desktop_enabled = !matches.is_present("NO_DESKTOP");
    let trust_executables = matches.is_present("TRUST_EXECUTABLES");
    let ui_enabled = ui_path.is_ok() && !matches.is_present("NO_UI");
    let browser_enabled = desktop_enabled && !matches.is_present("NO_BROWSER");

    let preview_path = upend.path.join("previews");
    #[cfg(feature = "previews")]
    let preview_store = Some(Arc::new(crate::previews::PreviewStore::new(
        preview_path.clone(),
        store.clone(),
    )));
    #[cfg(feature = "previews")]
    let preview_pool = Some(Arc::new(
        rayon::ThreadPoolBuilder::new()
            .num_threads(num_cpus::get() / 2)
            .build()
            .unwrap(),
    ));

    if matches.is_present("CLEAN") {
        info!("Cleaning temporary directories...");
        if preview_path.exists() {
            std::fs::remove_dir_all(&preview_path).unwrap();
            debug!("Removed {preview_path:?}");
        } else {
            debug!("No preview path exists, continuing...");
        }
    }

    #[cfg(not(feature = "previews"))]
    let preview_store = None;
    #[cfg(not(feature = "previews"))]
    let preview_pool = None;

    let mut bind: SocketAddr = matches
        .value_of("BIND")
        .unwrap()
        .parse()
        .expect("Incorrect bind format.");

    let secret = matches
        .value_of("SECRET")
        .map(String::from)
        .unwrap_or_else(|| {
            warn!("No secret supplied, generating one at random.");

            thread_rng()
                .sample_iter(&rand::distributions::Alphanumeric)
                .take(32)
                .map(char::from)
                .collect()
        });

    let key = matches.value_of("KEY").map(String::from);

    let state = routes::State {
        upend: upend.clone(),
        store,
        job_container: job_container.clone(),
        preview_store,
        preview_pool,
        config: UpEndConfig {
            vault_name: Some(
                matches
                    .value_of("VAULT_NAME")
                    .map(|s| s.to_string())
                    .unwrap_or_else(|| {
                        vault_path
                            .iter()
                            .last()
                            .unwrap()
                            .to_string_lossy()
                            .into_owned()
                    }),
            ),
            desktop_enabled,
            trust_executables,
            secret,
            key,
        },
    };

    // Start HTTP server

    let mut cnt = 0;
    let ui_path = ui_path.ok();
    let server = loop {
        let state = state.clone();
        let ui_path = ui_path.clone();

        let server = HttpServer::new(move || {
            let cors = Cors::default()
                .allowed_origin("http://localhost")
                .allowed_origin_fn(|origin, _req_head| {
                    origin.as_bytes().starts_with(b"http://localhost:")
                })
                .allow_any_method();

            let app = App::new()
                .wrap(cors)
                .app_data(actix_web::web::PayloadConfig::new(4_294_967_296))
                .data(state.clone())
                .wrap(middleware::Logger::default().exclude("/api/jobs"))
                .service(routes::login)
                .service(routes::get_raw)
                .service(routes::get_thumbnail)
                .service(routes::get_query)
                .service(routes::get_object)
                .service(routes::put_object)
                .service(routes::put_object_attribute)
                .service(routes::delete_object)
                .service(routes::get_all_attributes)
                .service(routes::api_refresh)
                .service(routes::list_hier)
                .service(routes::list_hier_roots)
                .service(routes::store_info)
                .service(routes::get_jobs)
                .service(routes::get_info);

            if let Some(ui_path) = &ui_path {
                app.service(actix_files::Files::new("/", ui_path).index_file("index.html"))
            } else {
                app
            }
        });

        let bind_result = server.bind(&bind);
        if let Ok(server) = bind_result {
            break server;
        } else {
            warn!("Failed to bind at {:?}, trying next port number...", bind);
            bind.set_port(bind.port() + 1);
        }

        if cnt > 10 {
            panic!("Couldn't start server.")
        } else {
            cnt += 1;
        }
    };

    info!("Starting server at: {}", &bind);
    server.run();

    if !matches.is_present("NO_INITIAL_UPDATE") {
        info!("Running initial update...");
        let initial = open_result.new;
        block_background::<_, _, anyhow::Error>(move || {
            let _ = state.store.update(&upend, job_container.clone(), initial);
            let _ = extractors::extract_all(upend, state.store, job_container);
            Ok(())
        })
    }

    #[cfg(feature = "desktop")]
    {
        if browser_enabled && ui_enabled {
            let ui_result = webbrowser::open(&format!("http://localhost:{}", bind.port()));
            if ui_result.is_err() {
                warn!("Could not open UI in browser!");
            }
        }
    }

    Ok(sys.run()?)
}