Skip to main content

memory_serve/
lib.rs

1#![allow(clippy::needless_doctest_main)]
2#![doc = include_str!("../README.md")]
3use axum::{
4    http::{HeaderMap, StatusCode},
5    routing::get,
6};
7use std::future::ready;
8use tracing::{info, warn};
9
10mod asset;
11mod build;
12mod cache_control;
13mod load;
14mod options;
15mod util;
16
17pub use crate::{
18    asset::Asset,
19    build::{assets_to_code, load_directory, load_directory_with_embed, load_names_directories},
20    cache_control::CacheControl,
21};
22
23/// Helper struct to create and configure an axum to serve static files from
24/// memory.
25#[derive(Debug, Default)]
26pub struct MemoryServe {
27    options: options::ServeOptions,
28    assets: &'static [Asset],
29    aliases: Vec<(&'static str, &'static str)>,
30}
31
32impl MemoryServe {
33    /// Initiate a `MemoryServe` instance, taking the output of the `load!`
34    /// macro as an argument. `load!` selects the assets prepared during the
35    /// build step.
36    pub fn new(assets: &'static [Asset]) -> Self {
37        Self {
38            assets,
39            ..Default::default()
40        }
41    }
42
43    /// Which static file to serve on the route "/" (the index)
44    /// The path (or route) should be relative to the asset directory passed
45    /// to `load_directory` (or similar) in your `build.rs`, but prepended with
46    /// a slash.
47    /// By default this is `Some("/index.html")`
48    pub fn index_file(mut self, index_file: Option<&'static str>) -> Self {
49        self.options.index_file = index_file;
50
51        self
52    }
53
54    /// Whether to serve the corresponding index.html file when a route
55    /// matches a subdirectory
56    pub fn index_on_subdirectories(mut self, enable: bool) -> Self {
57        self.options.index_on_subdirectories = enable;
58
59        self
60    }
61
62    /// Which static file to serve when no other routes are matched, also see
63    /// [fallback](https://docs.rs/axum/latest/axum/routing/struct.Router.html#method.fallback)
64    /// The path (or route) should be relative to the asset directory passed
65    /// to `load_directory` (or similar) in your `build.rs`, but prepended with
66    /// a slash.
67    /// By default this is `None`, which means axum will return an empty
68    /// response with a HTTP 404 status code when no route matches.
69    pub fn fallback(mut self, fallback: Option<&'static str>) -> Self {
70        self.options.fallback = fallback;
71
72        self
73    }
74
75    /// What HTTP status code to return when a static file is returned by the
76    /// fallback handler.
77    pub fn fallback_status(mut self, fallback_status: StatusCode) -> Self {
78        self.options.fallback_status = fallback_status;
79
80        self
81    }
82
83    /// Whether to enable gzip compression. When set to `true`, clients that
84    /// accept gzip compressed files, but not brotli compressed files,
85    /// are served gzip compressed files.
86    pub fn enable_gzip(mut self, enable_gzip: bool) -> Self {
87        self.options.enable_gzip = enable_gzip;
88
89        self
90    }
91
92    /// Whether to enable brotli compression. When set to `true`, clients that
93    /// accept brotli compressed files are served brotli compressed files.
94    pub fn enable_brotli(mut self, enable_brotli: bool) -> Self {
95        self.options.enable_brotli = enable_brotli;
96
97        self
98    }
99
100    /// Whether to enable clean URLs. When set to `true`, the routing path for
101    /// HTML files will not include the extension so that a file located at
102    /// "/about.html" maps to "/about" instead of "/about.html".
103    pub fn enable_clean_url(mut self, enable_clean_url: bool) -> Self {
104        self.options.enable_clean_url = enable_clean_url;
105
106        self
107    }
108
109    /// The Cache-Control header to set for HTML files.
110    /// See [Cache control](index.html#cache-control) for options.
111    pub fn html_cache_control(mut self, html_cache_control: CacheControl) -> Self {
112        self.options.html_cache_control = html_cache_control;
113
114        self
115    }
116
117    /// Cache header to non-HTML files.
118    /// See [Cache control](index.html#cache-control) for options.
119    pub fn cache_control(mut self, cache_control: CacheControl) -> Self {
120        self.options.cache_control = cache_control;
121
122        self
123    }
124
125    /// Create an alias for a route / file
126    pub fn add_alias(mut self, from: &'static str, to: &'static str) -> Self {
127        self.aliases.push((from, to));
128
129        self
130    }
131
132    /// Create an axum `Router` instance that will serve the included static assets
133    /// Caution! This method leaks memory. It should only be called once (at startup).
134    pub fn into_router<S>(self) -> axum::Router<S>
135    where
136        S: Clone + Send + Sync + 'static,
137    {
138        let mut router = axum::Router::new();
139        let options = Box::leak(Box::new(self.options));
140
141        // Warn about configuration that silently does nothing because it
142        // references a route that no asset provides.
143        let route_exists = |route: &str| self.assets.iter().any(|a| a.route == route);
144
145        if let Some(index) = options.index_file
146            && !route_exists(index)
147        {
148            warn!("index_file {index} does not match any asset route, \"/\" will not be served");
149        }
150        if let Some(fallback) = options.fallback
151            && !route_exists(fallback)
152        {
153            warn!("fallback {fallback} does not match any asset route, it will be ignored");
154        }
155        for (from, to) in self.aliases.iter() {
156            if !route_exists(to) {
157                warn!("alias {from} points to {to}, which does not match any asset route");
158            }
159        }
160
161        for asset in self.assets {
162            let (uncompressed_bytes, brotli_bytes, gzip_bytes) = asset.leak_bytes(options);
163
164            if !uncompressed_bytes.is_empty() {
165                if asset.is_compressed {
166                    info!(
167                        "serving {} {} -> {} bytes (compressed)",
168                        asset.route,
169                        uncompressed_bytes.len(),
170                        brotli_bytes.len()
171                    );
172                } else {
173                    info!("serving {} {} bytes", asset.route, uncompressed_bytes.len());
174                }
175            } else {
176                info!("serving {} (dynamically)", asset.route);
177            }
178
179            let handler = |headers: HeaderMap| {
180                ready(asset.handler(
181                    &headers,
182                    StatusCode::OK,
183                    uncompressed_bytes,
184                    brotli_bytes,
185                    gzip_bytes,
186                    options,
187                ))
188            };
189
190            if Some(asset.route) == options.fallback {
191                info!("serving {} as fallback", asset.route);
192
193                router = router.fallback(|headers: HeaderMap| {
194                    ready(asset.handler(
195                        &headers,
196                        options.fallback_status,
197                        uncompressed_bytes,
198                        brotli_bytes,
199                        gzip_bytes,
200                        options,
201                    ))
202                });
203            }
204
205            if let Some(index) = options.index_file {
206                if asset.route == index {
207                    info!("serving {} as index on /", asset.route);
208
209                    router = router.route("/", get(handler));
210                } else if options.index_on_subdirectories && asset.route.ends_with(index) {
211                    let path = &asset.route[..asset.route.len() - index.len()];
212                    info!("serving {} as index on {}", asset.route, path);
213
214                    router = router.route(path, get(handler));
215                }
216            }
217
218            let path = if options.enable_clean_url && asset.route.ends_with(".html") {
219                &asset.route[..asset.route.len() - 5]
220            } else {
221                asset.route
222            };
223            router = router.route(path, get(handler));
224
225            // add all aliases that point to the asset route
226            for (from, to) in self.aliases.iter() {
227                if *to == asset.route {
228                    info!("serving {} on alias {}", asset.route, from);
229
230                    router = router.route(from, get(handler));
231                }
232            }
233        }
234
235        router
236    }
237}
238
239#[cfg(test)]
240mod tests {
241    use axum::{
242        Router,
243        body::Body,
244        http::{
245            self, HeaderMap, HeaderName, HeaderValue, Request, StatusCode,
246            header::{self, CACHE_CONTROL, CONTENT_ENCODING, CONTENT_LENGTH},
247        },
248    };
249    use std::sync::LazyLock;
250    use tower::ServiceExt;
251
252    use crate::{self as memory_serve, Asset, CacheControl, MemoryServe};
253
254    static ASSETS: LazyLock<&'static [Asset]> =
255        LazyLock::new(|| memory_serve::build::load_test_assets("./static"));
256
257    macro_rules! test_load {
258        () => {{ MemoryServe::new(*ASSETS) }};
259    }
260
261    async fn get(
262        router: Router,
263        path: &str,
264        key: &str,
265        value: &str,
266    ) -> (StatusCode, HeaderMap<HeaderValue>) {
267        let response = router
268            .oneshot(
269                Request::builder()
270                    .method(http::Method::GET)
271                    .header(key, value)
272                    .uri(path)
273                    .body(Body::empty())
274                    .unwrap(),
275            )
276            .await
277            .unwrap();
278
279        (response.status(), response.headers().to_owned())
280    }
281
282    fn get_header<'s>(headers: &'s HeaderMap, name: &HeaderName) -> &'s str {
283        headers.get(name).unwrap().to_str().unwrap()
284    }
285
286    #[tokio::test]
287    async fn test_load_assets() {
288        let routes: Vec<&str> = ASSETS.iter().map(|a| a.route).collect();
289        let content_types: Vec<&str> = ASSETS.iter().map(|a| a.content_type).collect();
290        let etags: Vec<&str> = ASSETS.iter().map(|a| a.etag).collect();
291
292        assert_eq!(
293            routes,
294            [
295                "/about.html",
296                "/assets/icon.jpg",
297                "/assets/index.css",
298                "/assets/index.js",
299                "/assets/stars.svg",
300                "/blog/index.html",
301                "/index.html"
302            ]
303        );
304        assert_eq!(
305            content_types,
306            [
307                "text/html",
308                "image/jpeg",
309                "text/css",
310                "text/javascript",
311                "image/svg+xml",
312                "text/html",
313                "text/html"
314            ]
315        );
316        if cfg!(debug_assertions) && !cfg!(feature = "force-embed") {
317            assert_eq!(etags, ["", "", "", "", "", "", ""]);
318        } else {
319            assert_eq!(
320                etags,
321                [
322                    "56a0dcb83ec56b6c967966a1c06c7b1392e261069d0844aa4e910ca5c1e8cf58",
323                    "e64f4683bf82d854df40b7246666f6f0816666ad8cd886a8e159535896eb03d6",
324                    "ec4edeea111c854901385011f403e1259e3f1ba016dcceabb6d566316be3677b",
325                    "86a7fdfd19700843e5f7344a63d27e0b729c2554c8572903ceee71f5658d2ecf",
326                    "bd9dccc152de48cb7bedc35b9748ceeade492f6f904710f9c5d480bd6299cc7d",
327                    "89e9873a8e49f962fe83ad2bfe6ac9b21ef7c1b4040b99c34eb783dccbadebc5",
328                    "0639dc8aac157b58c74f65bbb026b2fd42bc81d9a0a64141df456fa23c214537"
329                ]
330            );
331        }
332    }
333
334    #[tokio::test]
335    async fn if_none_match_handling() {
336        let memory_router = test_load!().into_router();
337        let (code, headers) =
338            get(memory_router.clone(), "/index.html", "accept", "text/html").await;
339        let etag: &str = headers.get(header::ETAG).unwrap().to_str().unwrap();
340
341        assert_eq!(code, 200);
342        assert_eq!(
343            etag,
344            "\"0639dc8aac157b58c74f65bbb026b2fd42bc81d9a0a64141df456fa23c214537\""
345        );
346
347        let (code, headers) = get(memory_router, "/index.html", "If-None-Match", etag).await;
348        let length = get_header(&headers, &CONTENT_LENGTH);
349
350        assert_eq!(code, 304);
351        assert_eq!(length.parse::<i32>().unwrap(), 0);
352    }
353
354    #[tokio::test]
355    async fn brotli_compression() {
356        let memory_router = test_load!().enable_brotli(true).into_router();
357        let (code, headers) = get(
358            memory_router.clone(),
359            "/index.html",
360            "accept-encoding",
361            "br",
362        )
363        .await;
364        let encoding = get_header(&headers, &CONTENT_ENCODING);
365        let length = get_header(&headers, &CONTENT_LENGTH);
366
367        assert_eq!(code, 200);
368        assert_eq!(encoding, "br");
369        assert_eq!(length.parse::<i32>().unwrap(), 178);
370
371        // check disable compression
372        let memory_router = test_load!().enable_brotli(false).into_router();
373        let (code, headers) = get(
374            memory_router.clone(),
375            "/index.html",
376            "accept-encoding",
377            "br",
378        )
379        .await;
380        let length: &str = get_header(&headers, &CONTENT_LENGTH);
381
382        assert_eq!(code, 200);
383        assert_eq!(length.parse::<i32>().unwrap(), 437);
384    }
385
386    #[tokio::test]
387    async fn gzip_compression() {
388        let memory_router = test_load!().enable_gzip(true).into_router();
389        let (code, headers) = get(
390            memory_router.clone(),
391            "/index.html",
392            "accept-encoding",
393            "gzip",
394        )
395        .await;
396
397        let encoding = get_header(&headers, &CONTENT_ENCODING);
398        let length = get_header(&headers, &CONTENT_LENGTH);
399
400        assert_eq!(code, 200);
401        assert_eq!(encoding, "gzip");
402        assert_eq!(length.parse::<i32>().unwrap(), 274);
403
404        // check disable compression
405        let memory_router = test_load!().enable_gzip(false).into_router();
406        let (code, headers) = get(
407            memory_router.clone(),
408            "/index.html",
409            "accept-encoding",
410            "gzip",
411        )
412        .await;
413        let length: &str = get_header(&headers, &CONTENT_LENGTH);
414
415        assert_eq!(code, 200);
416        assert_eq!(length.parse::<i32>().unwrap(), 437);
417    }
418
419    #[tokio::test]
420    async fn encoding_preference() {
421        // Both encodings enabled; the client prefers gzip over brotli.
422        let memory_router = test_load!()
423            .enable_brotli(true)
424            .enable_gzip(true)
425            .into_router();
426        let (code, headers) = get(
427            memory_router.clone(),
428            "/index.html",
429            "accept-encoding",
430            "br;q=0.1, gzip;q=1.0",
431        )
432        .await;
433        assert_eq!(code, 200);
434        assert_eq!(get_header(&headers, &CONTENT_ENCODING), "gzip");
435
436        // The reverse preference selects brotli.
437        let (code, headers) = get(
438            memory_router.clone(),
439            "/index.html",
440            "accept-encoding",
441            "br;q=0.8, gzip;q=0.7",
442        )
443        .await;
444        assert_eq!(code, 200);
445        assert_eq!(get_header(&headers, &CONTENT_ENCODING), "br");
446    }
447
448    #[tokio::test]
449    async fn index_file() {
450        let memory_router = test_load!().index_file(None).into_router();
451
452        let (code, _) = get(memory_router.clone(), "/", "accept", "*").await;
453        assert_eq!(code, 404);
454
455        let memory_router = test_load!().index_file(Some("/index.html")).into_router();
456
457        let (code, _) = get(memory_router.clone(), "/", "accept", "*").await;
458        assert_eq!(code, 200);
459    }
460
461    #[tokio::test]
462    async fn index_file_on_subdirs() {
463        let memory_router = test_load!()
464            .index_file(Some("/index.html"))
465            .index_on_subdirectories(false)
466            .into_router();
467
468        let (code, _) = get(memory_router.clone(), "/blog", "accept", "*").await;
469        assert_eq!(code, 404);
470
471        let memory_router = test_load!()
472            .index_file(Some("/index.html"))
473            .index_on_subdirectories(true)
474            .into_router();
475
476        let (code, _) = get(memory_router.clone(), "/blog", "accept", "*").await;
477        assert_eq!(code, 200);
478    }
479
480    #[tokio::test]
481    async fn clean_url() {
482        let memory_router = test_load!().enable_clean_url(true).into_router();
483
484        let (code, _) = get(memory_router.clone(), "/about.html", "accept", "*").await;
485        assert_eq!(code, 404);
486
487        let (code, _) = get(memory_router.clone(), "/about", "accept", "*").await;
488        assert_eq!(code, 200);
489    }
490
491    #[tokio::test]
492    async fn fallback() {
493        let memory_router = test_load!().into_router();
494        let (code, _) = get(memory_router.clone(), "/foobar", "accept", "*").await;
495        assert_eq!(code, 404);
496
497        let memory_router = test_load!().fallback(Some("/index.html")).into_router();
498        let (code, headers) = get(memory_router.clone(), "/foobar", "accept", "*").await;
499        let length = get_header(&headers, &CONTENT_LENGTH);
500        assert_eq!(code, 404);
501        assert_eq!(length.parse::<i32>().unwrap(), 437);
502
503        let memory_router = test_load!()
504            .fallback(Some("/index.html"))
505            .fallback_status(StatusCode::OK)
506            .into_router();
507        let (code, headers) = get(memory_router.clone(), "/foobar", "accept", "*").await;
508        let length = get_header(&headers, &CONTENT_LENGTH);
509        assert_eq!(code, 200);
510        assert_eq!(length.parse::<i32>().unwrap(), 437);
511    }
512
513    #[tokio::test]
514    async fn cache_control() {
515        async fn check_cache_control(cache_control: CacheControl, expected: &str) {
516            let memory_router = test_load!().cache_control(cache_control).into_router();
517
518            let (code, headers) =
519                get(memory_router.clone(), "/assets/icon.jpg", "accept", "*").await;
520
521            let cache_control = get_header(&headers, &CACHE_CONTROL);
522            assert_eq!(code, 200);
523            assert_eq!(cache_control, expected);
524        }
525
526        check_cache_control(
527            CacheControl::NoCache,
528            CacheControl::NoCache.as_header().1.to_str().unwrap(),
529        )
530        .await;
531        check_cache_control(
532            CacheControl::Short,
533            CacheControl::Short.as_header().1.to_str().unwrap(),
534        )
535        .await;
536        check_cache_control(
537            CacheControl::Medium,
538            CacheControl::Medium.as_header().1.to_str().unwrap(),
539        )
540        .await;
541        check_cache_control(
542            CacheControl::Long,
543            CacheControl::Long.as_header().1.to_str().unwrap(),
544        )
545        .await;
546
547        async fn check_html_cache_control(cache_control: CacheControl, expected: &str) {
548            let memory_router = test_load!().html_cache_control(cache_control).into_router();
549
550            let (code, headers) = get(memory_router.clone(), "/index.html", "accept", "*").await;
551            let cache_control = get_header(&headers, &CACHE_CONTROL);
552            assert_eq!(code, 200);
553            assert_eq!(cache_control, expected);
554        }
555
556        check_html_cache_control(
557            CacheControl::NoCache,
558            CacheControl::NoCache.as_header().1.to_str().unwrap(),
559        )
560        .await;
561        check_html_cache_control(
562            CacheControl::Short,
563            CacheControl::Short.as_header().1.to_str().unwrap(),
564        )
565        .await;
566        check_html_cache_control(
567            CacheControl::Medium,
568            CacheControl::Medium.as_header().1.to_str().unwrap(),
569        )
570        .await;
571        check_html_cache_control(
572            CacheControl::Long,
573            CacheControl::Long.as_header().1.to_str().unwrap(),
574        )
575        .await;
576    }
577
578    #[tokio::test]
579    async fn aliases() {
580        let memory_router = test_load!()
581            .add_alias("/foobar", "/index.html")
582            .add_alias("/baz", "/index.html")
583            .into_router();
584        let (code, _) = get(memory_router.clone(), "/foobar", "accept", "*").await;
585        assert_eq!(code, 200);
586
587        let (code, _) = get(memory_router.clone(), "/baz", "accept", "*").await;
588        assert_eq!(code, 200);
589
590        let (code, _) = get(memory_router.clone(), "/index.html", "accept", "*").await;
591        assert_eq!(code, 200);
592
593        let (code, _) = get(memory_router.clone(), "/barfoo", "accept", "*").await;
594        assert_eq!(code, 404);
595    }
596}