monorepo/file-service/src/main.rs

extern crate log;

use cookie::Cookie;
use handlers::{file, handle_auth, handle_css, handle_delete, handle_upload, thumbnail};
use std::{
    collections::{HashMap, HashSet},
    convert::Infallible,
    net::{IpAddr, Ipv4Addr, SocketAddr},
    path::PathBuf,
    sync::Arc,
};
use tokio::sync::RwLock;
use warp::{Filter, Rejection};

mod handlers;
mod html;
mod pages;

const MAX_UPLOAD: u64 = 15 * 1024 * 1024;

pub use file_service::{
    AuthDB, AuthError, AuthToken, DeleteFileError, FileHandle, FileId, FileInfo, ReadFileError,
    SessionToken, Store, Username, WriteFileError,
};
pub use handlers::handle_index;

#[derive(Clone)]
pub struct App {
    authdb: Arc<RwLock<AuthDB>>,
    store: Arc<RwLock<Store>>,
}

impl App {
    pub fn new(authdb: AuthDB, store: Store) -> Self {
        Self {
            authdb: Arc::new(RwLock::new(authdb)),
            store: Arc::new(RwLock::new(store)),
        }
    }

    pub async fn authenticate(&self, token: AuthToken) -> Result<Option<SessionToken>, AuthError> {
        self.authdb.read().await.authenticate(token).await
    }

    pub async fn validate_session(
        &self,
        token: SessionToken,
    ) -> Result<Option<Username>, AuthError> {
        self.authdb.read().await.validate_session(token).await
    }

    pub async fn list_files(&self) -> Result<HashSet<FileId>, ReadFileError> {
        self.store.read().await.list_files()
    }

    pub async fn get_file(&self, id: &FileId) -> Result<FileHandle, ReadFileError> {
        self.store.read().await.get_file(id)
    }

    pub async fn add_file(
        &self,
        filename: String,
        content: Vec<u8>,
    ) -> Result<FileHandle, WriteFileError> {
        self.store.write().await.add_file(filename, content)
    }

    pub async fn delete_file(&self, id: FileId) -> Result<(), DeleteFileError> {
        self.store.write().await.delete_file(&id)?;
        Ok(())
    }
}

fn with_app(app: App) -> impl Filter<Extract = (App,), Error = Infallible> + Clone {
    warp::any().map(move || app.clone())
}

fn parse_cookies(cookie_str: &str) -> Result<HashMap<String, String>, cookie::ParseError> {
    Cookie::split_parse(cookie_str)
        .map(|c| c.map(|c| (c.name().to_owned(), c.value().to_owned())))
        .collect::<Result<HashMap<String, String>, cookie::ParseError>>()
}

fn get_session_token(cookies: HashMap<String, String>) -> Option<SessionToken> {
    cookies.get("session").cloned().map(SessionToken::from)
}

fn maybe_with_session() -> impl Filter<Extract = (Option<SessionToken>,), Error = Rejection> + Copy
{
    warp::any()
        .and(warp::header::optional::<String>("cookie"))
        .map(|cookie_str: Option<String>| match cookie_str {
            Some(cookie_str) => parse_cookies(&cookie_str).ok().and_then(get_session_token),
            None => None,
        })
}

fn with_session() -> impl Filter<Extract = (SessionToken,), Error = Rejection> + Copy {
    warp::any()
        .and(warp::header::<String>("cookie"))
        .and_then(|cookie_str: String| async move {
            match parse_cookies(&cookie_str).ok().and_then(get_session_token) {
                Some(session_token) => Ok(session_token),
                None => Err(warp::reject()),
            }
        })
}

#[tokio::main]
pub async fn main() {
    pretty_env_logger::init();

    let authdb = AuthDB::new(PathBuf::from(&std::env::var("AUTHDB").unwrap()))
        .await
        .unwrap();
    let store = Store::new(PathBuf::from(&std::env::var("FILE_SHARE_DIR").unwrap()));

    let app = App::new(authdb, store);

    let log = warp::log("file_service");
    let root = warp::path!()
        .and(warp::get())
        .and(with_app(app.clone()))
        .and(maybe_with_session())
        .then(handle_index);

    let styles = warp::path!("css").and(warp::get()).then(handle_css);

    let auth = warp::path!("auth")
        .and(warp::post())
        .and(with_app(app.clone()))
        .and(warp::filters::body::form())
        .then(handle_auth);

    let upload_via_form = warp::path!("upload")
        .and(warp::post())
        .and(with_app(app.clone()))
        .and(with_session())
        .and(warp::multipart::form().max_length(MAX_UPLOAD))
        .then(handle_upload);

    let delete_via_form = warp::path!("delete" / String)
        .and(warp::post())
        .and(with_app(app.clone()))
        .and(with_session())
        .then(|id, app, token| handle_delete(app, token, FileId::from(id)));

    let thumbnail = warp::path!(String / "tn")
        .and(warp::get())
        .and(warp::header::optional::<String>("if-none-match"))
        .and(with_app(app.clone()))
        .then(move |id, old_etags, app: App| thumbnail(app, id, old_etags));

    let file = warp::path!(String)
        .and(warp::get())
        .and(warp::header::optional::<String>("if-none-match"))
        .and(with_app(app.clone()))
        .then(move |id, old_etags, app: App| file(app, id, old_etags));

    let server = warp::serve(
        root.or(styles)
            .or(auth)
            .or(upload_via_form)
            .or(delete_via_form)
            .or(thumbnail)
            .or(file)
            .with(log),
    );

    server
        .run(SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 8002))
        .await;
}