ImageApi/src/main.rs

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

use actix_files::NamedFile;
use actix_multipart as mp;
use actix_web::web::{HttpRequest, HttpResponse, Json};
use actix_web::{get, post, web, App, HttpServer, Responder};
use chrono::{Duration, Utc};
use data::{AddFavoriteRequest, LoginRequest, ThumbnailRequest};
use futures::stream::StreamExt;
use jsonwebtoken::{encode, EncodingKey, Header};
use notify::{watcher, DebouncedEvent, RecursiveMode, Watcher};
use rayon::prelude::*;
use serde::Serialize;
use std::fs::File;
use std::io::prelude::*;
use std::path::{Path, PathBuf};
use std::sync::mpsc::channel;

use crate::data::{secret_key, Claims, CreateAccountRequest, Token};
use crate::database::{add_favorite, create_user, get_favorites, get_user, user_exists};
use crate::files::{is_valid_path, list_files};
use crate::video::*;

mod data;
mod database;
mod files;
mod video;

#[post("/register")]
async fn register(user: Json<CreateAccountRequest>) -> impl Responder {
    if !user.username.is_empty() && user.password.len() > 5 && user.password == user.confirmation {
        if user_exists(&user.username) {
            HttpResponse::BadRequest()
        } else if let Some(_user) = create_user(&user.username, &user.password) {
            HttpResponse::Ok()
        } else {
            HttpResponse::InternalServerError()
        }
    } else {
        HttpResponse::BadRequest()
    }
}

#[post("/login")]
async fn login(creds: Json<LoginRequest>) -> impl Responder {
    println!("Logging in: {}", creds.username);
    if let Some(user) = get_user(&creds.username, &creds.password) {
        let claims = Claims {
            sub: user.id.to_string(),
            exp: (Utc::now() + Duration::days(3)).timestamp(),
        };
        let token = encode(
            &Header::default(),
            &claims,
            &EncodingKey::from_secret(secret_key().as_bytes()),
        )
        .unwrap();
        HttpResponse::Ok().json(Token { token: &token })
    } else {
        HttpResponse::NotFound().finish()
    }
}

#[post("/photos")]
async fn list_photos(_claims: Claims, req: Json<ThumbnailRequest>) -> impl Responder {
    println!("{}", req.path);

    let path = &req.path;
    if let Some(path) = is_valid_path(path) {
        let files = list_files(path).unwrap_or_default();

        let photos = &files
            .iter()
            .filter(|f| !f.extension().unwrap_or_default().is_empty())
            .map(|f| f.to_str().unwrap().to_string())
            .collect::<Vec<String>>();

        let dirs = &files
            .iter()
            .filter(|f| f.extension().unwrap_or_default().is_empty())
            .map(|f| f.to_str().unwrap().to_string())
            .collect::<Vec<String>>();

        HttpResponse::Ok().json(PhotosResponse { photos, dirs })
    } else {
        HttpResponse::BadRequest().finish()
    }
}

#[derive(Serialize)]
struct PhotosResponse<'a> {
    photos: &'a [String],
    dirs: &'a [String],
}

#[get("/image")]
async fn get_image(
    _claims: Claims,
    request: HttpRequest,
    req: web::Query<ThumbnailRequest>,
) -> impl Responder {
    if let Some(path) = is_valid_path(&req.path) {
        if req.size.is_some() {
            let thumbs = dotenv::var("THUMBNAILS").unwrap();
            let relative_path = path
                .strip_prefix(dotenv::var("BASE_PATH").unwrap())
                .expect("Error stripping prefix");
            let thumb_path = Path::new(&thumbs).join(relative_path);

            println!("{:?}", thumb_path);
            if let Ok(file) = NamedFile::open(&thumb_path) {
                file.into_response(&request).unwrap()
            } else {
                HttpResponse::NotFound().finish()
            }
        } else if let Ok(file) = NamedFile::open(path) {
            file.into_response(&request).unwrap()
        } else {
            HttpResponse::NotFound().finish()
        }
    } else {
        HttpResponse::BadRequest().finish()
    }
}

#[post("/image")]
async fn upload_image(_: Claims, mut payload: mp::Multipart) -> impl Responder {
    let mut file_content: Vec<_> = Vec::new();
    let mut file_name: Option<String> = None;
    let mut file_path: Option<String> = None;

    while let Some(Ok(mut part)) = payload.next().await {
        if let Some(content_type) = part.content_disposition() {
            println!("{:?}", content_type);
            if let Some(filename) = content_type.get_filename() {
                println!("Name: {:?}", filename);
                file_name = Some(filename.to_string());

                while let Some(Ok(data)) = part.next().await {
                    file_content.extend_from_slice(data.as_ref());
                }
            } else if content_type.get_name().map_or(false, |name| name == "path") {
                while let Some(Ok(data)) = part.next().await {
                    if let Ok(path) = std::str::from_utf8(&data) {
                        file_path = Some(path.to_string())
                    }
                }
            }
        }
    }

    let path = file_path.unwrap_or_else(|| dotenv::var("BASE_PATH").unwrap());
    if !file_content.is_empty() {
        let full_path = PathBuf::from(&path);

        if let Some(mut full_path) = is_valid_path(full_path.to_str().unwrap_or("")) {
            // TODO: Validate this file_name as is subject to path traversals which could lead to
            // writing outside the base dir.
            full_path = full_path.join(file_name.unwrap());

            if !full_path.is_file() {
                let mut file = File::create(full_path).unwrap();
                file.write_all(&file_content).unwrap();
            } else {
                return HttpResponse::BadRequest().body("File already exists");
            }
        } else {
            return HttpResponse::BadRequest().body("Path was not valid");
        }
    } else {
        return HttpResponse::BadRequest().body("No file body read");
    }
    HttpResponse::Ok().finish()
}

#[post("/video/generate")]
async fn generate_video(_claims: Claims, body: web::Json<ThumbnailRequest>) -> impl Responder {
    let filename = PathBuf::from(&body.path);

    if let Some(name) = filename.file_stem() {
        let filename = name.to_str().expect("Filename should convert to string");
        let playlist = format!("tmp/{}.m3u8", filename);
        if let Some(path) = is_valid_path(&body.path) {
            create_playlist(&path.to_str().unwrap(), &playlist);
        } else {
            return HttpResponse::BadRequest().finish();
        }

        HttpResponse::Ok().json(playlist)
    } else {
        HttpResponse::BadRequest().finish()
    }
}

#[get("/video/stream")]
async fn stream_video(
    request: HttpRequest,
    _: Claims,
    path: web::Query<ThumbnailRequest>,
) -> impl Responder {
    let playlist = &path.path;
    println!("Playlist: {}", playlist);

    // Extract video playlist dir to dotenv
    if !playlist.starts_with("tmp") || playlist.contains("..") {
        HttpResponse::NotFound().finish()
    }
    else if let Ok(file) = NamedFile::open(playlist) {
        file.into_response(&request).unwrap()
    } else {
        HttpResponse::NotFound().finish()
    }
}

#[get("/video/{path}")]
async fn get_video_part(
    request: HttpRequest,
    _: Claims,
    path: web::Path<ThumbnailRequest>,
) -> impl Responder {
    let part = &path.path;
    println!("Video part: {}", part);

    if let Ok(file) = NamedFile::open(String::from("tmp/") + part) {
        file.into_response(&request).unwrap()
    } else {
        HttpResponse::NotFound().finish()
    }
}

#[get("image/favorites")]
async fn favorites(claims: Claims) -> impl Responder {
    let favorites = get_favorites(claims.sub.parse::<i32>().unwrap())
        .into_iter()
        .map(|favorite| favorite.path)
        .collect::<Vec<String>>();
    HttpResponse::Ok().json(PhotosResponse {
        photos: &favorites,
        dirs: &Vec::new(),
    })
}

#[post("image/favorites")]
async fn post_add_favorite(claims: Claims, body: web::Json<AddFavoriteRequest>) -> impl Responder {
    if let Ok(user_id) = claims.sub.parse::<i32>() {
        add_favorite(user_id, body.path.clone());
        HttpResponse::Ok()
    } else {
        HttpResponse::BadRequest()
    }
}

async fn create_thumbnails() {
    let thumbs = &dotenv::var("THUMBNAILS").expect("THUMBNAILS not defined");
    let thumbnail_directory: &Path = Path::new(thumbs);

    let images = PathBuf::from(dotenv::var("BASE_PATH").unwrap());

    walkdir::WalkDir::new(&images)
        .into_iter()
        .collect::<Vec<Result<_, _>>>()
        .into_par_iter()
        .filter_map(|entry| entry.ok())
        .filter(|entry| {
            println!("{:?}", entry.path());
            if let Some(ext) = entry
                .path()
                .extension()
                .and_then(|ext| ext.to_str().map(|ext| ext.to_lowercase()))
            {
                if ext == "mp4" || ext == "mov" {
                    let relative_path = &entry.path().strip_prefix(&images).unwrap();
                    let thumb_path = Path::new(thumbnail_directory).join(relative_path);
                    std::fs::create_dir_all(&thumb_path.parent().unwrap())
                        .expect("Error creating directory");
                    generate_video_thumbnail(entry.path(), &thumb_path);
                    false
                } else {
                    ext == "jpg" || ext == "jpeg" || ext == "png"
                }
            } else {
                false
            }
        })
        .filter(|entry| {
            let path = entry.path();
            let relative_path = &path.strip_prefix(&images).unwrap();
            let thumb_path = Path::new(thumbnail_directory).join(relative_path);
            !thumb_path.exists()
        })
        .map(|entry| (image::open(entry.path()), entry.path().to_path_buf()))
        .filter(|(img, _)| img.is_ok())
        .map(|(img, path)| (img.unwrap(), path))
        .map(|(image, path)| (image.thumbnail(200, 200), path))
        .map(|(image, path)| {
            let relative_path = &path.strip_prefix(&images).unwrap();
            let thumb_path = Path::new(thumbnail_directory).join(relative_path);
            std::fs::create_dir_all(&thumb_path.parent().unwrap())
                .expect("There was an issue creating directory");
            println!("{:?}", thumb_path);
            image.save(thumb_path).expect("Failure saving thumbnail");
        })
        .for_each(drop);

    println!("Finished");
}

#[actix_rt::main]
async fn main() -> std::io::Result<()> {
    create_thumbnails().await;

    tokio::spawn(async {
        let (wtx, wrx) = channel();
        let mut watcher = watcher(wtx, std::time::Duration::from_secs(10)).unwrap();
        watcher
            .watch(dotenv::var("BASE_PATH").unwrap(), RecursiveMode::Recursive)
            .unwrap();

        loop {
            let ev = wrx.recv_timeout(std::time::Duration::from_secs(5));
            match ev {
                Ok(event) => {
                    match event {
                        DebouncedEvent::Create(_) => create_thumbnails().await,
                        DebouncedEvent::Rename(_, _) => create_thumbnails().await,
                        _ => continue,
                    };
                }
                Err(e) => {
                    println!("Event: {:?}", e);
                    // break;
                }
            }
        }
    });

    HttpServer::new(|| {
        App::new()
            .service(register)
            .service(login)
            .service(list_photos)
            .service(get_image)
            .service(upload_image)
            .service(generate_video)
            .service(stream_video)
            .service(get_video_part)
            .service(favorites)
            .service(post_add_favorite)
    })
    .bind(dotenv::var("BIND_URL").unwrap())?
    .bind("localhost:8088")?
    .run()
    .await
}