poem_http_common/
embed_file.rs

1use std::marker::PhantomData;
2
3use poem::{
4    http::{header, Method, StatusCode},
5    Endpoint, Error, Request, Response,
6};
7use rust_embed::RustEmbed;
8
9/// An endpoint that wraps a single file from a `rust-embed` bundle.
10pub struct EmbeddedFileEndpoint<E: RustEmbed + Send + Sync> {
11    _embed: PhantomData<E>,
12    path: String,
13}
14
15impl<E: RustEmbed + Send + Sync> EmbeddedFileEndpoint<E> {
16    /// Create a new `EmbeddedFileEndpoint` from a `rust-embed` bundle.
17    ///
18    /// `path` - relative path within the bundle.
19    pub fn new(path: &str) -> Self {
20        EmbeddedFileEndpoint {
21            _embed: PhantomData,
22            path: path.to_owned(),
23        }
24    }
25}
26
27impl<E: RustEmbed + Send + Sync> Endpoint for EmbeddedFileEndpoint<E> {
28    type Output = Response;
29
30    async fn call(&self, req: Request) -> Result<Self::Output, Error> {
31        if req.method() != Method::GET {
32            return Err(StatusCode::METHOD_NOT_ALLOWED.into());
33        }
34
35        match E::get(&self.path) {
36            Some(content) => {
37                let hash = hex::encode(content.metadata.sha256_hash());
38                if req
39                    .headers()
40                    .get(header::IF_NONE_MATCH)
41                    .map(|etag| etag.to_str().unwrap_or("000000").eq(&hash))
42                    .unwrap_or(false)
43                {
44                    return Err(StatusCode::NOT_MODIFIED.into());
45                }
46
47                // otherwise, return 200 with etag hash
48                let body: Vec<u8> = content.data.into();
49                let mime = mime_guess::from_path(&self.path).first_or_octet_stream();
50                Ok(Response::builder()
51                    .header(header::CONTENT_TYPE, mime.as_ref())
52                    .header(header::ETAG, hash)
53                    .body(body))
54            }
55            None => Err(StatusCode::NOT_FOUND.into()),
56        }
57    }
58}
59
60/// An endpoint that wraps a `rust-embed` bundle.
61pub struct EmbeddedFilesEndpoint<E: RustEmbed + Send + Sync> {
62    _embed: PhantomData<E>,
63}
64
65impl<E: RustEmbed + Sync + Send> Default for EmbeddedFilesEndpoint<E> {
66    #[inline]
67    fn default() -> Self {
68        Self::new()
69    }
70}
71
72impl<E: RustEmbed + Send + Sync> EmbeddedFilesEndpoint<E> {
73    /// Create a new `EmbeddedFilesEndpoint` from a `rust-embed` bundle.
74    pub fn new() -> Self {
75        EmbeddedFilesEndpoint { _embed: PhantomData }
76    }
77}
78
79impl<E: RustEmbed + Send + Sync> Endpoint for EmbeddedFilesEndpoint<E> {
80    type Output = Response;
81
82    async fn call(&self, req: Request) -> Result<Self::Output, Error> {
83        let path = req.uri().path().trim_start_matches('/');
84        let original_path = req.original_uri().path();
85        let original_end_with_slash = original_path.ends_with('/');
86
87        use header::LOCATION;
88
89        log::info!("path: {}", path);
90
91        if path.is_empty() && !original_end_with_slash {
92            Ok(Response::builder()
93                .status(StatusCode::FOUND)
94                .header(LOCATION, format!("{}/", original_path))
95                .finish())
96        } else if original_end_with_slash {
97            let path = format!("{}index.html", path);
98            EmbeddedFileEndpoint::<E>::new(&path).call(req).await
99        } else if E::get(path).is_some() {
100            EmbeddedFileEndpoint::<E>::new(path).call(req).await
101        } else if E::get(&format!("{}/index.html", path)).is_some() {
102            Ok(Response::builder()
103                .status(StatusCode::FOUND)
104                .header(LOCATION, format!("{}/", original_path))
105                .finish())
106        } else {
107            log::info!("path: {path} not found, fallback to index.html");
108            EmbeddedFileEndpoint::<E>::new(&format!("index.html")).call(req).await
109        }
110    }
111}