Fix JS asset proxy review findings

Author: ChristianPavilonis

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

@@ -100,6 +100,10 @@ enum HandlerOutcome {
         response: HttpResponse,
         body: EdgeBody,
     },
+    RawStreaming {
+        response: HttpResponse,
+        body: EdgeBody,
+    },
 }
 
 impl HandlerOutcome {
@@ -108,11 +112,23 @@ impl HandlerOutcome {
         match self {
             HandlerOutcome::Buffered(resp) | HandlerOutcome::AuthChallenge(resp) => resp.status(),
             HandlerOutcome::Streaming { response, .. }
-            | HandlerOutcome::AssetStreaming { response, .. } => response.status(),
+            | HandlerOutcome::AssetStreaming { response, .. }
+            | HandlerOutcome::RawStreaming { response, .. } => response.status(),
         }
     }
 }
 
+fn response_to_legacy_outcome(response: HttpResponse) -> HandlerOutcome {
+    let (parts, body) = response.into_parts();
+    match body {
+        EdgeBody::Stream(_) => HandlerOutcome::RawStreaming {
+            response: HttpResponse::from_parts(parts, EdgeBody::empty()),
+            body,
+        },
+        body => HandlerOutcome::Buffered(HttpResponse::from_parts(parts, body)),
+    }
+}
+
 /// Returns `true` if the raw config-store value represents an enabled flag.
 ///
 /// Accepted values (after whitespace trimming): `"1"` or `"true"` in any ASCII case.
@@ -912,6 +928,46 @@ fn legacy_main(mut req: FastlyRequest) {
                 log::error!("failed to finish asset streaming body: {e}");
             }
         }
+        HandlerOutcome::RawStreaming { mut response, body } => {
+            finalize_response(&state.settings, geo_info.as_ref(), &mut response);
+            asset_cache_policy.apply_after_route_finalization(&mut response);
+            if should_finalize_ec {
+                ec_finalize_response(
+                    &state.settings,
+                    &ec_context,
+                    finalize_kv_graph.as_ref(),
+                    &partner_registry,
+                    eids_cookie.as_deref(),
+                    sharedid_cookie.as_deref(),
+                    &mut response,
+                );
+            }
+            request_filter_effects.apply_to_response(&mut response);
+            let fastly_resp = compat::to_fastly_response_skeleton(response);
+            let mut streaming_body = fastly_resp.stream_to_client();
+            let mut stream_succeeded = false;
+            if let Err(e) =
+                futures::executor::block_on(stream_asset_body(body, &mut streaming_body))
+            {
+                log::error!("raw response streaming failed: {e:?}");
+                drop(streaming_body);
+            } else if let Err(e) = streaming_body.finish() {
+                log::error!("failed to finish raw streaming body: {e}");
+            } else {
+                stream_succeeded = true;
+            }
+
+            if is_real_browser && stream_succeeded {
+                if let Some(context) = build_pull_sync_context(&ec_context) {
+                    run_pull_sync_after_send(
+                        &state.settings,
+                        &partner_registry,
+                        &context,
+                        &runtime_services,
+                    );
+                }
+            }
+        }
     }
 }
 
@@ -1359,7 +1415,7 @@ async fn route_request(
     let _ = organic_route;
 
     let outcome = result
-        .map(HandlerOutcome::Buffered)
+        .map(response_to_legacy_outcome)
         .unwrap_or_else(|e| HandlerOutcome::Buffered(http_error_response(&e)));
 
     Ok(RouteResult {

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

@@ -1551,6 +1551,84 @@ fn asset_routes_stream_asset_responses_directly() {
     // take the `response: None` direct-send path with EC finalization skipped.
 }
 
+#[test]
+fn integration_routes_preserve_streaming_proxy_responses() {
+    let mut settings = create_test_settings();
+    settings
+        .integrations
+        .insert_config(
+            "js_asset_proxy",
+            &json!({
+                "enabled": true,
+                "assets": [{
+                    "path": "/assets/vendor.js",
+                    "origin_url": "https://cdn.example.com/vendor.js",
+                    "proxy": "enabled"
+                }]
+            }),
+        )
+        .expect("should insert JS asset proxy config");
+    let orchestrator = build_orchestrator(&settings).expect("should build auction orchestrator");
+    let integration_registry =
+        IntegrationRegistry::new(&settings).expect("should create integration registry");
+    let partner_registry = test_partner_registry(&settings);
+
+    let fastly_req = Request::get("https://test.com/assets/vendor.js");
+    let http_client = Arc::new(StreamingRecordingHttpClient::new());
+    let services = test_runtime_services_with_secret_http_client_and_geo(
+        &fastly_req,
+        Arc::new(FixedBackend),
+        Arc::new(NoopSecretStore),
+        Arc::clone(&http_client) as Arc<dyn PlatformHttpClient>,
+        Arc::new(FixedGeo(us_california_geo())),
+    );
+    let req = compat::from_fastly_request(fastly_req);
+
+    let outcome = futures::executor::block_on(route_request(
+        &settings,
+        &orchestrator,
+        &integration_registry,
+        &partner_registry,
+        &services,
+        req,
+        browser_device_signals(),
+    ))
+    .expect("should route streaming integration request");
+
+    let (response, body) = match outcome.outcome {
+        HandlerOutcome::RawStreaming { response, body } => Some((response, body)),
+        _ => None,
+    }
+    .expect("should preserve streaming integration response");
+    assert_eq!(response.status(), StatusCode::OK);
+    assert!(
+        matches!(body, EdgeBody::Stream(_)),
+        "should preserve the integration response stream"
+    );
+
+    let calls = http_client
+        .calls
+        .lock()
+        .expect("should lock recorded calls");
+    assert_eq!(
+        calls.len(),
+        1,
+        "should send exactly one integration request"
+    );
+    assert!(
+        calls[0].stream_response,
+        "JS asset proxy should request a streaming origin response from the platform"
+    );
+    assert_eq!(
+        calls[0].backend_name, "https-cdn.example.com",
+        "should resolve the configured JS asset origin backend"
+    );
+    assert_eq!(
+        calls[0].uri, "https://cdn.example.com/vendor.js",
+        "should send the request to the configured JS asset origin URL"
+    );
+}
+
 #[test]
 fn asset_origin_failure_does_not_fall_back_to_publisher_origin() {
     let mut settings = create_test_settings();

crates/trusted-server-cli/src/audit.rs

@@ -426,7 +426,7 @@ fn build_js_asset_proxy_section(
         "# Audit note: some discovered scripts may be runtime-injected and may not appear\n",
     );
     toml.push_str(
-        "# in origin HTML. JS Asset Proxy rewrites only exact script src values present in\n",
+        "# in origin HTML. JS Asset Proxy rewrites only matching script src URLs present in\n",
     );
     toml.push_str("# HTML processed by Trusted Server.\n");