simbld_http/helpers/
response_middleware_helper.rs

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
use actix_service::{Service, Transform};
use actix_web::http::header::{HeaderName, HeaderValue};
use actix_web::{dev::ServiceRequest, dev::ServiceResponse, Error};
use futures_util::future::{ok, LocalBoxFuture, Ready};
use std::task::{Context, Poll};
use std::time::Instant;

pub struct ResponseMiddleware;

impl<S, B> Transform<S, ServiceRequest> for ResponseMiddleware
where
  S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error>,
  S::Future: 'static,
  B: 'static,
{
  type Response = ServiceResponse<B>;
  type Error = Error;
  type Transform = ResponseMiddlewareService<S>;
  type InitError = ();
  type Future = Ready<Result<Self::Transform, Self::InitError>>;

  fn new_transform(&self, service: S) -> Self::Future {
    ok(ResponseMiddlewareService {
      service,
    })
  }
}

pub struct ResponseMiddlewareService<S> {
  service: S,
}

impl<S, B> Service<ServiceRequest> for ResponseMiddlewareService<S>
where
  S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error>,
  S::Future: 'static,
  B: 'static,
{
  type Response = ServiceResponse<B>;
  type Error = Error;
  type Future = LocalBoxFuture<'static, Result<Self::Response, Self::Error>>;

  fn poll_ready(&self, ctx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
    self.service.poll_ready(ctx)
  }

  fn call(&self, req: ServiceRequest) -> Self::Future {
    let fut = self.service.call(req);
    let start_time = Instant::now();

    Box::pin(async move {
      let mut res = fut.await?;

      // Add description
      let status_code = res.status().as_u16();
      if let Some(description) =
        crate::helpers::response_helpers::get_description_by_code(status_code)
      {
        res.headers_mut().insert(
          HeaderName::from_static("x-status-description"),
          HeaderValue::from_str(description).unwrap(),
        );
      }

      // Add processing time
      let duration = start_time.elapsed();
      res.headers_mut().insert(
        HeaderName::from_static("x-response-time-ms"),
        HeaderValue::from_str(&format!("{}", duration.as_millis())).unwrap(),
      );

      Ok(res)
    })
  }
}