1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112

//! Static files serving and embedding.

use std::{borrow::Cow, marker::PhantomData};

use rust_embed::{EmbeddedFile, RustEmbed};
use viz_core::{
    async_trait,
    header::{HeaderMap, CONTENT_TYPE, ETAG, IF_NONE_MATCH},
    types::Params,
    Handler, IntoResponse, Method, Request, Response, Result, StatusCode,
};

/// Serve a single embedded file.
#[derive(Debug)]
pub struct File<E>(Cow<'static, str>, PhantomData<E>);

impl<E> Clone for File<E> {
    fn clone(&self) -> Self {
        Self(self.0.to_owned(), PhantomData)
    }
}

impl<E> File<E> {
    /// Serve a new file by the specified path.
    pub fn new(path: &'static str) -> Self {
        Self(path.into(), PhantomData)
    }
}

#[async_trait]
impl<E> Handler<Request> for File<E>
where
    E: RustEmbed + Send + Sync + 'static,
{
    type Output = Result<Response>;

    async fn call(&self, req: Request) -> Self::Output {
        serve::<E>(&self.0, req.method(), req.headers()).await
    }
}

/// Serve a embedded directory.
#[derive(Debug)]
pub struct Dir<E>(PhantomData<E>);

impl<E> Clone for Dir<E> {
    fn clone(&self) -> Self {
        Self(PhantomData)
    }
}

impl<E> Default for Dir<E> {
    fn default() -> Self {
        Self(PhantomData)
    }
}

#[async_trait]
impl<E> Handler<Request> for Dir<E>
where
    E: RustEmbed + Send + Sync + 'static,
{
    type Output = Result<Response>;

    async fn call(&self, req: Request) -> Self::Output {
        let path = match req
            .extensions()
            .get::<Params>()
            .and_then(|params| params.first().map(|(_, v)| v))
        {
            Some(p) => p,
            None => "index.html",
        };

        serve::<E>(path, req.method(), req.headers()).await
    }
}

async fn serve<E>(path: &str, method: &Method, headers: &HeaderMap) -> Result<Response>
where
    E: RustEmbed + Send + Sync + 'static,
{
    if method != Method::GET {
        Err(StatusCode::METHOD_NOT_ALLOWED.into_error())?;
    }

    match E::get(&path) {
        Some(EmbeddedFile { data, metadata }) => {
            let hash = hex::encode(metadata.sha256_hash());

            if headers
                .get(IF_NONE_MATCH)
                .map(|etag| etag.to_str().unwrap_or("000000").eq(&hash))
                .unwrap_or(false)
            {
                Err(StatusCode::NOT_MODIFIED.into_error())?;
            }

            Response::builder()
                .header(
                    CONTENT_TYPE,
                    mime_guess::from_path(&path)
                        .first_or_octet_stream()
                        .as_ref(),
                )
                .header(ETAG, hash)
                .body(data.into())
                .map_err(Into::into)
        }
        None => Err(StatusCode::NOT_FOUND.into_error()),
    }
}