crow/src/routes/image.rs

use crate::*;
use blake2::{Blake2s256, Digest};
use image::{
    codecs::{jpeg::JpegEncoder, png::PngEncoder},
    ColorType, ImageEncoder,
};
use lmdb::LmdbResultExt;
use lmdb_zero as lmdb;
use multipart::server::Multipart;

use std::{
    io::{ErrorKind, Read},
    ops::Deref,
};
use tiny_http::{Header, Request, Response};

fn image_upload<'u>(
    mut request: Request,
    db_context: &DatabaseContext,
    _: matchit::Params<'u, 'u>,
) -> CrowResult<()> {
    let mut entry = some_or_response!(
        single_multipart!(&mut request), or respond to request with Response::empty(400)
    );

    let mut data: Vec<u8> = Vec::with_capacity(20000);
    entry.data.read_to_end(&mut data)?;

    let upload_format = image::guess_format(&data)?;
    let img = image::load_from_memory_with_format(&data, upload_format)?;

    let data_hash = Blake2s256::digest(&img.as_bytes());

    let txn = db_context.write_txn()?;

    let mut accessor = txn.access();

    if accessor
        .get::<[u8], [u8]>(&db_context.image_store, data_hash.as_slice())
        .to_opt()?
        .is_some()
    {
        request.respond(Response::from_string(base64_url::encode(&data_hash)))?;
        return Ok(());
    }

    let store_format = match upload_format {
        image::ImageFormat::Jpeg => {
            accessor.put(
                &db_context.image_store,
                data_hash.as_slice(),
                &data,
                lmdb::put::Flags::empty(),
            )?;

            ImageFormat::JPG
        }
        _ => {
            // TODO: allow lossy compression via option
            let encoded = webp::Encoder::from_image(&img).unwrap().encode_lossless();
            accessor.put(
                &db_context.image_store,
                data_hash.as_slice(),
                encoded.deref(),
                lmdb::put::Flags::empty(),
            )?;

            ImageFormat::WEBP
        }
    };

    accessor.put(
        &db_context.image_meta_store,
        data_hash.as_slice(),
        rkyv::to_bytes::<_, 256>(&ImageHeader { store_format })
            .unwrap()
            .as_slice(),
        lmdb::put::Flags::empty(),
    )?;

    request.respond(Response::from_string(base64_url::encode(&data_hash)))?;

    drop(accessor);

    txn.commit()?;

    Ok(())
}

fn image_get<'u>(
    mut request: Request,
    db_context: &DatabaseContext,
    params: matchit::Params<'u, 'u>,
) -> CrowResult<()> {
    let (get_id, requested_format) = some_or_response!(
        parse_base64_hash!(params.get("id")).zip(
            params
                .get("format")
                .and_then(|s| ImageFormat::from_str(s).ok())
        ), or respond to request with Response::empty(400)
    );

    let content_type = requested_format.to_mime();

    let read = db_context.read_txn()?;
    let access = read.access();

    let metadata: &ArchivedImageHeader = some_or_response!(
        access
            .get::<[u8], ArchivedImageHeader>(&db_context.image_meta_store, &get_id)
            .to_opt()?,
        or respond to request with
        Response::empty(404)
    );

    let data = access.get::<[u8], [u8]>(&db_context.image_store, &get_id)?;

    if metadata.store_format == requested_format {
        request.respond(Response::new(
            200.into(),
            vec![],
            data,
            Some(data.len()),
            None,
        ))?;

        return Ok(());
    }

    let stored_image = image::load_from_memory_with_format(data, metadata.store_format.into())?;

    use ColorType::*;

    // catch unsupported transcoding (because of unsupported color format)
    // and return a 500
    match (stored_image.color(), requested_format) {
        (Rgb8 | Rgba8, ImageFormat::WEBP) => (),
        (Rgb8 | Rgba8 | L8 | La8, ImageFormat::JPG) => (),
        (Rgb8 | Rgba8 | L8 | La8 | Rgba16 | Rgb16 | L16 | La16, ImageFormat::PNG) => (),
        _ => {
            request.respond(
                response!(err 500 "unsupported color channels for the requested image format"),
            )?;
            return Ok(());
        }
    }

    let mut req_writer = request.extract_writer_impl();

    let response = Response::new(
        200.into(),
        vec![Header::from_bytes(&b"Content-Type"[..], content_type.as_bytes()).unwrap()],
        std::io::empty(),
        None,
        None,
    );
    Request::ignore_client_closing_errors(response.print_and_write(
        &mut req_writer,
        request.http_version().clone(),
        request.headers(),
        false,
        None,
        None,
        |w, _| {
            match requested_format {
                ImageFormat::WEBP => {
                    // TODO: make compression quality here configurable
                    let mem = webp::Encoder::from_image(&stored_image)
                        .unwrap()
                        .encode(95.0);
                    w.write_all(mem.deref())
                }
                ImageFormat::JPG => JpegEncoder::new(w)
                    .encode_image(&stored_image)
                    .map_err(|e| std::io::Error::new(ErrorKind::Other, e)),
                ImageFormat::PNG => PngEncoder::new(w)
                    .write_image(
                        stored_image.as_bytes(),
                        stored_image.width(),
                        stored_image.height(),
                        stored_image.color(),
                    )
                    .map_err(|e| std::io::Error::new(ErrorKind::Other, e)),
            }
        },
    ))?;

    Request::ignore_client_closing_errors(req_writer.flush())?;
    if let Some(sender) = request.notify_when_responded.take() {
        sender.send(()).unwrap();
    }

    drop(req_writer);
    drop(request);

    Ok(())
}

fn_to_handler!(image_upload: ImageUploadHandler);
fn_to_handler!(image_get: ImageGetHandler);