image-gallery/src/main.rs

use std::{net::SocketAddr, path::{Path, PathBuf}, ffi::OsStr, fs::File, io::{BufReader, Seek, SeekFrom, Cursor, BufRead}, time::SystemTime, env::args_os};

use axum::{Router, routing::{get, post}, response::{IntoResponse, Redirect}, http::{StatusCode, header}, extract::{self, State}, Form, handler::Handler};
use axum_sessions::{async_session::CookieStore, SessionLayer, extractors::{ReadableSession, WritableSession}};
use axum_template::{RenderHtml, engine::Engine};
use exif::{Tag, Value, Field, In, Exif};
use image::{imageops::FilterType, DynamicImage};
use rand::Rng;
use serde::{Serialize, Deserialize};
use tokio::task;
use tower_http::{trace::{self, TraceLayer}, compression::CompressionLayer};
use tracing::{Level, debug_span};
use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt, fmt::format::FmtSpan};

// TODO
// * Proper error handling
// * Cache generated images (+ cache cleanup)
// * Add configuration file
// * Support for headlines
// * Polish the css

type TemplateEngine = Engine<minijinja::Environment<'static>>;
type ImageDir = PathBuf;

#[derive(Clone, extract::FromRef)]
struct ApplicationState {
    engine: TemplateEngine,
    image_dir: ImageDir,
}

#[tokio::main]
async fn main() {
    let image_path = args_os().skip(1).next().expect("Usage: image-gallery IMAGE_DIRECTORY");

    let default_tracing = "image_gallery=debug,tower_http=info".into();
    let tracing_filter = EnvFilter::try_from_default_env().unwrap_or(default_tracing);
    let tracing_formatter = tracing_subscriber::fmt::layer()
        .with_target(false)
        .with_span_events(FmtSpan::CLOSE) // TODO: maybe drop this
        .compact();
    tracing_subscriber::registry()
        .with(tracing_filter)
        .with(tracing_formatter)
        .init();

    let mut jinja = minijinja::Environment::new();
    jinja.add_template("base", include_str!("../templates/base.html")).unwrap();
    jinja.add_template("index", include_str!("../templates/index.html")).unwrap();
    jinja.add_template("login", include_str!("../templates/login.html")).unwrap();

    let secret = rand::thread_rng().gen::<[u8; 128]>();
    let session_layer = SessionLayer::new(CookieStore::new(), &secret)
        .with_cookie_name("session")
        .with_session_ttl(None);

    let app = Router::new()
        .route("/", get(index.layer(CompressionLayer::new())))
        .route("/authenticate", post(authenticate))
        .route("/image/:image", get(converted_image))
        .layer(TraceLayer::new_for_http()
            .make_span_with(trace::DefaultMakeSpan::new().level(Level::INFO))
            .on_response(trace::DefaultOnResponse::new().level(Level::INFO))
        )
        .layer(session_layer)
        .with_state(ApplicationState {
            engine: Engine::from(jinja),
            image_dir: PathBuf::from(image_path),
        });

    let addr = SocketAddr::from(([0, 0, 0, 0], 3000));
    tracing::debug!("listening on {}", addr);
    axum::Server::bind(&addr)
        .serve(app.into_make_service())
        .await
        .unwrap();
}

#[derive(Debug, Serialize)]
pub struct ImageInfo {
    width: u32,
    height: u32,
    name: String,
}

#[derive(Debug, Serialize)]
pub struct IndexTempalte {
    title: String,
    images: Vec<ImageInfo>,
}

async fn index(
    engine: TemplateEngine,
    State(image_dir): State<ImageDir>,
    session: ReadableSession,
) -> impl IntoResponse {

    let logged_in = session.get("logged_in").unwrap_or(false);

    if logged_in {
        let images = read_images(&image_dir);

        let images = images.into_iter().rev().take(30).collect();

        RenderHtml("index", engine, IndexTempalte {
            title: "Some pictures".into(),
            images: images,
        })
    } else {
        RenderHtml("login", engine, IndexTempalte {
            title: "Some pictures".into(),
            images: vec![],
        })
    }
}

#[derive(Deserialize)]
struct AuthenticateForm {
    password: String,
}

async fn authenticate(
    mut session: WritableSession,
    Form(form): Form<AuthenticateForm>
) -> Redirect {
    if form.password == "testle" {
        session.insert("logged_in", true).ok();
    }
    Redirect::to("/")
}

async fn converted_image(
    extract::Path(image): extract::Path<String>,
    State(image_dir): State<ImageDir>,
    session: ReadableSession,
) -> Result<impl IntoResponse, impl IntoResponse> {
    if ! session.get("logged_in").unwrap_or(false) {
        return Err(StatusCode::FORBIDDEN);
    }

    let image_path = image_dir.join(image);

    if ! image_path.exists() {
        return Err(StatusCode::NOT_FOUND);
    }

    let image_buffer = task::spawn_blocking(move || -> Result<_, StatusCode> {
        let file = File::open(&image_path).unwrap();
        let mut file = BufReader::new(file);

        // Read exif data first, as it only reads a tiny amount of data
        let exif = read_exif_data(&mut file);

        let mut image = debug_span!("decode_image",
            image=?image_path.file_name(),
        ).in_scope(|| -> Result<_, StatusCode>{
            file.seek(SeekFrom::Start(0)).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
            image::io::Reader::new(&mut file)
                .with_guessed_format()
                .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
                .decode()
                .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)
        })?;

        // Fix image orientation
        if let Ok(exif) = exif {
            image = fix_image_orientation(image, &exif);
        }

        let small_image = debug_span!("resize_image",
            image=?image_path.file_name(),
        ).in_scope(|| {
            image.resize(600, 600, FilterType::Lanczos3)
        });

        let mut image_buffer = Cursor::new(Vec::new());

        debug_span!("encode_image",
            image=?image_path.file_name(),
        ).in_scope(|| {
            small_image.write_to(&mut image_buffer, image::ImageOutputFormat::WebP)
        }).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

        Ok(image_buffer)
    }).await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)??;

    return Ok((
        [
            (header::CONTENT_TYPE, "image/webp"),
        ],
        image_buffer.into_inner()
    ))
}

fn read_exif_data(mut file: impl BufRead+Seek) -> Result<Exif, ()> {
    file.seek(SeekFrom::Start(0)).map_err(|_| ())?;
    let exifreader = exif::Reader::new();
    exifreader.read_from_container(&mut file)
        .map_err(|_| ())
}

fn read_image_size(mut file: impl BufRead+Seek, exif: Option<&Exif>) -> Result<(u32, u32), ()> {
    file.seek(SeekFrom::Start(0)).map_err(|_| ())?;

    let (mut width, mut height) = image::io::Reader::new(&mut file)
        .with_guessed_format()
        .map_err(|_| ())?
        .into_dimensions()
        .map_err(|_| ())?;

    if let Some(exif) = exif {
        match exif.get_field(Tag::Orientation, In::PRIMARY) {
            Some(orientation) =>
                match orientation.value.get_uint(0) {
                    Some(5 | 6 | 7 | 8) => std::mem::swap(&mut width, &mut height),
                    _ => {},
                },
            None => {},
        };
    };

    Ok((width, height))
}

// How many degrees to rotate clockwise, does not support flipped images
fn fix_image_orientation(image: DynamicImage, exif: &Exif) -> DynamicImage {
    match exif.get_field(Tag::Orientation, In::PRIMARY) {
        Some(orientation) =>
            match orientation.value.get_uint(0) {
                Some(1) => image,
                Some(2) => image.fliph(),
                Some(3) => image.rotate180(),
                Some(4) => image.rotate180().fliph(),
                Some(5) => image.rotate90().fliph(),
                Some(6) => image.rotate90(),
                Some(7) => image.rotate270().fliph(),
                Some(8) => image.rotate270(),
                _ => image,
            },
        None => image,
    }
}


fn read_images(directory: &Path) -> Vec<ImageInfo> {
    let mut files = vec![];
    for file in directory.read_dir().expect("read_dir call failed") {
        if let Ok(file_entry) = file {
            let path = file_entry.path();
            if path.extension() == Some(OsStr::new("jpg")) {
                let file = File::open(&path).unwrap();
                let mut file = BufReader::new(file);

                let exif = read_exif_data(&mut file);

                // Check if we need to flip the coordinates
                let (width, height) = read_image_size(&mut file, exif.as_ref().ok()).unwrap();

                let datetime = 'datetime: {
                    // First, try to read creation date from EXIF data
                    if let Ok(ref exif) = exif {
                        match exif.get_field(Tag::DateTimeOriginal, In::PRIMARY) {
                            Some(Field { value: Value::Ascii(value), ..}) if !value.is_empty() => {
                                break 'datetime String::from_utf8_lossy(&value[0]).into_owned()
                            }
                            _ => {}
                        };
                    }

                    // If that doesn't work, fall back to the file modification time
                    format!("{:?}", std::fs::metadata(&path).unwrap().modified().unwrap().duration_since(SystemTime::UNIX_EPOCH).unwrap())
                };

                let image_info = ImageInfo {
                    width: width,
                    height: height,
                    name: path.file_name().expect("invalid file path").to_string_lossy().to_string(),
                };

                files.push(image_info);
            }
        }
    }
    files
}