poem_middleware/
no_cache.rs

1
2use poem::{middleware::Middleware, http::header::{CACHE_CONTROL, EXPIRES, PRAGMA},
3http::HeaderValue, Endpoint, IntoResponse, Request, Response, Result};
4
5
6#[derive(Default)]
7#[allow(clippy::type_complexity)]
8pub struct NoCacheMiddleware;
9
10
11impl<E: Endpoint> Middleware<E> for NoCacheMiddleware {
12    type Output = NoCacheEndpoint<E>;
13
14    fn transform(&self, ep: E) -> Self::Output {
15        NoCacheEndpoint { ep }
16    }
17}
18
19pub struct NoCacheEndpoint<E> {
20    ep: E,
21}
22
23impl<E: Endpoint> Endpoint for NoCacheEndpoint<E> {
24    type Output = Response;
25
26    async fn call(&self, req: Request) -> Result<Self::Output> {
27        let mut response = self.ep.call(req).await?.into_response();
28
29        // setting Cache-Control, avoid cache store
30        response.headers_mut().insert(
31            CACHE_CONTROL,
32            HeaderValue::from_static("no-store, no-cache, must-revalidate, max-age=0"),
33        );
34
35        // setting EXPIRES
36        response
37            .headers_mut()
38            .insert(EXPIRES, HeaderValue::from_static("0"));
39
40        // setting PRAGMA, compatible for HTTP/1.0
41        response
42            .headers_mut()
43            .insert(PRAGMA, HeaderValue::from_static("no-cache"));
44
45        Ok(response)
46    }
47}
48
49
50#[cfg(test)]
51mod tests {
52    use poem::{endpoint::make_sync, http::header::{CACHE_CONTROL, EXPIRES, PRAGMA}, test::TestClient, EndpointExt as _};
53
54    use crate::no_cache::NoCacheMiddleware;
55
56    #[tokio::test]
57    async fn test_rm_cache() {
58        let ep = make_sync(|_| "hello").with(NoCacheMiddleware);
59        let cli = TestClient::new(ep);
60        let resp = cli.get("/").send().await;
61        resp.assert_status_is_ok();
62        resp.assert_header(CACHE_CONTROL, "no-store, no-cache, must-revalidate, max-age=0");
63        resp.assert_header(EXPIRES, 0);
64        resp.assert_header(PRAGMA, "no-cache")
65    }   
66}