Preserve Next.js and Google Tag Manager rewrites in streamed pages

crates/trusted-server-adapter-axum/tests/routes.rs

@@ -7,8 +7,10 @@
 use axum::body::Body as AxumBody;
 use axum::http::Request;
 use edgezero_adapter_axum::service::EdgeZeroAxumService;
+use std::sync::Arc;
 use tower::{Service as _, ServiceExt as _};
 use trusted_server_adapter_axum::app::TrustedServerApp;
+use trusted_server_core::test_support::nextjs_auction;
 
 const LEGACY_ADMIN_DENY_METHODS: &[&str] =
     &["GET", "POST", "HEAD", "OPTIONS", "PUT", "PATCH", "DELETE"];
@@ -939,6 +941,161 @@ async fn nextjs_auction_output_holds_until_the_structural_body_close() {
     );
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn nextjs_gtm_rejects_malformed_request_hosts() {
+    for header in ["host", "x-forwarded-host", "forwarded"] {
+        let origin = Arc::new(nextjs_auction::NextJsAuctionOrigin::with_html_response(
+            &nextjs_auction::script_composition_html(),
+            false,
+            true,
+        ));
+        let router = TrustedServerApp::routes_with_settings_and_services(
+            nextjs_auction::script_composition_settings(),
+            nextjs_auction::services(origin),
+        )
+        .expect("should build script composition router");
+        let mut service = EdgeZeroAxumService::new(router);
+        let hostile_host = "example.com</script><script>alert(1)</script>";
+        let value = if header == "forwarded" {
+            format!("host=\"{hostile_host}\"")
+        } else {
+            hostile_host.to_owned()
+        };
+        let mut request = Request::builder()
+            .uri("/script-composition")
+            .header("host", nextjs_auction::PUBLISHER_HOST)
+            .body(AxumBody::empty())
+            .expect("should build hostile host request");
+        request.headers_mut().insert(
+            header,
+            value.parse().expect("should encode hostile host header"),
+        );
+        let response = service
+            .ready()
+            .await
+            .expect("should be ready")
+            .call(request)
+            .await
+            .expect("should respond to hostile host request");
+        assert_eq!(
+            response.status().as_u16(),
+            400,
+            "should reject malformed {header} rather than rewrite script source"
+        );
+        let body = axum::body::to_bytes(response.into_body(), 1024)
+            .await
+            .expect("should consume host rejection response");
+        assert!(
+            !String::from_utf8_lossy(&body).contains(hostile_host),
+            "should not reflect the rejected host in the response"
+        );
+    }
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn nextjs_gtm_accepts_local_request_hosts_and_ports() {
+    for host in [
+        "localhost:3000",
+        "example.com:8443",
+        "[::1]:3000",
+        "my_service:8080",
+        "_service_.example.com:8443",
+    ] {
+        let origin = Arc::new(nextjs_auction::NextJsAuctionOrigin::with_html_response(
+            &nextjs_auction::script_composition_html(),
+            false,
+            true,
+        ));
+        let router = TrustedServerApp::routes_with_settings_and_services(
+            nextjs_auction::script_composition_settings(),
+            nextjs_auction::services(origin),
+        )
+        .expect("should build script composition router");
+        let mut service = EdgeZeroAxumService::new(router);
+        let request = Request::builder()
+            .uri("/script-composition")
+            .header("host", host)
+            .body(AxumBody::empty())
+            .expect("should build local host request");
+        let response = service
+            .ready()
+            .await
+            .expect("should be ready")
+            .call(request)
+            .await
+            .expect("should serve local host request");
+        assert_eq!(
+            response.status().as_u16(),
+            200,
+            "should accept host: {host}"
+        );
+        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
+            .await
+            .expect("should consume local host response");
+        let html = String::from_utf8(body.to_vec()).expect("should emit UTF-8 HTML");
+        assert!(
+            html.contains(&format!("http://{host}/app")),
+            "should preserve host and port in rewritten Flight URL: {host}"
+        );
+        assert!(
+            !html.contains("__ts_rsc_"),
+            "should resolve every local host payload"
+        );
+        assert!(
+            html.ends_with("</body></html>"),
+            "should complete local response"
+        );
+    }
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn nextjs_gtm_html_response_is_complete() {
+    for gzip in [false, true] {
+        let origin = Arc::new(nextjs_auction::NextJsAuctionOrigin::with_html_response(
+            &nextjs_auction::script_composition_html(),
+            gzip,
+            false,
+        ));
+        let router = TrustedServerApp::routes_with_settings_and_services(
+            nextjs_auction::script_composition_settings(),
+            nextjs_auction::services(origin),
+        )
+        .expect("should build mixed integration router");
+        let mut service = EdgeZeroAxumService::new(router);
+        let request = Request::builder()
+            .method("GET")
+            .uri("/article")
+            .header("host", nextjs_auction::PUBLISHER_HOST)
+            .header("accept", "text/html")
+            .body(AxumBody::empty())
+            .expect("should build navigation");
+        let response = service
+            .ready()
+            .await
+            .expect("should be ready")
+            .call(request)
+            .await
+            .expect("should serve HTML");
+        assert_eq!(
+            response.status().as_u16(),
+            200,
+            "should return a successful composed publisher response"
+        );
+        assert_eq!(
+            response
+                .headers()
+                .get("content-encoding")
+                .map(|value| value.to_str().expect("should have ASCII coding")),
+            gzip.then_some("gzip"),
+            "should retain the requested response compression"
+        );
+        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
+            .await
+            .expect("should consume entire adapter body");
+        nextjs_auction::assert_script_composition_response(&body, gzip);
+    }
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn health_reports_git_version() {
     let mut service = make_service();

crates/trusted-server-adapter-fastly/src/app.rs

@@ -1390,6 +1390,7 @@ mod tests {
     use std::net::{IpAddr, Ipv4Addr};
 
     use error_stack::Report;
+    use futures::StreamExt as _;
     use futures::executor::block_on;
     use trusted_server_core::constants::HEADER_X_GEO_INFO_AVAILABLE;
     use trusted_server_core::ec::device::DeviceSignals;
@@ -1406,6 +1407,7 @@ mod tests {
         TemplateCacheMiss, TemplateCacheReservation, TemplateEntry, TemplateMetadata,
     };
     use trusted_server_core::settings::Settings;
+    use trusted_server_core::test_support::nextjs_auction;
 
     #[test]
     fn hooks_store_metadata_matches_edgezero_manifest() {
@@ -3147,6 +3149,58 @@ mod tests {
         );
     }
 
+    #[test]
+    fn dispatch_fallback_nextjs_gtm_response_is_complete() {
+        for gzip in [false, true] {
+            let state = build_state_from_settings(nextjs_auction::script_composition_settings())
+                .expect("should build mixed state");
+            let origin = Arc::new(nextjs_auction::NextJsAuctionOrigin::with_html_response(
+                &nextjs_auction::script_composition_html(),
+                gzip,
+                true,
+            ));
+            let services = nextjs_auction::services(origin);
+            let request = request_builder()
+                .method(Method::GET)
+                .uri(format!(
+                    "https://{}/article",
+                    nextjs_auction::PUBLISHER_HOST
+                ))
+                .header("host", nextjs_auction::PUBLISHER_HOST)
+                .header("accept", "text/html")
+                .body(Body::empty())
+                .expect("should build navigation");
+            let response = block_on(super::dispatch_fallback(&state, &services, request));
+            assert_eq!(
+                response.status(),
+                StatusCode::OK,
+                "should return a successful composed publisher response"
+            );
+            assert!(
+                !response.headers().contains_key(header::CONTENT_LENGTH),
+                "should drop the stale content length after streaming rewrites"
+            );
+            assert_eq!(
+                response
+                    .headers()
+                    .get(header::CONTENT_ENCODING)
+                    .map(|value| value.to_str().expect("should have ASCII coding")),
+                gzip.then_some("gzip"),
+                "should retain the requested response compression"
+            );
+            let Body::Stream(mut stream) = response.into_body() else {
+                panic!("should retain lazy Fastly body");
+            };
+            let mut bytes = Vec::new();
+            block_on(async {
+                while let Some(chunk) = stream.next().await {
+                    bytes.extend_from_slice(&chunk.expect("should process every stream item"));
+                }
+            });
+            nextjs_auction::assert_script_composition_response(&bytes, gzip);
+        }
+    }
+
     #[test]
     fn dispatch_fallback_streams_publisher_body_without_buffering() {
         // Regression guard for the publisher streaming cutover (#849): a