Add asset backend proxy pipeline through authenticated s3 and fastly io (#668)

Author: ChristianPavilonis

Cargo.lock

@@ -9,6 +9,7 @@ workspace = true
 [dependencies]
 async-trait = { workspace = true }
 base64 = { workspace = true }
+bytes = { workspace = true }
 chrono = { workspace = true }
 edgezero-adapter-fastly = { workspace = true, features = ["fastly"] }
 edgezero-core = { workspace = true }

crates/integration-tests/Cargo.lock

@@ -27,8 +27,9 @@ use trusted_server_core::geo::GeoInfo;
 use trusted_server_core::integrations::{IntegrationRegistry, ProxyDispatchInput};
 use trusted_server_core::platform::RuntimeServices;
 use trusted_server_core::proxy::{
-    handle_first_party_click, handle_first_party_proxy, handle_first_party_proxy_rebuild,
-    handle_first_party_proxy_sign,
+    handle_asset_proxy_request, handle_first_party_click, handle_first_party_proxy,
+    handle_first_party_proxy_rebuild, handle_first_party_proxy_sign, stream_asset_body,
+    AssetProxyCachePolicy,
 };
 use trusted_server_core::publisher::{
     handle_publisher_request, handle_tsjs_dynamic, stream_publisher_body, PublisherResponse,
@@ -325,6 +326,9 @@ async fn route_request(
     let path = req.get_path().to_string();
     let method = req.get_method().clone();
 
+    let mut asset_cache_policy = AssetProxyCachePolicy::OriginControlled;
+    let mut should_finalize_ec = true;
+
     // Match known routes and handle them
     let (result, organic_route) = match (method, path.as_str()) {
         // Serve the tsjs library
@@ -387,18 +391,22 @@ async fn route_request(
         }
 
         // tsjs endpoints
-        (Method::GET, "/first-party/proxy") => {
-            (handle_first_party_proxy(settings, req).await, false)
-        }
-        (Method::GET, "/first-party/click") => {
-            (handle_first_party_click(settings, req).await, false)
-        }
-        (Method::GET, "/first-party/sign") | (Method::POST, "/first-party/sign") => {
-            (handle_first_party_proxy_sign(settings, req).await, false)
-        }
-        (Method::POST, "/first-party/proxy-rebuild") => {
-            (handle_first_party_proxy_rebuild(settings, req).await, false)
-        }
+        (Method::GET, "/first-party/proxy") => (
+            handle_first_party_proxy(settings, runtime_services, req).await,
+            false,
+        ),
+        (Method::GET, "/first-party/click") => (
+            handle_first_party_click(settings, runtime_services, req).await,
+            false,
+        ),
+        (Method::GET, "/first-party/sign") | (Method::POST, "/first-party/sign") => (
+            handle_first_party_proxy_sign(settings, runtime_services, req).await,
+            false,
+        ),
+        (Method::POST, "/first-party/proxy-rebuild") => (
+            handle_first_party_proxy_rebuild(settings, runtime_services, req).await,
+            false,
+        ),
         (m, path) if integration_registry.has_route(&m, path) => {
             let result = integration_registry
                 .handle_proxy(ProxyDispatchInput {
@@ -407,6 +415,7 @@ async fn route_request(
                     settings,
                     kv: kv_graph.as_ref(),
                     ec_context: &mut ec_context,
+                    services: runtime_services,
                     req,
                 })
                 .await
@@ -418,78 +427,121 @@ async fn route_request(
             (result, true)
         }
 
-        // No known route matched, proxy to publisher origin as fallback
-        _ => {
-            log::info!(
-                "No known route matched for path: {}, proxying to publisher origin",
-                path
-            );
+        // No known route matched, proxy to an asset origin or publisher origin as fallback
+        (method, _) => {
+            let matched_asset_route = matches!(method, Method::GET | Method::HEAD)
+                .then(|| settings.asset_route_for_path(&path))
+                .flatten();
+
+            if let Some(asset_route) = matched_asset_route {
+                should_finalize_ec = false;
+                log::info!("No explicit route matched; proxying via configured asset route");
+                let result =
+                    match handle_asset_proxy_request(settings, runtime_services, req, asset_route)
+                        .await
+                    {
+                        Ok(asset_response) => {
+                            asset_cache_policy = asset_response.cache_policy();
+                            let (mut response, stream_body) =
+                                asset_response.into_response_and_body();
+                            if let Some(body) = stream_body {
+                                finalize_response(settings, geo_info.as_ref(), &mut response);
+                                asset_cache_policy.apply_after_route_finalization(&mut response);
+
+                                let mut streaming_body = response.stream_to_client();
+                                if let Err(err) = stream_asset_body(body, &mut streaming_body).await
+                                {
+                                    log::error!("Asset streaming failed: {err:?}");
+                                    drop(streaming_body);
+                                } else if let Err(err) = streaming_body.finish() {
+                                    log::error!("Failed to finish asset streaming body: {err}");
+                                }
+
+                                return Ok(RouteOutcome {
+                                    response: None,
+                                    pull_sync_context: None,
+                                });
+                            }
+                            Ok(response)
+                        }
+                        Err(e) => {
+                            asset_cache_policy = AssetProxyCachePolicy::NoStorePrivate;
+                            Err(e)
+                        }
+                    };
+                (result, false)
+            } else {
+                log::info!(
+                    "No known route matched for path: {}, proxying to publisher origin",
+                    path
+                );
 
-            match handle_publisher_request(
-                settings,
-                integration_registry,
-                runtime_services,
-                kv_graph.as_ref(),
-                &mut ec_context,
-                req,
-            ) {
-                Ok(PublisherResponse::Stream {
-                    mut response,
-                    body,
-                    params,
-                }) => {
-                    // Publisher fallback has multiple delivery modes.
-                    // EC finalization is header-only, so it must happen before
-                    // headers are committed on the streaming path.
-                    ec_finalize_response(
-                        settings,
-                        &ec_context,
-                        finalize_kv_graph.as_ref(),
-                        partner_registry,
-                        eids_cookie.as_deref(),
-                        sharedid_cookie.as_deref(),
-                        &mut response,
-                    );
-                    finalize_response(settings, geo_info.as_ref(), &mut response);
-
-                    let mut streaming_body = response.stream_to_client();
-                    let mut stream_succeeded = false;
-                    if let Err(err) = stream_publisher_body(
+                match handle_publisher_request(
+                    settings,
+                    integration_registry,
+                    runtime_services,
+                    kv_graph.as_ref(),
+                    &mut ec_context,
+                    req,
+                ) {
+                    Ok(PublisherResponse::Stream {
+                        mut response,
                         body,
-                        &mut streaming_body,
-                        &params,
-                        settings,
-                        integration_registry,
-                    ) {
-                        // Headers are already committed. Log and abort rather
-                        // than trying to replace the response mid-stream.
-                        log::error!("Streaming processing failed: {err:?}");
-                        drop(streaming_body);
-                    } else if let Err(err) = streaming_body.finish() {
-                        log::error!("Failed to finish streaming body: {err}");
-                    } else {
-                        stream_succeeded = true;
+                        params,
+                    }) => {
+                        // Publisher fallback has multiple delivery modes.
+                        // EC finalization is header-only, so it must happen before
+                        // headers are committed on the streaming path.
+                        ec_finalize_response(
+                            settings,
+                            &ec_context,
+                            finalize_kv_graph.as_ref(),
+                            partner_registry,
+                            eids_cookie.as_deref(),
+                            sharedid_cookie.as_deref(),
+                            &mut response,
+                        );
+                        finalize_response(settings, geo_info.as_ref(), &mut response);
+
+                        let mut streaming_body = response.stream_to_client();
+                        let mut stream_succeeded = false;
+                        if let Err(err) = stream_publisher_body(
+                            body,
+                            &mut streaming_body,
+                            &params,
+                            settings,
+                            integration_registry,
+                        ) {
+                            // Headers are already committed. Log and abort rather
+                            // than trying to replace the response mid-stream.
+                            log::error!("Streaming processing failed: {err:?}");
+                            drop(streaming_body);
+                        } else if let Err(err) = streaming_body.finish() {
+                            log::error!("Failed to finish streaming body: {err}");
+                        } else {
+                            stream_succeeded = true;
+                        }
+
+                        let pull_sync_context = if is_real_browser && stream_succeeded {
+                            build_pull_sync_context(&ec_context)
+                        } else {
+                            None
+                        };
+
+                        return Ok(RouteOutcome {
+                            response: None,
+                            pull_sync_context,
+                        });
+                    }
+                    Ok(PublisherResponse::PassThrough { mut response, body }) => {
+                        response.set_body(body);
+                        (Ok(response), true)
+                    }
+                    Ok(PublisherResponse::Buffered(response)) => (Ok(response), true),
+                    Err(e) => {
+                        log::error!("Failed to proxy to publisher origin: {:?}", e);
+                        (Err(e), true)
                     }
-
-                    let pull_sync_context = if is_real_browser && stream_succeeded {
-                        build_pull_sync_context(&ec_context)
-                    } else {
-                        None
-                    };
-
-                    return Ok(RouteOutcome {
-                        response: None,
-                        pull_sync_context,
-                    });
-                }
-                Ok(PublisherResponse::PassThrough { mut response, body }) => {
-                    response.set_body(body);
-                    (Ok(response), true)
-                }
-                Ok(PublisherResponse::Buffered(response)) => (Ok(response), true),
-                Err(e) => {
-                    log::error!("Failed to proxy to publisher origin: {:?}", e);
-                    (Err(e), true)
                 }
             }
         }
@@ -503,17 +555,20 @@ async fn route_request(
     // Bot gate still suppresses generated EC writes and pull sync via
     // `kv_graph = None`, but withdrawal finalization uses `finalize_kv_graph`
     // so revocation tombstones are not lost for bot-classified clients.
-    ec_finalize_response(
-        settings,
-        &ec_context,
-        finalize_kv_graph.as_ref(),
-        partner_registry,
-        eids_cookie.as_deref(),
-        sharedid_cookie.as_deref(),
-        &mut response,
-    );
+    if should_finalize_ec {
+        ec_finalize_response(
+            settings,
+            &ec_context,
+            finalize_kv_graph.as_ref(),
+            partner_registry,
+            eids_cookie.as_deref(),
+            sharedid_cookie.as_deref(),
+            &mut response,
+        );
+    }
 
     finalize_response(settings, geo_info.as_ref(), &mut response);
+    asset_cache_policy.apply_after_route_finalization(&mut response);
 
     let pull_sync_context = if is_real_browser && organic_route && route_succeeded {
         // Pull sync is intentionally refreshed only from successful organic