From 35a1cf12b6b2f2cf4c699153f4f7938ca289bd7d Mon Sep 17 00:00:00 2001
From: Tau <git@tau.garden>
Date: Thu, 21 Mar 2024 21:36:20 +0100
Subject: [PATCH] Fix and always set Content-Length header.

Previously, the header was erroneously set as 'Content-Type' and not
visible to local clients. This commit fixes both of these issues.
---
 core/lib/src/lifecycle.rs        |  8 ++++++++
 core/lib/src/server.rs           | 17 +++--------------
 core/lib/tests/content-length.rs | 15 +++++++++++++++
 3 files changed, 26 insertions(+), 14 deletions(-)
 create mode 100644 core/lib/tests/content-length.rs

diff --git a/core/lib/src/lifecycle.rs b/core/lib/src/lifecycle.rs
index 9859f1f643..6300f15e48 100644
--- a/core/lib/src/lifecycle.rs
+++ b/core/lib/src/lifecycle.rs
@@ -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
     }
diff --git a/core/lib/src/server.rs b/core/lib/src/server.rs
index ff31791734..8d26604f33 100644
--- a/core/lib/src/server.rs
+++ b/core/lib/src/server.rs
@@ -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)
         });
@@ -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);
@@ -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();
 
diff --git a/core/lib/tests/content-length.rs b/core/lib/tests/content-length.rs
new file mode 100644
index 0000000000..c9e40e44e0
--- /dev/null
+++ b/core/lib/tests/content-length.rs
@@ -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());
+}