memory_serve/
asset.rs

1use axum::{
2    http::{
3        HeaderMap, HeaderName, HeaderValue, StatusCode,
4        header::{CONTENT_ENCODING, CONTENT_TYPE, ETAG, IF_NONE_MATCH},
5    },
6    response::{IntoResponse, Response},
7};
8use memory_serve_core::COMPRESS_TYPES;
9use tracing::debug;
10
11use crate::{
12    ServeOptions,
13    util::{compress_brotli, compress_gzip, content_length, decompress_brotli, supports_encoding},
14};
15
16const BROTLI_ENCODING: &str = "br";
17#[allow(clippy::declare_interior_mutable_const)]
18const BROTLI_HEADER: (HeaderName, HeaderValue) =
19    (CONTENT_ENCODING, HeaderValue::from_static(BROTLI_ENCODING));
20const GZIP_ENCODING: &str = "gzip";
21#[allow(clippy::declare_interior_mutable_const)]
22const GZIP_HEADER: (HeaderName, HeaderValue) =
23    (CONTENT_ENCODING, HeaderValue::from_static(GZIP_ENCODING));
24
25/// Represents a static asset that can be served
26#[derive(Debug)]
27pub struct Asset {
28    pub route: &'static str,
29    pub path: &'static str,
30    pub etag: &'static str,
31    pub content_type: &'static str,
32    pub bytes: Option<&'static [u8]>,
33    pub is_compressed: bool,
34}
35
36struct AssetResponse<'t, B> {
37    options: &'t ServeOptions,
38    headers: &'t HeaderMap,
39    status: StatusCode,
40    asset: &'t Asset,
41    etag: &'t str,
42    bytes: B,
43    bytes_len: usize,
44    brotli_bytes: B,
45    brotli_bytes_len: usize,
46    gzip_bytes: B,
47    gzip_bytes_len: usize,
48}
49
50impl<B: IntoResponse> AssetResponse<'_, B> {
51    fn into_response(self) -> Response {
52        let content_type = self.asset.content_type();
53        let cache_control = self.asset.cache_control(self.options);
54        let etag_header = (ETAG, HeaderValue::from_str(self.etag).unwrap());
55
56        if let Some(if_none_match) = self.headers.get(IF_NONE_MATCH)
57            && if_none_match == self.etag
58        {
59            return (
60                StatusCode::NOT_MODIFIED,
61                [content_type, cache_control, etag_header],
62            )
63                .into_response();
64        }
65
66        if self.options.enable_brotli
67            && self.brotli_bytes_len > 0
68            && supports_encoding(self.headers, BROTLI_ENCODING)
69        {
70            return (
71                self.status,
72                [
73                    content_length(self.brotli_bytes_len),
74                    BROTLI_HEADER,
75                    content_type,
76                    cache_control,
77                    etag_header,
78                ],
79                self.brotli_bytes,
80            )
81                .into_response();
82        }
83
84        if self.options.enable_gzip
85            && self.gzip_bytes_len > 0
86            && supports_encoding(self.headers, GZIP_ENCODING)
87        {
88            return (
89                self.status,
90                [
91                    content_length(self.gzip_bytes_len),
92                    GZIP_HEADER,
93                    content_type,
94                    cache_control,
95                    etag_header,
96                ],
97                self.gzip_bytes,
98            )
99                .into_response();
100        }
101
102        (
103            self.status,
104            [
105                content_length(self.bytes_len),
106                content_type,
107                cache_control,
108                etag_header,
109            ],
110            self.bytes,
111        )
112            .into_response()
113    }
114}
115
116impl Asset {
117    fn cache_control(&self, options: &ServeOptions) -> (HeaderName, HeaderValue) {
118        match self.content_type {
119            "text/html" => options.html_cache_control.as_header(),
120            _ => options.cache_control.as_header(),
121        }
122    }
123
124    fn content_type(&self) -> (HeaderName, HeaderValue) {
125        (CONTENT_TYPE, HeaderValue::from_static(self.content_type))
126    }
127
128    /// Get the bytes for the asset, which is possibly compressed in the binary
129    pub(crate) fn leak_bytes(
130        &self,
131        options: &'static ServeOptions,
132    ) -> (&'static [u8], &'static [u8], &'static [u8]) {
133        let mut uncompressed = self.bytes.unwrap_or_default();
134
135        if self.is_compressed {
136            uncompressed = Box::new(decompress_brotli(uncompressed).unwrap_or_default()).leak()
137        }
138
139        let gzip_bytes = if self.is_compressed && options.enable_gzip {
140            Box::new(compress_gzip(uncompressed).unwrap_or_default()).leak()
141        } else {
142            Default::default()
143        };
144
145        let brotli_bytes = if self.is_compressed {
146            self.bytes.unwrap_or_default()
147        } else {
148            Default::default()
149        };
150
151        (uncompressed, brotli_bytes, gzip_bytes)
152    }
153
154    fn dynamic_handler(
155        &self,
156        headers: &HeaderMap,
157        status: StatusCode,
158        options: &ServeOptions,
159    ) -> Response {
160        let Ok(bytes) = std::fs::read(self.path) else {
161            return StatusCode::NOT_FOUND.into_response();
162        };
163
164        let brotli_bytes = if options.enable_brotli && COMPRESS_TYPES.contains(&self.content_type) {
165            compress_brotli(&bytes).unwrap_or_default()
166        } else {
167            Default::default()
168        };
169
170        let gzip_bytes = if options.enable_gzip && COMPRESS_TYPES.contains(&self.content_type) {
171            compress_gzip(&bytes).unwrap_or_default()
172        } else {
173            Default::default()
174        };
175
176        let etag = sha256::digest(&bytes);
177
178        AssetResponse {
179            options,
180            headers,
181            status,
182            asset: self,
183            etag: &etag,
184            bytes_len: bytes.len(),
185            bytes,
186            brotli_bytes_len: brotli_bytes.len(),
187            brotli_bytes,
188            gzip_bytes_len: gzip_bytes.len(),
189            gzip_bytes,
190        }
191        .into_response()
192    }
193
194    pub(super) fn handler(
195        &self,
196        headers: &HeaderMap,
197        status: StatusCode,
198        bytes: &'static [u8],
199        brotli_bytes: &'static [u8],
200        gzip_bytes: &'static [u8],
201        options: &ServeOptions,
202    ) -> Response {
203        if bytes.is_empty() {
204            debug!("using dynamic handler for {}", self.path);
205
206            return self.dynamic_handler(headers, status, options);
207        }
208
209        AssetResponse {
210            options,
211            headers,
212            status,
213            asset: self,
214            etag: self.etag,
215            bytes_len: bytes.len(),
216            bytes,
217            brotli_bytes_len: brotli_bytes.len(),
218            brotli_bytes,
219            gzip_bytes_len: gzip_bytes.len(),
220            gzip_bytes,
221        }
222        .into_response()
223    }
224}