Skip to content

Commit f02edbd

Browse files
GlenDCGlenDC
authored
Patch/header removal improvements (#831)
1 parent 8fb088b commit f02edbd

File tree

7 files changed

+307
-25
lines changed

7 files changed

+307
-25
lines changed

‎examples/http_record_har.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -171,7 +171,7 @@ async fn main() -> Result<(), BoxError> {
171171
tracing::debug!(
172172
"force a stop recording so the file immediately flushes (DX etc)"
173173
);
174-
state.har_layer.recorder.stop_record().await;
174+
state.har_layer.recorder().stop_record().await;
175175
}
176176
StatusCode::OK
177177
})),

‎rama-http/src/layer/har/layer.rs‎

Lines changed: 28 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,13 +5,35 @@ use rama_core::Layer;
55
#[non_exhaustive]
66
#[derive(Debug, Clone)]
77
pub struct HARExportLayer<R, T> {
8-
pub recorder: R,
9-
pub toggle: T,
8+
recorder: R,
9+
toggle: T,
10+
11+
preserve_sensitive: bool,
1012
}
1113

1214
impl<R, T> HARExportLayer<R, T> {
1315
pub fn new(recorder: R, toggle: T) -> Self {
14-
Self { recorder, toggle }
16+
Self {
17+
recorder,
18+
toggle,
19+
preserve_sensitive: false,
20+
}
21+
}
22+
23+
pub fn recorder(&self) -> &R {
24+
&self.recorder
25+
}
26+
27+
pub fn toggle(&self) -> &T {
28+
&self.toggle
29+
}
30+
31+
rama_utils::macros::generate_set_and_with! {
32+
/// Sets whether to preserve sensitive headers (false by default).
33+
pub fn preserve_sensitive(mut self) -> Self {
34+
self.preserve_sensitive = true;
35+
self
36+
}
1537
}
1638
}
1739

@@ -27,6 +49,7 @@ where
2749
service,
2850
toggle: self.toggle.clone(),
2951
recorder: self.recorder.clone(),
52+
preserve_sensitive: self.preserve_sensitive,
3053
}
3154
}
3255

@@ -35,6 +58,7 @@ where
3558
service,
3659
toggle: self.toggle,
3760
recorder: self.recorder,
61+
preserve_sensitive: self.preserve_sensitive,
3862
}
3963
}
4064
}
@@ -80,10 +104,7 @@ mod tests {
80104

81105
impl<T: Toggle> HARExportLayer<InMemoryRecorder, T> {
82106
pub fn new_test(toggle: T) -> Self {
83-
Self {
84-
recorder: InMemoryRecorder::new(),
85-
toggle,
86-
}
107+
Self::new(InMemoryRecorder::new(), toggle)
87108
}
88109
}
89110

‎rama-http/src/layer/har/service.rs‎

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,26 @@ pub struct HARExportService<R, S, T> {
1818
pub(super) recorder: R,
1919
pub(super) service: S,
2020
pub(super) toggle: T,
21+
22+
pub(super) preserve_sensitive: bool,
23+
}
24+
25+
impl<R, S, T> HARExportService<R, S, T> {
26+
pub fn recorder(&self) -> &R {
27+
&self.recorder
28+
}
29+
30+
pub fn toggle(&self) -> &T {
31+
&self.toggle
32+
}
33+
34+
rama_utils::macros::generate_set_and_with! {
35+
/// Sets whether to preserve sensitive headers (false by default).
36+
pub fn preserve_sensitive(mut self) -> Self {
37+
self.preserve_sensitive = true;
38+
self
39+
}
40+
}
2141
}
2242

2343
impl<R, S, W, ReqBody, ResBody> Service<Request<ReqBody>> for HARExportService<R, S, W>
@@ -49,7 +69,11 @@ where
4969
.context("collect request body for HAR recording and inner svc")?
5070
.to_bytes();
5171

52-
let har_req_result = HarRequest::from_http_request_parts(&req_parts, &req_body_bytes);
72+
let har_req_result = HarRequest::from_http_request_parts(
73+
&req_parts,
74+
&req_body_bytes,
75+
self.preserve_sensitive,
76+
);
5377
let request = Request::from_parts(req_parts, Body::from(req_body_bytes));
5478

5579
match har_req_result {
@@ -88,6 +112,7 @@ where
88112
let maybe_response = match HarResponse::from_http_response_parts(
89113
&resp_parts,
90114
&resp_body_bytes,
115+
self.preserve_sensitive,
91116
) {
92117
Err(err) => {
93118
tracing::debug!(

‎rama-http/src/layer/har/spec.rs‎

Lines changed: 34 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,9 @@ use std::net::IpAddr;
55
use std::str::FromStr;
66

77
use crate::layer::har::extensions::RequestComment;
8+
use crate::layer::remove_header::{
9+
remove_sensitive_request_headers, remove_sensitive_response_headers,
10+
};
811
use crate::proto::HeaderByteLength;
912
use crate::request::Parts as ReqParts;
1013
use crate::response::Parts as RespParts;
@@ -409,7 +412,11 @@ impl TryFrom<Request> for crate::Request {
409412
}
410413

411414
impl Request {
412-
pub fn from_http_request_parts(parts: &ReqParts, payload: &[u8]) -> Result<Self, BoxError> {
415+
pub fn from_http_request_parts(
416+
parts: &ReqParts,
417+
payload: &[u8],
418+
preserve_sensitive: bool,
419+
) -> Result<Self, BoxError> {
413420
let post_data = if !payload.is_empty() {
414421
let mime_type = get_mime(&parts.headers);
415422
let params = if mime_type
@@ -463,7 +470,13 @@ impl Request {
463470

464471
let query_string = into_query_string(parts);
465472
let headers_order = parts.extensions.get().cloned().unwrap_or_default();
466-
let header_map = Http1HeaderMap::from_parts(parts.headers.clone(), headers_order);
473+
474+
let mut headers = parts.headers.clone();
475+
if !preserve_sensitive {
476+
remove_sensitive_request_headers(&mut headers);
477+
}
478+
479+
let header_map = Http1HeaderMap::from_parts(headers, headers_order);
467480

468481
let headers_size_ext = parts.extensions.get::<HeaderByteLength>();
469482
let headers_size = headers_size_ext.map(|v| v.0 as i64).unwrap_or(-1);
@@ -568,7 +581,11 @@ impl TryFrom<Response> for crate::Response {
568581
}
569582

570583
impl Response {
571-
pub fn from_http_response_parts(parts: &RespParts, payload: &[u8]) -> Result<Self, BoxError> {
584+
pub fn from_http_response_parts(
585+
parts: &RespParts,
586+
payload: &[u8],
587+
preserve_sensitive: bool,
588+
) -> Result<Self, BoxError> {
572589
let content = Content {
573590
size: payload.len() as i64,
574591
compression: None,
@@ -606,7 +623,13 @@ impl Response {
606623
.unwrap_or_default();
607624

608625
let headers_order = parts.extensions.get().cloned().unwrap_or_default();
609-
let header_map = Http1HeaderMap::from_parts(parts.headers.clone(), headers_order);
626+
627+
let mut headers = parts.headers.clone();
628+
if !preserve_sensitive {
629+
remove_sensitive_response_headers(&mut headers);
630+
}
631+
632+
let header_map = Http1HeaderMap::from_parts(headers, headers_order);
610633

611634
let headers_size_ext = parts.extensions.get::<HeaderByteLength>();
612635
let headers_size = headers_size_ext.map(|v| v.0 as i64).unwrap_or(-1);
@@ -882,7 +905,7 @@ mod tests {
882905
assert_eq!("www.igvita.com", host);
883906

884907
// rama Request to HAR Request
885-
let req0_back = Request::from_http_request_parts(&req0_parts, &[]).unwrap();
908+
let req0_back = Request::from_http_request_parts(&req0_parts, &[], false).unwrap();
886909
assert_eq!(entry0.request.method, req0_back.method);
887910
assert_eq!(entry0.request.url, req0_back.url);
888911
assert!(matches!(req0_back.http_version, HttpVersion::Http11));
@@ -901,7 +924,7 @@ mod tests {
901924
let (req5_parts, req5_body) = req5.into_parts();
902925
drop(req5_body);
903926

904-
let req5_back = Request::from_http_request_parts(&req5_parts, &[]).unwrap();
927+
let req5_back = Request::from_http_request_parts(&req5_parts, &[], false).unwrap();
905928
assert_eq!(5, req5_back.query_string.len(), "req: {req5_back:?}");
906929
assert!(
907930
req5_back
@@ -952,7 +975,7 @@ mod tests {
952975
assert_eq!("gzip", ce);
953976

954977
// rama Response to HAR Response
955-
let res0_back = Response::from_http_response_parts(&res0_parts, &[]).unwrap();
978+
let res0_back = Response::from_http_response_parts(&res0_parts, &[], false).unwrap();
956979
assert_eq!(200, res0_back.status);
957980
assert_eq!(Some("OK"), res0_back.status_text.as_deref());
958981
assert!(matches!(res0_back.http_version, HttpVersion::Http11));
@@ -981,7 +1004,8 @@ mod tests {
9811004
assert_eq!(req_payload, req_bytes);
9821005

9831006
// rama request parts + payload -> HAR request payload
984-
let har_req_back = Request::from_http_request_parts(&req_parts, &req_payload).unwrap();
1007+
let har_req_back =
1008+
Request::from_http_request_parts(&req_parts, &req_payload, false).unwrap();
9851009
let post_data = har_req_back.post_data.unwrap();
9861010
assert_eq!(
9871011
Some(Mime::from_str("application/octet-stream").unwrap()),
@@ -1000,7 +1024,8 @@ mod tests {
10001024
assert_eq!(res_payload, res_bytes);
10011025

10021026
// rama response parts + payload -> HAR response payload
1003-
let har_res_back = Response::from_http_response_parts(&res_parts, &res_payload).unwrap();
1027+
let har_res_back =
1028+
Response::from_http_response_parts(&res_parts, &res_payload, false).unwrap();
10041029
assert_eq!(res_payload.len() as i64, har_res_back.content.size);
10051030
assert_eq!(
10061031
Some(Mime::from_str("application/octet-stream").unwrap()),

0 commit comments

Comments
 (0)