Skip to content

Commit

Permalink
Fix and always set Content-Length header.
Browse files Browse the repository at this point in the history
Previously, the header was erroneously set as 'Content-Type' and not
visible to local clients. This commit fixes both of these issues.
  • Loading branch information
bash authored and SergioBenitez committed Mar 21, 2024
1 parent 8f3061b commit 35a1cf1
Show file tree
Hide file tree
Showing 3 changed files with 26 additions and 14 deletions.
8 changes: 8 additions & 0 deletions core/lib/src/lifecycle.rs
Original file line number Diff line number Diff line change
Expand Up @@ -145,6 +145,14 @@ impl Rocket<Orbit> {
response.strip_body();
}

if let Some(size) = response.body_mut().size().await {
response.set_raw_header("Content-Length", size.to_string());
}

if let Some(alt_svc) = request.rocket().alt_svc() {
response.set_raw_header("Alt-Svc", alt_svc);
}

// TODO: Should upgrades be handled here? We miss them on local clients.
response
}
Expand Down
17 changes: 3 additions & 14 deletions core/lib/src/server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,6 @@ impl Rocket<Orbit> {
upgrade: Option<hyper::upgrade::OnUpgrade>,
connection: ConnectionMeta,
) -> Result<hyper::Response<ReaderStream<ErasedResponse>>, http::Error> {
let alt_svc = self.alt_svc();
let request = ErasedRequest::new(self, parts, |rocket, parts| {
Request::from_hyp(rocket, parts, connection).unwrap_or_else(|e| e)
});
Expand All @@ -41,12 +40,11 @@ impl Rocket<Orbit> {
return rocket.dispatch_error(Status::BadRequest, request).await;
}

let mut response = rocket.dispatch(token, request, data).await;
response.body_mut().size().await;
response
rocket.dispatch(token, request, data).await
})
).await;

// TODO: Should upgrades be handled in dispatch?
let io_handler = response.make_io_handler(Rocket::extract_io_handler);
if let (Some(handler), Some(upgrade)) = (io_handler, upgrade) {
let upgrade = upgrade.map_ok(IoStream::from).map_err(io::Error::other);
Expand All @@ -59,20 +57,11 @@ impl Rocket<Orbit> {
builder = builder.header(header.name().as_str(), header.value());
}

if let Some(size) = response.inner().body().preset_size() {
builder = builder.header(http::header::CONTENT_TYPE, size);
}

if let Some(alt_svc) = alt_svc {
let value = http::HeaderValue::from_static(alt_svc);
builder = builder.header(http::header::ALT_SVC, value);
}

let chunk_size = response.inner().body().max_chunk_size();
builder.body(ReaderStream::with_capacity(response, chunk_size))
}

fn alt_svc(&self) -> Option<&'static str> {
pub(crate) fn alt_svc(&self) -> Option<&'static str> {
cfg!(feature = "http3-preview").then(|| {
static ALT_SVC: state::InitCell<Option<String>> = state::InitCell::new();

Expand Down
15 changes: 15 additions & 0 deletions core/lib/tests/content-length.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
#[macro_use]
extern crate rocket;

#[get("/")]
fn index() -> String {
"Hello, world!".into()
}

#[test]
fn content_length_header() {
let rocket = rocket::build().mount("/", routes![index]);
let client = rocket::local::blocking::Client::debug(rocket).unwrap();
let response = client.get("/").dispatch();
assert!(response.headers().get_one("Content-Length").is_some());
}

0 comments on commit 35a1cf1

Please sign in to comment.