use std::sync::Arc;
use plecto_host::{
Header, HttpRequest, HttpResponse, LoadedFilter, RequestBodyDecision, RequestDecision,
RequestEdit, RequestTrace, ResponseDecision, ResponseEdit,
};
pub enum ChainOutcome {
Respond(HttpResponse),
Forward(HttpRequest),
}
pub enum RequestBodyOutcome {
Respond(HttpResponse),
Forward(Vec<u8>),
}
pub enum ResponseOutcome {
Forward(HttpResponse),
Respond(HttpResponse),
}
pub(crate) fn dispatch_request(
chain: &[Arc<LoadedFilter>],
mut request: HttpRequest,
trace: &RequestTrace,
) -> ChainOutcome {
for filter in chain {
match filter.on_request(&request, trace) {
Ok((RequestDecision::Continue, _logs)) => {}
Ok((RequestDecision::Modified(edit), _logs)) => apply_request_edit(&mut request, edit),
Ok((RequestDecision::ShortCircuit(response), _logs)) => {
return ChainOutcome::Respond(response);
}
Err(err) => return ChainOutcome::Respond(err.fail_closed_response()),
}
}
ChainOutcome::Forward(request)
}
pub(crate) fn dispatch_request_body(
chain: &[Arc<LoadedFilter>],
mut body: Vec<u8>,
trace: &RequestTrace,
) -> RequestBodyOutcome {
for filter in chain {
match filter.on_request_body(&body, trace) {
Ok((RequestBodyDecision::Continue(next), _logs)) => body = next,
Ok((RequestBodyDecision::ShortCircuit(response), _logs)) => {
return RequestBodyOutcome::Respond(response);
}
Err(err) => return RequestBodyOutcome::Respond(err.fail_closed_response()),
}
}
RequestBodyOutcome::Forward(body)
}
pub(crate) fn dispatch_response(
chain: &[Arc<LoadedFilter>],
request: &HttpRequest,
mut response: HttpResponse,
trace: &RequestTrace,
) -> ResponseOutcome {
for filter in chain.iter().rev() {
match filter.on_response(request, &response, trace) {
Ok((ResponseDecision::Continue, _logs)) => {}
Ok((ResponseDecision::Modified(edit), _logs)) => {
apply_response_edit(&mut response, edit)
}
Ok((ResponseDecision::Replace(replacement), _logs)) => {
return ResponseOutcome::Respond(replacement);
}
Err(err) => return ResponseOutcome::Respond(err.fail_closed_response()),
}
}
ResponseOutcome::Forward(response)
}
fn apply_request_edit(request: &mut HttpRequest, edit: RequestEdit) {
remove_headers(&mut request.headers, &edit.remove_headers);
set_headers(&mut request.headers, edit.set_headers);
}
fn apply_response_edit(response: &mut HttpResponse, edit: ResponseEdit) {
if let Some(status) = edit.set_status {
response.status = status;
}
remove_headers(&mut response.headers, &edit.remove_headers);
set_headers(&mut response.headers, edit.set_headers);
}
fn remove_headers(headers: &mut Vec<Header>, names: &[String]) {
for name in names {
headers.retain(|h| !h.name.eq_ignore_ascii_case(name));
}
}
fn set_headers(headers: &mut Vec<Header>, set: Vec<Header>) {
for header in set {
if !is_repeatable(&header.name) {
headers.retain(|existing| !existing.name.eq_ignore_ascii_case(&header.name));
}
headers.push(header);
}
}
fn is_repeatable(name: &str) -> bool {
name.eq_ignore_ascii_case("set-cookie")
}
#[cfg(test)]
mod tests {
use super::*;
fn h(name: &str, value: &str) -> Header {
Header {
name: name.to_string(),
value: value.as_bytes().to_vec(),
}
}
#[test]
fn request_edit_sets_replaces_and_removes_case_insensitively() {
let mut request = HttpRequest {
method: "GET".to_string(),
path: "/".to_string(),
authority: "a".to_string(),
scheme: "https".to_string(),
headers: vec![h("X-Keep", "1"), h("X-Drop", "old"), h("X-Replace", "old")],
};
apply_request_edit(
&mut request,
RequestEdit {
set_headers: vec![h("x-replace", "new"), h("X-Add", "1")],
remove_headers: vec!["x-drop".to_string()],
},
);
assert!(
!request
.headers
.iter()
.any(|x| x.name.eq_ignore_ascii_case("x-drop")),
"removed header gone"
);
let replaced: Vec<_> = request
.headers
.iter()
.filter(|x| x.name.eq_ignore_ascii_case("x-replace"))
.collect();
assert_eq!(replaced.len(), 1, "set replaces, not duplicates");
assert_eq!(replaced[0].value.as_slice(), b"new");
assert!(
request
.headers
.iter()
.any(|x| x.name.eq_ignore_ascii_case("x-add"))
);
assert!(request.headers.iter().any(|x| x.name == "X-Keep"));
}
#[test]
fn response_edit_sets_status_and_headers() {
let mut response = HttpResponse {
status: 200,
headers: vec![h("X-Old", "1")],
body: vec![],
};
apply_response_edit(
&mut response,
ResponseEdit {
set_status: Some(503),
set_headers: vec![h("X-New", "1")],
remove_headers: vec!["x-old".to_string()],
},
);
assert_eq!(response.status, 503);
assert!(response.headers.iter().any(|x| x.name == "X-New"));
assert!(
!response
.headers
.iter()
.any(|x| x.name.eq_ignore_ascii_case("x-old"))
);
}
#[test]
fn response_edit_appends_repeatable_set_cookie_but_replaces_others() {
let mut response = HttpResponse {
status: 200,
headers: vec![h("x-keep", "1")],
body: vec![],
};
apply_response_edit(
&mut response,
ResponseEdit {
set_status: None,
set_headers: vec![
h("set-cookie", "session=abc"),
h("set-cookie", "flag=1"),
h("x-keep", "2"),
],
remove_headers: vec![],
},
);
let cookies: Vec<&str> = response
.headers
.iter()
.filter(|x| x.name.eq_ignore_ascii_case("set-cookie"))
.map(|x| std::str::from_utf8(&x.value).expect("utf-8"))
.collect();
assert_eq!(
cookies,
vec!["session=abc", "flag=1"],
"both Set-Cookie field lines must survive (append, not collapse)"
);
let keep: Vec<&Header> = response
.headers
.iter()
.filter(|x| x.name.eq_ignore_ascii_case("x-keep"))
.collect();
assert_eq!(keep.len(), 1, "a non-repeatable header still replaces");
assert_eq!(keep[0].value.as_slice(), b"2");
}
}