rama_http/layer/har/
service.rs

1use crate::layer::har::recorder::{
2    HttpRequestCapture, HttpResponseCapture, RecorderSession, StreamingRecorder, WebSocketCapture,
3    body_capture_channel,
4};
5use crate::layer::har::spec::{Request as HarRequest, Response as HarResponse};
6use crate::layer::har::toggle::Toggle;
7use crate::{Body, Request, Response, StreamingBody};
8
9use jiff::Timestamp;
10use rama_core::error::BoxError;
11use rama_core::extensions::ExtensionsRef;
12use rama_core::telemetry::tracing;
13use rama_core::{Service, bytes::Bytes};
14use rama_http_types::proto::h2::ext::Protocol;
15use tokio::time::Instant;
16
17#[derive(Clone)]
18pub struct HARExportService<R, S, T> {
19    pub(super) recorder: R,
20    pub(super) service: S,
21    pub(super) toggle: T,
22
23    pub(super) preserve_sensitive: bool,
24}
25
26struct ActiveRecording<S> {
27    session: S,
28    web_socket_capture: Option<WebSocketCapture>,
29}
30
31impl<R, S, T> HARExportService<R, S, T> {
32    pub fn recorder(&self) -> &R {
33        &self.recorder
34    }
35
36    pub fn toggle(&self) -> &T {
37        &self.toggle
38    }
39
40    rama_utils::macros::generate_set_and_with! {
41        /// Sets whether to preserve sensitive headers (false by default).
42        pub fn preserve_sensitive(mut self) -> Self {
43            self.preserve_sensitive = true;
44            self
45        }
46    }
47}
48
49impl<R, S, W, ReqBody, ResBody> Service<Request<ReqBody>> for HARExportService<R, S, W>
50where
51    R: StreamingRecorder,
52    S: Service<Request, Output = Response<ResBody>, Error: Into<BoxError> + Send + Sync + 'static>,
53    W: Toggle,
54    ReqBody: StreamingBody<Data = Bytes, Error: Into<BoxError>> + Send + Sync + 'static,
55    ResBody: StreamingBody<Data = Bytes, Error: Into<BoxError>> + Send + Sync + 'static,
56{
57    type Output = Response;
58    type Error = BoxError;
59
60    async fn serve(&self, req: Request<ReqBody>) -> Result<Self::Output, Self::Error> {
61        let (request, recording) = if self.toggle.status().await {
62            self.start_recording(req).await
63        } else {
64            self.recorder.stop_record().await;
65            (req.map(Body::new), None)
66        };
67
68        let result = self.service.serve(request).await;
69        let Some(recording) = recording else {
70            return result
71                .map(|response| response.map(Body::new))
72                .map_err(Into::into);
73        };
74
75        match result {
76            Ok(response) => self.record_response(recording, response).await,
77            Err(err) => {
78                if let Some(capture) = recording.web_socket_capture {
79                    capture.close();
80                }
81                _ = recording.session.record_request_only().await;
82                Err(err.into())
83            }
84        }
85    }
86}
87
88impl<R, S, W> HARExportService<R, S, W>
89where
90    R: StreamingRecorder,
91{
92    async fn start_recording<ReqBody>(
93        &self,
94        request: Request<ReqBody>,
95    ) -> (Request, Option<ActiveRecording<R::Session>>)
96    where
97        ReqBody: StreamingBody<Data = Bytes, Error: Into<BoxError>> + Send + Sync + 'static,
98    {
99        let start_time = Timestamp::now();
100        let begin = Instant::now();
101        let (parts, body) = request.into_parts();
102        let har_request =
103            match HarRequest::from_http_request_parts(&parts, &[], self.preserve_sensitive) {
104                Ok(request) => request,
105                Err(err) => {
106                    tracing::debug!(
107                        "failed to create HAR request from incoming HTTP Request: {err}"
108                    );
109                    return (Request::from_parts(parts, Body::new(body)), None);
110                }
111            };
112
113        let mime_type = super::spec::get_mime(&parts.headers);
114        let web_socket = is_web_socket_request(&parts);
115        let (sink, body_stream) = body_capture_channel();
116        let capture = HttpRequestCapture {
117            started_date_time: start_time,
118            begin,
119            request: har_request,
120            body_mime_type: mime_type,
121            body: body_stream,
122            web_socket,
123        };
124
125        match self.recorder.start_http_recording(capture).await {
126            Some(session) => {
127                let web_socket_capture = if web_socket {
128                    session.web_socket_capture()
129                } else {
130                    None
131                };
132                if let Some(capture) = &web_socket_capture {
133                    parts.extensions.insert(capture.clone());
134                }
135                (
136                    Request::from_parts(parts, Body::new(body).capture(sink)),
137                    Some(ActiveRecording {
138                        session,
139                        web_socket_capture,
140                    }),
141                )
142            }
143            None => (Request::from_parts(parts, Body::new(body)), None),
144        }
145    }
146
147    async fn record_response<ResBody>(
148        &self,
149        recording: ActiveRecording<R::Session>,
150        response: Response<ResBody>,
151    ) -> Result<Response, BoxError>
152    where
153        ResBody: StreamingBody<Data = Bytes, Error: Into<BoxError>> + Send + Sync + 'static,
154    {
155        let (parts, body) = response.into_parts();
156        let har_response =
157            match HarResponse::from_http_response_parts(&parts, &[], self.preserve_sensitive) {
158                Ok(response) => response,
159                Err(err) => {
160                    tracing::debug!(
161                        "failed to create HAR response from returned HTTP Response: {err}"
162                    );
163                    if let Some(capture) = recording.web_socket_capture {
164                        capture.close();
165                    }
166                    let extensions = recording.session.record_request_only().await;
167                    let response = Response::from_parts(parts, Body::new(body));
168                    extend_response(&response, extensions);
169                    return Ok(response);
170                }
171            };
172
173        let (sink, body_stream) = body_capture_channel();
174        if let Some(capture) = recording.web_socket_capture {
175            if is_successful_web_socket_response(&parts) {
176                parts.extensions.insert(capture);
177            } else {
178                capture.close();
179            }
180        }
181        let extensions = recording
182            .session
183            .record_response(HttpResponseCapture {
184                response: har_response,
185                body: body_stream,
186            })
187            .await;
188        let response = Response::from_parts(parts, Body::new(body).capture(sink));
189        extend_response(&response, extensions);
190        Ok(response)
191    }
192}
193
194fn is_web_socket_request(parts: &rama_http_types::request::Parts) -> bool {
195    parts
196        .extensions
197        .get_ref::<Protocol>()
198        .is_some_and(|protocol| protocol.as_str().eq_ignore_ascii_case("websocket"))
199        || parts
200            .headers
201            .get_all(rama_http_types::header::UPGRADE)
202            .iter()
203            .filter_map(|value| value.to_str().ok())
204            .flat_map(|value| value.split(','))
205            .any(|protocol| protocol.trim().eq_ignore_ascii_case("websocket"))
206}
207
208fn is_successful_web_socket_response(parts: &rama_http_types::response::Parts) -> bool {
209    match parts.version {
210        rama_http_types::Version::HTTP_2 | rama_http_types::Version::HTTP_3 => {
211            parts.status.is_success()
212        }
213        _ => parts.status == rama_http_types::StatusCode::SWITCHING_PROTOCOLS,
214    }
215}
216
217fn extend_response(response: &Response, extensions: Option<rama_core::extensions::Extensions>) {
218    if let Some(extensions) = extensions {
219        tracing::trace!("extend response with HAR recorder extensions");
220        response.extensions().extend(&extensions);
221    }
222}
223
224#[cfg(test)]
225mod tests {
226    use super::{is_successful_web_socket_response, is_web_socket_request};
227    use crate::{Request, Response, StatusCode, Version, proto::h2::ext::Protocol};
228    use rama_core::extensions::ExtensionsRef;
229
230    fn response_parts(status: StatusCode, version: Version) -> rama_http_types::response::Parts {
231        let (mut parts, _) = Response::new(()).into_parts();
232        parts.status = status;
233        parts.version = version;
234        parts
235    }
236
237    #[test]
238    fn validates_web_socket_response_by_http_version() {
239        assert!(is_successful_web_socket_response(&response_parts(
240            StatusCode::SWITCHING_PROTOCOLS,
241            Version::HTTP_11,
242        )));
243        assert!(!is_successful_web_socket_response(&response_parts(
244            StatusCode::OK,
245            Version::HTTP_11,
246        )));
247        assert!(is_successful_web_socket_response(&response_parts(
248            StatusCode::OK,
249            Version::HTTP_2,
250        )));
251        assert!(is_successful_web_socket_response(&response_parts(
252            StatusCode::CREATED,
253            Version::HTTP_2,
254        )));
255        assert!(!is_successful_web_socket_response(&response_parts(
256            StatusCode::BAD_REQUEST,
257            Version::HTTP_2,
258        )));
259    }
260
261    #[test]
262    fn detects_h1_upgrade_and_h2_extended_connect_requests() {
263        let (h1, _) = Request::builder()
264            .header("upgrade", "h2c, WebSocket")
265            .body(())
266            .unwrap()
267            .into_parts();
268        assert!(is_web_socket_request(&h1));
269
270        let mut h2 = Request::new(());
271        *h2.version_mut() = Version::HTTP_2;
272        h2.extensions().insert(Protocol::from_static("websocket"));
273        let (h2, _) = h2.into_parts();
274        assert!(is_web_socket_request(&h2));
275
276        let (plain, _) = Request::new(()).into_parts();
277        assert!(!is_web_socket_request(&plain));
278    }
279}
rama_http/layer/har/service.rs

rama_http/layer/har/
service.rs
