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
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
/*!
# Salvo handlers for etag and last-modified-since headers.
This crate provides three handlers: [`ETag`], [`Modified`], and
[`CachingHeaders`].
Unless you are sure that you _don't_ want either etag or last-modified
behavior, please use the combined [`CachingHeaders`] handler.
 */

use etag::EntityTag;
use salvo_core::http::header::{ETAG, IF_NONE_MATCH};
use salvo_core::http::headers::{self, HeaderMapExt};
use salvo_core::http::{ResBody, StatusCode};
use salvo_core::{async_trait, Depot, FlowCtrl, Handler, Request, Response};

/**
# Etag and If-None-Match header handler

Salvo handler that provides an outbound [`etag
header`](https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/ETag)
after other handlers have been run, and if the request includes an
[`if-none-match`](https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/If-None-Match)
header, compares these values and sends a
[`304 not modified`](https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/304) status,
omitting the response body.

## Streamed bodies

Note that this handler does not currently provide an etag trailer for
streamed bodies, but may do so in the future.

## Strong vs weak comparison

Etags can be compared using a strong method or a weak
method. By default, this handler allows weak comparison. To change
this setting, construct your handler with `Etag::new().strong()`.
See [`etag::EntityTag`](https://docs.rs/etag/3.0.0/etag/struct.EntityTag.html#comparison)
for further documentation.

Read more: <https://salvo.rs>
*/
#[derive(Default, Clone, Copy, Debug)]
pub struct ETag {
    strong: bool,
}

impl ETag {
    /// constructs a new Etag handler
    pub fn new() -> Self {
        Self::default()
    }

    /// Configures this handler to use strong content-based etag comparison only. See
    /// [`etag::EntityTag`](https://docs.rs/etag/3.0.0/etag/struct.EntityTag.html#comparison)
    /// for further documentation on the differences between strong
    /// and weak etag comparison.
    pub fn strong(mut self) -> Self {
        self.strong = true;
        self
    }
}

#[async_trait]
impl Handler for ETag {
    async fn handle(&self, req: &mut Request, depot: &mut Depot, res: &mut Response, ctrl: &mut FlowCtrl) {
        ctrl.call_next(req, depot, res).await;
        if ctrl.is_ceased() {
            return;
        }

        let if_none_match = req
            .headers()
            .get(IF_NONE_MATCH)
            .and_then(|etag| etag.to_str().ok())
            .and_then(|etag| etag.parse::<EntityTag>().ok());

        let etag = req
            .headers()
            .get(ETAG)
            .and_then(|etag| etag.to_str().ok())
            .and_then(|etag| etag.parse().ok())
            .or_else(|| {
                let etag = match &res.body {
                    ResBody::Once(bytes) => Some(EntityTag::from_data(bytes)),
                    ResBody::Chunks(bytes) => {
                        let tags = bytes
                            .iter()
                            .map(|item| EntityTag::from_data(item).tag().to_owned())
                            .collect::<Vec<_>>()
                            .concat();
                        Some(EntityTag::from_data(tags.as_bytes()))
                    }
                    ResBody::Stream(_) => {
                        tracing::debug!("etag not supported for streaming body");
                        None
                    }
                    ResBody::None => {
                        tracing::debug!("etag not supported for empty body");
                        None
                    }
                    _ => None,
                };

                if let Some(etag) = &etag {
                    match etag.to_string().parse::<headers::ETag>() {
                        Ok(etag) => res.headers_mut().typed_insert(etag),
                        Err(e) => {
                            tracing::error!(error = ?e, "failed to parse etag");
                        }
                    }
                }
                etag
            });

        if let (Some(etag), Some(if_none_match)) = (etag, if_none_match) {
            let eq = if self.strong {
                etag.strong_eq(&if_none_match)
            } else {
                etag.weak_eq(&if_none_match)
            };

            if eq {
                res.body(ResBody::None);
                res.status_code(StatusCode::NOT_MODIFIED);
            }
        }
    }
}

/**
# A handler for the `Last-Modified` and `If-Modified-Since` header interaction.

This handler does not set a `Last-Modified` header on its own, but
relies on other handlers doing so.
*/
#[derive(Clone, Debug, Copy, Default)]
pub struct Modified {
    _private: (),
}

impl Modified {
    /// Constructs a new Modified handler
    pub fn new() -> Self {
        Self { _private: () }
    }
}

#[async_trait]
impl Handler for Modified {
    async fn handle(&self, req: &mut Request, depot: &mut Depot, res: &mut Response, ctrl: &mut FlowCtrl) {
        ctrl.call_next(req, depot, res).await;
        if ctrl.is_ceased() {
            return;
        }

        if let (Some(if_modified_since), Some(last_modified)) = (
            req.headers().typed_get::<headers::IfModifiedSince>(),
            res.headers().typed_get::<headers::LastModified>(),
        ) {
            if !if_modified_since.is_modified(last_modified.into()) {
                res.body(ResBody::None);
                res.status_code(StatusCode::NOT_MODIFIED);
            }
        }
    }
}

/**
A combined handler that provides both [`ETag`] and [`Modified`] behavior.
*/
#[derive(Clone, Debug, Copy, Default)]
pub struct CachingHeaders(Modified, ETag);

impl CachingHeaders {
    /// Constructs a new combination modified and etag handler
    pub fn new() -> Self {
        Self::default()
    }
}

#[async_trait]
impl Handler for CachingHeaders {
    async fn handle(&self, req: &mut Request, depot: &mut Depot, res: &mut Response, ctrl: &mut FlowCtrl) {
        self.0.handle(req, depot, res, ctrl).await;
        if res.status_code != Some(StatusCode::NOT_MODIFIED) {
            self.1.handle(req, depot, res, ctrl).await;
        }
    }
}

#[cfg(test)]
mod tests {
    use salvo_core::http::header::*;
    use salvo_core::prelude::*;
    use salvo_core::test::TestClient;

    use super::*;

    #[handler]
    async fn hello() -> &'static str {
        "Hello World"
    }

    #[tokio::test]
    async fn test_affix() {
        let router = Router::with_hoop(CachingHeaders::new()).get(hello);
        let service = Service::new(router);

        let respone = TestClient::get("http://127.0.0.1:5800/").send(&service).await;
        assert_eq!(respone.status_code, Some(StatusCode::OK));

        let etag = respone.headers().get(ETAG).unwrap();
        let respone = TestClient::get("http://127.0.0.1:5800/")
            .add_header(IF_NONE_MATCH, etag, true)
            .send(&service)
            .await;
        assert_eq!(respone.status_code, Some(StatusCode::NOT_MODIFIED));
        assert!(respone.body.is_none());
    }
}