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#[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 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}