memory-serve 2.4.0

#![allow(clippy::needless_doctest_main)]
#![doc = include_str!("../README.md")]
use axum::{
    http::{HeaderMap, StatusCode},
    routing::get,
};
use std::future::ready;
use tracing::{info, warn};

mod asset;
mod build;
mod cache_control;
mod load;
mod manifest;
mod options;
mod util;

use crate::util::route::hashed_route;
pub use crate::{
    asset::Asset,
    build::{assets_to_code, load_directory, load_directory_with_embed, load_names_directories},
    cache_control::CacheControl,
    manifest::Manifest,
};

/// Helper struct to create and configure an axum router to serve static
/// files from memory.
#[derive(Debug, Default)]
pub struct MemoryServe {
    options: options::ServeOptions,
    assets: &'static [Asset],
    aliases: Vec<(&'static str, &'static str)>,
}

impl MemoryServe {
    /// Initiate a `MemoryServe` instance, taking the output of the `load!`
    /// macro as an argument. `load!` selects the assets prepared during the
    /// build step.
    pub fn new(assets: &'static [Asset]) -> Self {
        Self {
            assets,
            ..Default::default()
        }
    }

    /// Which static file to serve on the route "/" (the index)
    /// The path (or route) should be relative to the asset directory passed
    /// to `load_directory` (or similar) in your `build.rs`, but prepended with
    /// a slash.
    /// By default this is `Some("/index.html")`
    pub fn index_file(mut self, index_file: Option<&'static str>) -> Self {
        self.options.index_file = index_file;

        self
    }

    /// Whether to serve the corresponding index.html file when a route
    /// matches a subdirectory
    pub fn index_on_subdirectories(mut self, enable: bool) -> Self {
        self.options.index_on_subdirectories = enable;

        self
    }

    /// Which static file to serve when no other routes are matched, also see
    /// [fallback](https://docs.rs/axum/latest/axum/routing/struct.Router.html#method.fallback)
    /// The path (or route) should be relative to the asset directory passed
    /// to `load_directory` (or similar) in your `build.rs`, but prepended with
    /// a slash.
    /// By default this is `None`, which means axum will return an empty
    /// response with a HTTP 404 status code when no route matches.
    pub fn fallback(mut self, fallback: Option<&'static str>) -> Self {
        self.options.fallback = fallback;

        self
    }

    /// What HTTP status code to return when a static file is returned by the
    /// fallback handler.
    pub fn fallback_status(mut self, fallback_status: StatusCode) -> Self {
        self.options.fallback_status = fallback_status;

        self
    }

    /// Whether to enable gzip compression. When set to `true`, clients that
    /// accept gzip compressed files, but not brotli compressed files,
    /// are served gzip compressed files.
    pub fn enable_gzip(mut self, enable_gzip: bool) -> Self {
        self.options.enable_gzip = enable_gzip;

        self
    }

    /// Whether to enable brotli compression. When set to `true`, clients that
    /// accept brotli compressed files are served brotli compressed files.
    pub fn enable_brotli(mut self, enable_brotli: bool) -> Self {
        self.options.enable_brotli = enable_brotli;

        self
    }

    /// Whether to enable clean URLs. When set to `true`, the routing path for
    /// HTML files will not include the extension so that a file located at
    /// "/about.html" maps to "/about" instead of "/about.html".
    pub fn enable_clean_url(mut self, enable_clean_url: bool) -> Self {
        self.options.enable_clean_url = enable_clean_url;

        self
    }

    /// The Cache-Control header to set for HTML files.
    /// See [Cache control](index.html#cache-control) for options.
    pub fn html_cache_control(mut self, html_cache_control: CacheControl) -> Self {
        self.options.html_cache_control = html_cache_control;

        self
    }

    /// The Cache-Control header to set for non-HTML files.
    /// See [Cache control](index.html#cache-control) for options.
    pub fn cache_control(mut self, cache_control: CacheControl) -> Self {
        self.options.cache_control = cache_control;

        self
    }

    /// Create an alias for a route / file
    pub fn add_alias(mut self, from: &'static str, to: &'static str) -> Self {
        self.aliases.push((from, to));

        self
    }

    /// Whether to also serve non-HTML assets on a cache-busted route containing
    /// a content hash, e.g. `/assets/index.3f9a1c2b7d84e6a0.css`. Embedded
    /// assets are served there with [`CacheControl::Long`].
    /// Use [`MemoryServe::manifest`] to look up the hashed routes.
    /// See [Cache busting](index.html#cache-busting).
    /// By default this is `false`.
    pub fn enable_hashed_routes(mut self, enable_hashed_routes: bool) -> Self {
        self.options.enable_hashed_routes = enable_hashed_routes;

        self
    }

    /// Create a [`Manifest`] mapping each asset route to the route it is served
    /// on: the hashed route when [`MemoryServe::enable_hashed_routes`] is set
    /// (except for HTML), the plain route otherwise.
    /// Dynamically served assets are hashed by reading the file, so call this
    /// once at startup and keep the result.
    pub fn manifest(&self) -> Manifest {
        let routes = self
            .assets
            .iter()
            .map(|asset| {
                let served = self
                    .hashed_route(asset)
                    .unwrap_or_else(|| asset.route.to_string());

                (asset.route, served)
            })
            .collect();

        Manifest::new(routes)
    }

    /// The hashed route for a non-HTML asset, if enabled and hashable.
    fn hashed_route(&self, asset: &Asset) -> Option<String> {
        if !self.options.enable_hashed_routes || asset.is_html() {
            return None;
        }

        let hash = asset.content_hash()?;

        Some(hashed_route(asset.route, &hash))
    }

    /// Create an axum `Router` instance that will serve the included static assets
    /// Caution! This method leaks memory. It should only be called once (at startup).
    pub fn into_router<S>(self) -> axum::Router<S>
    where
        S: Clone + Send + Sync + 'static,
    {
        let mut router = axum::Router::new();
        let options: &'static options::ServeOptions = Box::leak(Box::new(self.options));

        // hashed routes are immutable, so embedded assets get a long cache
        let hashed_options: &'static options::ServeOptions =
            Box::leak(Box::new(options::ServeOptions {
                cache_control: CacheControl::Long,
                ..*options
            }));

        // Warn about configuration that silently does nothing because it
        // references a route that no asset provides.
        let route_exists = |route: &str| self.assets.iter().any(|a| a.route == route);

        if let Some(index) = options.index_file
            && !route_exists(index)
        {
            warn!("index_file {index} does not match any asset route, \"/\" will not be served");
        }
        if let Some(fallback) = options.fallback
            && !route_exists(fallback)
        {
            warn!("fallback {fallback} does not match any asset route, it will be ignored");
        }
        for (from, to) in self.aliases.iter() {
            if !route_exists(to) {
                warn!("alias {from} points to {to}, which does not match any asset route");
            }
        }

        for asset in self.assets {
            let (uncompressed_bytes, brotli_bytes, gzip_bytes) = asset.leak_bytes(options);

            if !uncompressed_bytes.is_empty() {
                if !brotli_bytes.is_empty() {
                    info!(
                        "serving {} {} -> {} bytes (compressed)",
                        asset.route,
                        uncompressed_bytes.len(),
                        brotli_bytes.len()
                    );
                } else {
                    info!("serving {} {} bytes", asset.route, uncompressed_bytes.len());
                }
            } else {
                info!("serving {} (dynamically)", asset.route);
            }

            let handler = |headers: HeaderMap| {
                ready(asset.handler(
                    &headers,
                    StatusCode::OK,
                    uncompressed_bytes,
                    brotli_bytes,
                    gzip_bytes,
                    options,
                ))
            };

            if Some(asset.route) == options.fallback {
                info!("serving {} as fallback", asset.route);

                router = router.fallback(|headers: HeaderMap| {
                    ready(asset.handler(
                        &headers,
                        options.fallback_status,
                        uncompressed_bytes,
                        brotli_bytes,
                        gzip_bytes,
                        options,
                    ))
                });
            }

            if let Some(index) = options.index_file {
                if asset.route == index {
                    info!("serving {} as index on /", asset.route);

                    router = router.route("/", get(handler));
                } else if options.index_on_subdirectories && asset.route.ends_with(index) {
                    let path = &asset.route[..asset.route.len() - index.len()];
                    info!("serving {} as index on {}", asset.route, path);

                    router = router.route(path, get(handler));
                }
            }

            let path = if options.enable_clean_url && asset.route.ends_with(".html") {
                &asset.route[..asset.route.len() - 5]
            } else {
                asset.route
            };
            router = router.route(path, get(handler));

            // add all aliases that point to the asset route
            for (from, to) in self.aliases.iter() {
                if *to == asset.route {
                    info!("serving {} on alias {}", asset.route, from);

                    router = router.route(from, get(handler));
                }
            }

            if let Some(hashed) = self.hashed_route(asset) {
                info!("serving {} on hashed route {}", asset.route, hashed);

                // dynamically served files may change, keep the regular cache control
                let hashed_options = if uncompressed_bytes.is_empty() {
                    options
                } else {
                    hashed_options
                };

                router = router.route(
                    &hashed,
                    get(move |headers: HeaderMap| {
                        ready(asset.handler(
                            &headers,
                            StatusCode::OK,
                            uncompressed_bytes,
                            brotli_bytes,
                            gzip_bytes,
                            hashed_options,
                        ))
                    }),
                );
            }
        }

        router
    }
}

#[cfg(test)]
mod tests {
    use axum::{
        Router,
        body::Body,
        http::{
            self, HeaderMap, HeaderName, HeaderValue, Request, StatusCode,
            header::{self, CACHE_CONTROL, CONTENT_ENCODING, CONTENT_LENGTH},
        },
    };
    use std::sync::LazyLock;
    use tower::ServiceExt;

    use crate::{self as memory_serve, Asset, CacheControl, MemoryServe};

    static ASSETS: LazyLock<&'static [Asset]> =
        LazyLock::new(|| memory_serve::build::load_test_assets("./static"));

    macro_rules! test_load {
        () => {{ MemoryServe::new(*ASSETS) }};
    }

    async fn get(
        router: Router,
        path: &str,
        key: &str,
        value: &str,
    ) -> (StatusCode, HeaderMap<HeaderValue>) {
        let response = router
            .oneshot(
                Request::builder()
                    .method(http::Method::GET)
                    .header(key, value)
                    .uri(path)
                    .body(Body::empty())
                    .unwrap(),
            )
            .await
            .unwrap();

        (response.status(), response.headers().to_owned())
    }

    fn get_header<'s>(headers: &'s HeaderMap, name: &HeaderName) -> &'s str {
        headers.get(name).unwrap().to_str().unwrap()
    }

    #[tokio::test]
    async fn test_load_assets() {
        let routes: Vec<&str> = ASSETS.iter().map(|a| a.route).collect();
        let content_types: Vec<&str> = ASSETS.iter().map(|a| a.content_type).collect();
        let etags: Vec<&str> = ASSETS.iter().map(|a| a.etag).collect();

        assert_eq!(
            routes,
            [
                "/about.html",
                "/assets/icon.jpg",
                "/assets/index.css",
                "/assets/index.js",
                "/assets/stars.svg",
                "/blog/index.html",
                "/index.html"
            ]
        );
        assert_eq!(
            content_types,
            [
                "text/html",
                "image/jpeg",
                "text/css",
                "text/javascript",
                "image/svg+xml",
                "text/html",
                "text/html"
            ]
        );
        if cfg!(debug_assertions) && !cfg!(feature = "force-embed") {
            assert_eq!(etags, ["", "", "", "", "", "", ""]);
        } else {
            assert_eq!(
                etags,
                [
                    "56a0dcb83ec56b6c967966a1c06c7b1392e261069d0844aa4e910ca5c1e8cf58",
                    "e64f4683bf82d854df40b7246666f6f0816666ad8cd886a8e159535896eb03d6",
                    "ec4edeea111c854901385011f403e1259e3f1ba016dcceabb6d566316be3677b",
                    "86a7fdfd19700843e5f7344a63d27e0b729c2554c8572903ceee71f5658d2ecf",
                    "bd9dccc152de48cb7bedc35b9748ceeade492f6f904710f9c5d480bd6299cc7d",
                    "89e9873a8e49f962fe83ad2bfe6ac9b21ef7c1b4040b99c34eb783dccbadebc5",
                    "0639dc8aac157b58c74f65bbb026b2fd42bc81d9a0a64141df456fa23c214537"
                ]
            );
        }
    }

    #[tokio::test]
    async fn if_none_match_handling() {
        let memory_router = test_load!().into_router();
        let (code, headers) =
            get(memory_router.clone(), "/index.html", "accept", "text/html").await;
        let etag: &str = headers.get(header::ETAG).unwrap().to_str().unwrap();

        assert_eq!(code, 200);
        assert_eq!(
            etag,
            "\"0639dc8aac157b58c74f65bbb026b2fd42bc81d9a0a64141df456fa23c214537\""
        );

        let (code, headers) = get(memory_router, "/index.html", "If-None-Match", etag).await;
        let length = get_header(&headers, &CONTENT_LENGTH);

        assert_eq!(code, 304);
        assert_eq!(length.parse::<i32>().unwrap(), 0);
    }

    #[tokio::test]
    async fn brotli_compression() {
        let memory_router = test_load!().enable_brotli(true).into_router();
        let (code, headers) = get(
            memory_router.clone(),
            "/index.html",
            "accept-encoding",
            "br",
        )
        .await;
        let encoding = get_header(&headers, &CONTENT_ENCODING);
        let length = get_header(&headers, &CONTENT_LENGTH);

        assert_eq!(code, 200);
        assert_eq!(encoding, "br");
        assert_eq!(length.parse::<i32>().unwrap(), 178);

        // check disable compression
        let memory_router = test_load!().enable_brotli(false).into_router();
        let (code, headers) = get(
            memory_router.clone(),
            "/index.html",
            "accept-encoding",
            "br",
        )
        .await;
        let length: &str = get_header(&headers, &CONTENT_LENGTH);

        assert_eq!(code, 200);
        assert_eq!(length.parse::<i32>().unwrap(), 437);
    }

    #[tokio::test]
    async fn gzip_compression() {
        let memory_router = test_load!().enable_gzip(true).into_router();
        let (code, headers) = get(
            memory_router.clone(),
            "/index.html",
            "accept-encoding",
            "gzip",
        )
        .await;

        let encoding = get_header(&headers, &CONTENT_ENCODING);
        let length = get_header(&headers, &CONTENT_LENGTH);

        assert_eq!(code, 200);
        assert_eq!(encoding, "gzip");
        assert_eq!(length.parse::<i32>().unwrap(), 274);

        // check disable compression
        let memory_router = test_load!().enable_gzip(false).into_router();
        let (code, headers) = get(
            memory_router.clone(),
            "/index.html",
            "accept-encoding",
            "gzip",
        )
        .await;
        let length: &str = get_header(&headers, &CONTENT_LENGTH);

        assert_eq!(code, 200);
        assert_eq!(length.parse::<i32>().unwrap(), 437);
    }

    #[tokio::test]
    async fn encoding_preference() {
        // Both encodings enabled; the client prefers gzip over brotli.
        let memory_router = test_load!()
            .enable_brotli(true)
            .enable_gzip(true)
            .into_router();
        let (code, headers) = get(
            memory_router.clone(),
            "/index.html",
            "accept-encoding",
            "br;q=0.1, gzip;q=1.0",
        )
        .await;
        assert_eq!(code, 200);
        assert_eq!(get_header(&headers, &CONTENT_ENCODING), "gzip");

        // The reverse preference selects brotli.
        let (code, headers) = get(
            memory_router.clone(),
            "/index.html",
            "accept-encoding",
            "br;q=0.8, gzip;q=0.7",
        )
        .await;
        assert_eq!(code, 200);
        assert_eq!(get_header(&headers, &CONTENT_ENCODING), "br");
    }

    #[tokio::test]
    async fn index_file() {
        let memory_router = test_load!().index_file(None).into_router();

        let (code, _) = get(memory_router.clone(), "/", "accept", "*").await;
        assert_eq!(code, 404);

        let memory_router = test_load!().index_file(Some("/index.html")).into_router();

        let (code, _) = get(memory_router.clone(), "/", "accept", "*").await;
        assert_eq!(code, 200);
    }

    #[tokio::test]
    async fn index_file_on_subdirs() {
        let memory_router = test_load!()
            .index_file(Some("/index.html"))
            .index_on_subdirectories(false)
            .into_router();

        let (code, _) = get(memory_router.clone(), "/blog", "accept", "*").await;
        assert_eq!(code, 404);

        let memory_router = test_load!()
            .index_file(Some("/index.html"))
            .index_on_subdirectories(true)
            .into_router();

        let (code, _) = get(memory_router.clone(), "/blog", "accept", "*").await;
        assert_eq!(code, 200);
    }

    #[tokio::test]
    async fn clean_url() {
        let memory_router = test_load!().enable_clean_url(true).into_router();

        let (code, _) = get(memory_router.clone(), "/about.html", "accept", "*").await;
        assert_eq!(code, 404);

        let (code, _) = get(memory_router.clone(), "/about", "accept", "*").await;
        assert_eq!(code, 200);
    }

    #[tokio::test]
    async fn fallback() {
        let memory_router = test_load!().into_router();
        let (code, _) = get(memory_router.clone(), "/foobar", "accept", "*").await;
        assert_eq!(code, 404);

        let memory_router = test_load!().fallback(Some("/index.html")).into_router();
        let (code, headers) = get(memory_router.clone(), "/foobar", "accept", "*").await;
        let length = get_header(&headers, &CONTENT_LENGTH);
        assert_eq!(code, 404);
        assert_eq!(length.parse::<i32>().unwrap(), 437);

        let memory_router = test_load!()
            .fallback(Some("/index.html"))
            .fallback_status(StatusCode::OK)
            .into_router();
        let (code, headers) = get(memory_router.clone(), "/foobar", "accept", "*").await;
        let length = get_header(&headers, &CONTENT_LENGTH);
        assert_eq!(code, 200);
        assert_eq!(length.parse::<i32>().unwrap(), 437);
    }

    #[tokio::test]
    async fn cache_control() {
        async fn check_cache_control(cache_control: CacheControl, expected: &str) {
            let memory_router = test_load!().cache_control(cache_control).into_router();

            let (code, headers) =
                get(memory_router.clone(), "/assets/icon.jpg", "accept", "*").await;

            let cache_control = get_header(&headers, &CACHE_CONTROL);
            assert_eq!(code, 200);
            assert_eq!(cache_control, expected);
        }

        check_cache_control(
            CacheControl::NoCache,
            CacheControl::NoCache.as_header().1.to_str().unwrap(),
        )
        .await;
        check_cache_control(
            CacheControl::Short,
            CacheControl::Short.as_header().1.to_str().unwrap(),
        )
        .await;
        check_cache_control(
            CacheControl::Medium,
            CacheControl::Medium.as_header().1.to_str().unwrap(),
        )
        .await;
        check_cache_control(
            CacheControl::Long,
            CacheControl::Long.as_header().1.to_str().unwrap(),
        )
        .await;

        async fn check_html_cache_control(cache_control: CacheControl, expected: &str) {
            let memory_router = test_load!().html_cache_control(cache_control).into_router();

            let (code, headers) = get(memory_router.clone(), "/index.html", "accept", "*").await;
            let cache_control = get_header(&headers, &CACHE_CONTROL);
            assert_eq!(code, 200);
            assert_eq!(cache_control, expected);
        }

        check_html_cache_control(
            CacheControl::NoCache,
            CacheControl::NoCache.as_header().1.to_str().unwrap(),
        )
        .await;
        check_html_cache_control(
            CacheControl::Short,
            CacheControl::Short.as_header().1.to_str().unwrap(),
        )
        .await;
        check_html_cache_control(
            CacheControl::Medium,
            CacheControl::Medium.as_header().1.to_str().unwrap(),
        )
        .await;
        check_html_cache_control(
            CacheControl::Long,
            CacheControl::Long.as_header().1.to_str().unwrap(),
        )
        .await;
    }

    #[tokio::test]
    async fn hashed_routes() {
        let memory_serve = test_load!().enable_hashed_routes(true);
        let manifest = memory_serve.manifest();
        let memory_router = memory_serve.into_router();

        // non-HTML assets map to a route with a 16 character hex hash before
        // the extension
        let hashed = manifest.get("/assets/index.css").unwrap();
        let hash = hashed
            .strip_prefix("/assets/index.")
            .and_then(|rest| rest.strip_suffix(".css"))
            .unwrap();
        assert_eq!(hash.len(), 16);
        assert!(hash.chars().all(|c| c.is_ascii_hexdigit()));

        if !cfg!(debug_assertions) || cfg!(feature = "force-embed") {
            assert_eq!(hashed, "/assets/index.ec4edeea111c8549.css");
        }

        // HTML files keep their plain route, unknown routes are absent
        assert_eq!(manifest.get("/index.html"), Some("/index.html"));
        assert_eq!(manifest.get("/blog/index.html"), Some("/blog/index.html"));
        assert_eq!(manifest.get("/missing.css"), None);
        assert_eq!(manifest.len(), ASSETS.len());

        // the hashed route serves the same file
        let (code, headers) = get(memory_router.clone(), hashed, "accept", "*").await;
        assert_eq!(code, 200);
        assert_eq!(get_header(&headers, &header::CONTENT_TYPE), "text/css");
        assert_eq!(get_header(&headers, &CONTENT_LENGTH), "1552");

        // embedded assets are immutable on the hashed route, dynamically
        // served assets keep the regular cache control
        let expected_cache_control = if !cfg!(debug_assertions) || cfg!(feature = "force-embed") {
            CacheControl::Long
        } else {
            CacheControl::Medium
        };
        assert_eq!(
            get_header(&headers, &CACHE_CONTROL),
            expected_cache_control.as_header().1.to_str().unwrap()
        );

        // the plain route is still served with the configured cache control
        let (code, headers) = get(memory_router.clone(), "/assets/index.css", "accept", "*").await;
        assert_eq!(code, 200);
        assert_eq!(
            get_header(&headers, &CACHE_CONTROL),
            CacheControl::Medium.as_header().1.to_str().unwrap()
        );

        // every manifest entry resolves
        for (_, served) in manifest.iter() {
            let (code, _) = get(memory_router.clone(), served, "accept", "*").await;
            assert_eq!(code, 200, "{served} should be served");
        }
    }

    #[tokio::test]
    async fn hashed_routes_disabled() {
        let memory_serve = test_load!();
        let manifest = memory_serve.manifest();
        let memory_router = memory_serve.into_router();

        // without hashed routes every asset maps to its plain route
        assert_eq!(manifest.get("/assets/index.css"), Some("/assets/index.css"));
        assert_eq!(manifest.get("/index.html"), Some("/index.html"));
        assert_eq!(manifest.len(), ASSETS.len());

        let (code, _) = get(
            memory_router.clone(),
            "/assets/index.ec4edeea111c8549.css",
            "accept",
            "*",
        )
        .await;
        assert_eq!(code, 404);
    }

    #[tokio::test]
    async fn aliases() {
        let memory_router = test_load!()
            .add_alias("/foobar", "/index.html")
            .add_alias("/baz", "/index.html")
            .into_router();
        let (code, _) = get(memory_router.clone(), "/foobar", "accept", "*").await;
        assert_eq!(code, 200);

        let (code, _) = get(memory_router.clone(), "/baz", "accept", "*").await;
        assert_eq!(code, 200);

        let (code, _) = get(memory_router.clone(), "/index.html", "accept", "*").await;
        assert_eq!(code, 200);

        let (code, _) = get(memory_router.clone(), "/barfoo", "accept", "*").await;
        assert_eq!(code, 404);
    }
}