crow/src/lib.rs

use blake2::Digest;
use lmdb::{traits::FromLmdbBytes, Database, Environment};
use lmdb_zero as lmdb;

use tiny_http::Request;

use std::{str::FromStr, sync::Arc};
use thiserror::Error;

pub trait Handler: Send + Sync {
    fn respond<'u>(
        &self,
        req: Request,
        db: &DatabaseContext,
        params: matchit::Params<'u, 'u>,
    ) -> CrowResult<()>;
}

pub mod routes;

#[derive(Error, Debug)]
pub enum CrowError {
    #[error(transparent)]
    LmdbError(#[from] lmdb::Error),
    #[error(transparent)]
    IOError(#[from] std::io::Error),
    #[error(transparent)]
    ImageError(#[from] image::error::ImageError),
    #[error(transparent)]
    Other(#[from] anyhow::Error),
}

pub type CrowResult<T> = Result<T, CrowError>;

#[derive(rkyv::Archive, rkyv::Serialize)]
pub struct FileHeader {
    file_name: Option<String>,
    content_type: Option<String>,
}

impl FromLmdbBytes for ArchivedFileHeader {
    fn from_lmdb_bytes(bytes: &[u8]) -> Result<&Self, String> {
        Ok(unsafe { rkyv::archived_root::<FileHeader>(bytes) })
    }
}

#[derive(rkyv::Serialize, rkyv::Archive, Copy, Debug, Clone)]
#[archive(compare(PartialEq))]
#[archive_attr(derive(Copy, Clone))]
#[repr(u8)]
pub enum ImageFormat {
    PNG,
    WEBP,
    JPG,
}

impl FromStr for ImageFormat {
    type Err = ();

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Ok(match s {
            "png" | "PNG" => ImageFormat::PNG,
            "webp" | "WEBP" => ImageFormat::WEBP,
            "jpg" | "JPG" | "jpeg" | "JPEG" => ImageFormat::JPG,
            _ => return Err(()),
        })
    }
}

impl ImageFormat {
    pub const fn to_mime(&self) -> &'static str {
        match *self {
            ImageFormat::PNG => "image/png",
            ImageFormat::JPG => "image/jpeg",
            ImageFormat::WEBP => "image/webp",
        }
    }
}

impl From<ImageFormat> for image::ImageFormat {
    fn from(val: ImageFormat) -> Self {
        match val {
            ImageFormat::PNG => image::ImageFormat::Png,
            ImageFormat::WEBP => image::ImageFormat::WebP,
            ImageFormat::JPG => image::ImageFormat::Jpeg,
        }
    }
}

impl From<ArchivedImageFormat> for image::ImageFormat {
    fn from(val: ArchivedImageFormat) -> Self {
        match val {
            ArchivedImageFormat::PNG => image::ImageFormat::Png,
            ArchivedImageFormat::WEBP => image::ImageFormat::WebP,
            ArchivedImageFormat::JPG => image::ImageFormat::Jpeg,
        }
    }
}

#[derive(rkyv::Serialize, rkyv::Archive)]
pub struct ImageHeader {
    store_format: ImageFormat,
}

impl FromLmdbBytes for ArchivedImageHeader {
    fn from_lmdb_bytes(bytes: &[u8]) -> Result<&Self, String> {
        Ok(unsafe { rkyv::archived_root::<ImageHeader>(bytes) })
    }
}

#[derive(Clone)]
pub struct DatabaseContext {
    pub env: Arc<Environment>,
    pub metadata_store: Arc<Database<'static>>,
    pub binary_store: Arc<Database<'static>>,
    pub image_store: Arc<Database<'static>>,
    pub image_meta_store: Arc<Database<'static>>,
}

impl DatabaseContext {
    pub fn create(path: &str) -> CrowResult<DatabaseContext> {
        let env = Arc::new(unsafe {
            let mut builder = lmdb::EnvBuilder::new()?;
            builder.set_maxdbs(4)?;
            builder.open(path, lmdb::open::NOTLS, 0o600)?
        });

        let binary_store = Arc::new(lmdb::Database::open(
            env.clone(),
            Some("binary"),
            &lmdb::DatabaseOptions::new(lmdb::db::CREATE),
        )?);
        let image_store = Arc::new(lmdb::Database::open(
            env.clone(),
            Some("image"),
            &lmdb::DatabaseOptions::new(lmdb::db::CREATE),
        )?);
        let metadata_store = Arc::new(lmdb::Database::open(
            env.clone(),
            Some("metadata"),
            &lmdb::DatabaseOptions::new(lmdb::db::CREATE),
        )?);
        let image_meta_store = Arc::new(lmdb::Database::open(
            env.clone(),
            Some("metadata-image"),
            &lmdb::DatabaseOptions::new(lmdb::db::CREATE),
        )?);

        Ok(DatabaseContext {
            env,
            binary_store,
            metadata_store,
            image_meta_store,
            image_store,
        })
    }

    #[inline]
    pub fn write_txn(&self) -> CrowResult<lmdb::WriteTransaction<'static>> {
        lmdb::WriteTransaction::new(self.env.clone()).map_err(CrowError::from)
    }

    #[inline]
    pub fn read_txn(&self) -> CrowResult<lmdb::ReadTransaction<'static>> {
        lmdb::ReadTransaction::new(self.env.clone()).map_err(CrowError::from)
    }
}

mod macros {
    macro_rules! response {
        (err $status:literal $msg:literal) => {
            Response::new(
                tiny_http::StatusCode($status),
                vec![],
                $msg.as_bytes(),
                None,
                None,
            )
        };
    }

    macro_rules! single_multipart {
        ($req:expr) => {
            Multipart::from_request($req)
                .ok()
                .and_then(|v| v.into_entry().into_result().ok())
                .flatten()
        };
    }

    macro_rules! fn_to_handler {
        ($f:ident : $handler_name:ident) => {
            pub struct $handler_name;

            impl Handler for $handler_name {
                fn respond<'u>(
                    &self,
                    req: Request,
                    db: &DatabaseContext,
                    params: matchit::Params<'u, 'u>,
                ) -> CrowResult<()> {
                    $f(req, db, params)
                }
            }
        };
    }

    macro_rules! parse_base64_hash {
        ($fr:expr) => {
            $fr.and_then(|s| {
                let mut out: [u8; 32] = [0; 32];
                base64_url::decode_to_slice(s, &mut out).ok()?;
                Some(out)
            })
        };
    }

    macro_rules! some_or_response {
        ($opt:expr, or respond to $req:ident with $response:expr) => {
            match $opt {
                Some(v) => v,
                None => {
                    $req.respond($response)?;
                    return Ok(());
                }
            }
        };
    }

    pub(crate) use fn_to_handler;
    pub(crate) use parse_base64_hash;
    pub(crate) use response;
    pub(crate) use single_multipart;
    pub(crate) use some_or_response;
}

pub(crate) use macros::*;