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 tracing::debug;
9
10use crate::{
11    options::ServeOptions,
12    util::{
13        compression::{compress_brotli, compress_gzip, decompress_brotli},
14        headers::{content_length, supports_encoding},
15    },
16};
17
18const BROTLI_ENCODING: &str = "br";
19
20const BROTLI_HEADER: (HeaderName, HeaderValue) =
21    (CONTENT_ENCODING, HeaderValue::from_static(BROTLI_ENCODING));
22
23const GZIP_ENCODING: &str = "gzip";
24
25const GZIP_HEADER: (HeaderName, HeaderValue) =
26    (CONTENT_ENCODING, HeaderValue::from_static(GZIP_ENCODING));
27
28/// Preferred compression for a dynamically served asset.
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30enum OnDemandEncoding {
31    /// Send the original bytes without further compression.
32    Identity,
33    /// Compress the response using brotli.
34    Brotli,
35    /// Compress the response using gzip.
36    Gzip,
37}
38
39/// Represents a static asset that can be served
40#[derive(Debug)]
41pub struct Asset {
42    /// The HTTP route used to serve the asset, e.g. `/index.html`.
43    pub route: &'static str,
44    /// Absolute filesystem path pointing to the source asset on disk.
45    pub path: &'static str,
46    /// Strong validator (SHA-256) used for HTTP caching semantics.
47    pub etag: &'static str,
48    /// MIME type advertised for the asset.
49    pub content_type: &'static str,
50    /// Optional embedded bytes for the asset; `None` when dynamic loading is used.
51    pub bytes: Option<&'static [u8]>,
52    /// Indicates if the embedded bytes are already brotli compressed.
53    pub is_compressed: bool,
54    /// Whether the asset should be compressed before sending to clients.
55    pub should_compress: bool,
56}
57
58/// Aggregates response metadata and payloads for an asset request.
59struct AssetResponse<'t, B> {
60    options: &'t ServeOptions,
61    headers: &'t HeaderMap,
62    status: StatusCode,
63    asset: &'t Asset,
64    etag: &'t str,
65    bytes: B,
66    bytes_len: usize,
67    brotli_bytes: B,
68    brotli_bytes_len: usize,
69    gzip_bytes: B,
70    gzip_bytes_len: usize,
71}
72
73impl<B: IntoResponse> AssetResponse<'_, B> {
74    /// Construct an Axum `Response` from the gathered asset data.
75    fn into_response(self) -> Response {
76        let content_type = self.asset.content_type();
77        let cache_control = self.asset.cache_control(self.options);
78        let etag_header = (ETAG, HeaderValue::from_str(self.etag).unwrap());
79
80        if let Some(if_none_match) = self.headers.get(IF_NONE_MATCH)
81            && if_none_match == self.etag
82        {
83            return (
84                StatusCode::NOT_MODIFIED,
85                [content_type, cache_control, etag_header],
86            )
87                .into_response();
88        }
89
90        if self.options.enable_brotli
91            && self.brotli_bytes_len > 0
92            && supports_encoding(self.headers, BROTLI_ENCODING)
93        {
94            return (
95                self.status,
96                [
97                    content_length(self.brotli_bytes_len),
98                    BROTLI_HEADER,
99                    content_type,
100                    cache_control,
101                    etag_header,
102                ],
103                self.brotli_bytes,
104            )
105                .into_response();
106        }
107
108        if self.options.enable_gzip
109            && self.gzip_bytes_len > 0
110            && supports_encoding(self.headers, GZIP_ENCODING)
111        {
112            return (
113                self.status,
114                [
115                    content_length(self.gzip_bytes_len),
116                    GZIP_HEADER,
117                    content_type,
118                    cache_control,
119                    etag_header,
120                ],
121                self.gzip_bytes,
122            )
123                .into_response();
124        }
125
126        (
127            self.status,
128            [
129                content_length(self.bytes_len),
130                content_type,
131                cache_control,
132                etag_header,
133            ],
134            self.bytes,
135        )
136            .into_response()
137    }
138}
139
140impl Asset {
141    /// Pick the cache policy for the asset based on its MIME type.
142    fn cache_control(&self, options: &ServeOptions) -> (HeaderName, HeaderValue) {
143        match self.content_type {
144            "text/html" => options.html_cache_control.as_header(),
145            _ => options.cache_control.as_header(),
146        }
147    }
148
149    /// Produce the Content-Type header tuple for the asset.
150    fn content_type(&self) -> (HeaderName, HeaderValue) {
151        (CONTENT_TYPE, HeaderValue::from_static(self.content_type))
152    }
153
154    /// Get the bytes for the asset, which is possibly compressed in the binary
155    pub(crate) fn leak_bytes(
156        &self,
157        options: &'static ServeOptions,
158    ) -> (&'static [u8], &'static [u8], &'static [u8]) {
159        let mut uncompressed = self.bytes.unwrap_or_default();
160
161        if self.is_compressed {
162            uncompressed = Box::new(decompress_brotli(uncompressed).unwrap_or_default()).leak()
163        }
164
165        let gzip_bytes = if self.should_compress && options.enable_gzip {
166            Box::new(compress_gzip(uncompressed).unwrap_or_default()).leak()
167        } else {
168            Default::default()
169        };
170
171        let brotli_bytes = if self.should_compress && options.enable_brotli {
172            self.bytes.unwrap_or_default()
173        } else {
174            Default::default()
175        };
176
177        (uncompressed, brotli_bytes, gzip_bytes)
178    }
179
180    /// Load the asset bytes from disk, returning a `404` if the file is missing.
181    fn read_source_bytes(&self) -> Result<Vec<u8>, StatusCode> {
182        std::fs::read(self.path).map_err(|_| StatusCode::NOT_FOUND)
183    }
184
185    /// Decide which compression algorithm (if any) to use for a dynamic request.
186    fn negotiate_dynamic_encoding(
187        &self,
188        headers: &HeaderMap,
189        options: &ServeOptions,
190    ) -> OnDemandEncoding {
191        if !self.should_compress {
192            return OnDemandEncoding::Identity;
193        }
194
195        if options.enable_brotli && supports_encoding(headers, BROTLI_ENCODING) {
196            return OnDemandEncoding::Brotli;
197        }
198
199        if options.enable_gzip && supports_encoding(headers, GZIP_ENCODING) {
200            return OnDemandEncoding::Gzip;
201        }
202
203        OnDemandEncoding::Identity
204    }
205
206    /// Compress the provided bytes according to the negotiated encoding.
207    fn encode_dynamic_bytes(&self, bytes: &[u8], encoding: OnDemandEncoding) -> (Vec<u8>, Vec<u8>) {
208        match encoding {
209            OnDemandEncoding::Brotli => (compress_brotli(bytes).unwrap_or_default(), Vec::new()),
210            OnDemandEncoding::Gzip => (Vec::new(), compress_gzip(bytes).unwrap_or_default()),
211            OnDemandEncoding::Identity => (Vec::new(), Vec::new()),
212        }
213    }
214
215    /// Load an asset from disk and emit a response tailored to client encodings.
216    fn dynamic_handler(
217        &self,
218        headers: &HeaderMap,
219        status: StatusCode,
220        options: &ServeOptions,
221    ) -> Response {
222        let bytes = match self.read_source_bytes() {
223            Ok(bytes) => bytes,
224            Err(status) => return status.into_response(),
225        };
226
227        let encoding = self.negotiate_dynamic_encoding(headers, options);
228        let (brotli_bytes, gzip_bytes) = self.encode_dynamic_bytes(&bytes, encoding);
229
230        let etag = sha256::digest(&bytes);
231
232        AssetResponse {
233            options,
234            headers,
235            status,
236            asset: self,
237            etag: &etag,
238            bytes_len: bytes.len(),
239            bytes,
240            brotli_bytes_len: brotli_bytes.len(),
241            brotli_bytes,
242            gzip_bytes_len: gzip_bytes.len(),
243            gzip_bytes,
244        }
245        .into_response()
246    }
247
248    /// Serve an asset using either embedded bytes or on-demand loading.
249    pub(super) fn handler(
250        &self,
251        headers: &HeaderMap,
252        status: StatusCode,
253        bytes: &'static [u8],
254        brotli_bytes: &'static [u8],
255        gzip_bytes: &'static [u8],
256        options: &ServeOptions,
257    ) -> Response {
258        if bytes.is_empty() {
259            debug!("using dynamic handler for {}", self.path);
260
261            return self.dynamic_handler(headers, status, options);
262        }
263
264        AssetResponse {
265            options,
266            headers,
267            status,
268            asset: self,
269            etag: self.etag,
270            bytes_len: bytes.len(),
271            bytes,
272            brotli_bytes_len: brotli_bytes.len(),
273            brotli_bytes,
274            gzip_bytes_len: gzip_bytes.len(),
275            gzip_bytes,
276        }
277        .into_response()
278    }
279}