🛡 Initial shielding support for Viceroy

cli/tests/integration/common.rs

@@ -17,6 +17,7 @@ use viceroy_lib::{
     body::Body,
     config::{
         Acls, DeviceDetection, Dictionaries, FastlyConfig, Geolocation, ObjectStores, SecretStores,
+        ShieldingSites,
     },
     ExecuteCtx, ProfilingStrategy, ViceroyService,
 };
@@ -84,6 +85,7 @@ pub struct Test {
     geolocation: Geolocation,
     object_stores: ObjectStores,
     secret_stores: SecretStores,
+    shielding_sites: ShieldingSites,
     capture_logs: Arc<Mutex<dyn Write + Send>>,
     log_stdout: bool,
     log_stderr: bool,
@@ -107,6 +109,7 @@ impl Test {
             geolocation: Geolocation::new(),
             object_stores: ObjectStores::new(),
             secret_stores: SecretStores::new(),
+            shielding_sites: ShieldingSites::new(),
             capture_logs: Arc::new(Mutex::new(std::io::stdout())),
             log_stdout: false,
             log_stderr: false,
@@ -130,6 +133,7 @@ impl Test {
             geolocation: Geolocation::new(),
             object_stores: ObjectStores::new(),
             secret_stores: SecretStores::new(),
+            shielding_sites: ShieldingSites::new(),
             capture_logs: Arc::new(Mutex::new(std::io::stdout())),
             log_stdout: false,
             log_stderr: false,
@@ -150,6 +154,7 @@ impl Test {
             geolocation: config.geolocation().to_owned(),
             object_stores: config.object_stores().to_owned(),
             secret_stores: config.secret_stores().to_owned(),
+            shielding_sites: config.shielding_sites().to_owned(),
             ..self
         })
     }
@@ -339,6 +344,7 @@ impl Test {
         .with_geolocation(self.geolocation.clone())
         .with_object_stores(self.object_stores.clone())
         .with_secret_stores(self.secret_stores.clone())
+        .with_shielding_sites(self.shielding_sites.clone())
         .with_capture_logs(self.capture_logs.clone())
         .with_log_stderr(self.log_stderr)
         .with_log_stdout(self.log_stdout);

cli/tests/integration/main.rs

@@ -21,6 +21,7 @@ mod request;
 mod response;
 mod secret_store;
 mod sending_response;
+mod shielding;
 mod sleep;
 mod unknown_import_behavior;
 mod upstream;

cli/tests/integration/shielding.rs

@@ -0,0 +1,129 @@
+use crate::common::{Test, TestResult};
+use crate::viceroy_test;
+use hyper::service::{make_service_fn, service_fn};
+use hyper::{body, Body, Request, Response, Server, StatusCode};
+use std::convert::Infallible;
+use std::net::SocketAddr;
+
+viceroy_test!(shielding_running_on, |is_component| {
+    const FASTLY_TOML: &str = r#"
+        name = "shielding_running_on"
+        description = "test that running_on works"
+        authors = ["Test user <[email protected]>"]
+        language = "rust"
+        [local_server.shielding_sites]
+        "pdx-or-us" = "Local"
+        "bfi-wa-us".unencrypted = "http://localhost"
+        "bfi-wa-us".encrypted = "https://localhost"
+    "#;
+    let test = Test::using_fixture("shielding.wasm")
+        .using_fastly_toml(FASTLY_TOML)?
+        .adapt_component(is_component);
+
+    let resp1 = test
+        .against(
+            Request::get("/is-shield")
+                .header("shield", "waffle-cone")
+                .body("")?,
+        )
+        .await?;
+    assert_eq!(StatusCode::INTERNAL_SERVER_ERROR, resp1.status());
+
+    let resp2 = test
+        .against(
+            Request::get("/is-shield")
+                .header("shield", "pdx-or-us")
+                .body("")?,
+        )
+        .await?;
+    assert_eq!(StatusCode::OK, resp2.status());
+    let body = body::to_bytes(resp2.into_body()).await.unwrap().to_vec();
+    let string = String::from_utf8(body).unwrap();
+    assert_eq!("true", &string);
+
+    let resp3 = test
+        .against(
+            Request::get("/is-shield")
+                .header("shield", "bfi-wa-us")
+                .body("")?,
+        )
+        .await?;
+    assert_eq!(StatusCode::OK, resp3.status());
+    let body = body::to_bytes(resp3.into_body()).await.unwrap().to_vec();
+    let string = String::from_utf8(body).unwrap();
+    assert_eq!("false", &string);
+
+    Ok(())
+});
+
+viceroy_test!(shield_backends, |is_component| {
+    let blank_addr = SocketAddr::from(([0, 0, 0, 0], 0));
+
+    let make_service_unenc = make_service_fn(|_conn| async {
+        Ok::<_, Infallible>(service_fn(|_: Request<Body>| async {
+            Ok::<Response<Body>, std::io::Error>(Response::new(Body::from("unencrypted land")))
+        }))
+    });
+
+    let make_service_enc = make_service_fn(|_conn| async {
+        Ok::<_, Infallible>(service_fn(|_: Request<Body>| async {
+            Ok::<Response<Body>, std::io::Error>(Response::new(Body::from("encrypted land")))
+        }))
+    });
+
+    let unenc_server = Server::bind(&blank_addr).serve(make_service_unenc);
+    let enc_server = Server::bind(&blank_addr).serve(make_service_enc);
+
+    let unenc_addr = unenc_server.local_addr();
+    let enc_addr = enc_server.local_addr();
+
+    let mut server_set = tokio::task::JoinSet::new();
+    server_set.spawn(unenc_server);
+    server_set.spawn(enc_server);
+
+    let fastly_toml = format!(
+        r#"
+        name = "shielding_running_on"
+        description = "test that running_on works"
+        authors = ["Test user <[email protected]>"]
+        language = "rust"
+        [local_server.shielding_sites]
+        "pdx-or-us" = "Local"
+        "bfi-wa-us".unencrypted = "http://{}"
+        "bfi-wa-us".encrypted = "http://{}"
+    "#,
+        unenc_addr, enc_addr
+    );
+
+    let test = Test::using_fixture("shielding.wasm")
+        .using_fastly_toml(&fastly_toml)?
+        .adapt_component(is_component);
+
+    let resp1 = test
+        .against(
+            Request::get("/shield-to")
+                .header("shield", "bfi-wa-us")
+                .header("shield-type", "unencrypted")
+                .body("")?,
+        )
+        .await?;
+    assert_eq!(StatusCode::OK, resp1.status());
+    let body = body::to_bytes(resp1.into_body()).await.unwrap().to_vec();
+    let string = String::from_utf8(body).unwrap();
+    assert_eq!("unencrypted land", &string);
+
+    let resp2 = test
+        .against(
+            Request::get("/shield-to")
+                .header("shield", "bfi-wa-us")
+                .header("shield-type", "encrypted")
+                .body("")?,
+        )
+        .await?;
+    assert_eq!(StatusCode::OK, resp2.status());
+    let body = body::to_bytes(resp2.into_body()).await.unwrap().to_vec();
+    let string = String::from_utf8(body).unwrap();
+    assert_eq!("encrypted land", &string);
+
+    Ok(())
+});

crates/adapter/src/fastly/core.rs

@@ -3610,3 +3610,164 @@ pub mod fastly_purge {
         )
     }
 }
+
+pub mod fastly_shielding {
+    use super::*;
+    use crate::bindings::fastly::api::{shielding as host, types};
+    use std::slice;
+
+    bitflags::bitflags! {
+        #[derive(Default)]
+        #[repr(transparent)]
+        pub struct ShieldBackendOptions: u32 {
+            const RESERVED = 1 << 0;
+            const CACHE_KEY = 1 << 1;
+        }
+    }
+
+    #[repr(C)]
+    pub struct ShieldBackendConfig {
+        pub cache_key: *const u8,
+        pub cache_key_len: u32,
+    }
+
+    impl Default for ShieldBackendConfig {
+        fn default() -> Self {
+            ShieldBackendConfig {
+                cache_key: std::ptr::null(),
+                cache_key_len: 0,
+            }
+        }
+    }
+
+    #[export_name = "fastly_shielding#shield_info"]
+    pub fn shield_info(
+        name: *const u8,
+        name_len: usize,
+        info_block: *mut u8,
+        info_block_len: usize,
+        nwritten_out: *mut u32,
+    ) -> FastlyStatus {
+        let name = unsafe { slice::from_raw_parts(name, name_len) };
+        with_buffer!(
+            info_block,
+            info_block_len,
+            { host::shield_info(name, u64::try_from(info_block_len).trapping_unwrap()) },
+            |res| {
+                match res {
+                    Ok(res) => {
+                        unsafe {
+                            *nwritten_out = u32::try_from(res.len()).unwrap_or(0);
+                        }
+                        std::mem::forget(res);
+                    }
+
+                    Err(e) => {
+                        if let types::Error::BufferLen(needed) = e {
+                            unsafe {
+                                *nwritten_out = u32::try_from(needed).unwrap_or(0);
+                            }
+                        }
+
+                        return Err(e.into());
+                    }
+                }
+            }
+        )
+    }
+
+    impl From<ShieldBackendOptions> for host::ShieldBackendOptionsMask {
+        fn from(value: ShieldBackendOptions) -> Self {
+            let mut flags = Self::empty();
+
+            flags.set(
+                Self::RESERVED,
+                value.contains(ShieldBackendOptions::RESERVED),
+            );
+            flags.set(
+                Self::CACHE_KEY,
+                value.contains(ShieldBackendOptions::CACHE_KEY),
+            );
+
+            flags
+        }
+    }
+
+    fn shield_backend_options(
+        mask: ShieldBackendOptions,
+        options: *const ShieldBackendConfig,
+    ) -> (host::ShieldBackendOptionsMask, host::ShieldBackendOptions) {
+        let mask = host::ShieldBackendOptionsMask::from(mask);
+
+        // NOTE: this is only really safe because we never mutate the vectors -- we only need
+        // vectors to satisfy the interface produced by the DynamicBackendConfig record,
+        // `register_dynamic_backend` will never mutate the vectors it's given.
+        macro_rules! make_vec {
+            ($ptr_field:ident, $len_field:ident) => {
+                unsafe {
+                    let len = usize::try_from((*options).$len_field).trapping_unwrap();
+                    Vec::from_raw_parts((*options).$ptr_field as *mut _, len, len)
+                }
+            };
+        }
+
+        let options = host::ShieldBackendOptions {
+            cache_key: if mask.contains(host::ShieldBackendOptionsMask::CACHE_KEY) {
+                make_vec!(cache_key, cache_key_len)
+            } else {
+                Vec::new()
+            },
+        };
+
+        (mask, options)
+    }
+
+    /// Turn a pop name into a backend that we can send requests to.
+    #[export_name = "fastly_shielding#backend_for_shield"]
+    pub fn backend_for_shield(
+        name: *const u8,
+        name_len: usize,
+        options_mask: ShieldBackendOptions,
+        options: *const ShieldBackendConfig,
+        backend_name: *mut u8,
+        backend_name_len: usize,
+        nwritten_out: *mut u32,
+    ) -> FastlyStatus {
+        let name = unsafe { slice::from_raw_parts(name, name_len) };
+        let (mask, options) = shield_backend_options(options_mask, options);
+        with_buffer!(
+            backend_name,
+            backend_name_len,
+            {
+                let res = host::backend_for_shield(
+                    name,
+                    mask,
+                    &options,
+                    u64::try_from(backend_name_len).trapping_unwrap(),
+                );
+                std::mem::forget(options);
+                res
+            },
+            |res| {
+                match res {
+                    Ok(res) => {
+                        unsafe {
+                            *nwritten_out = u32::try_from(res.len()).unwrap_or(0);
+                        }
+                        std::mem::forget(res);
+                    }
+
+                    Err(e) => {
+                        if let types::Error::BufferLen(needed) = e {
+                            unsafe {
+                                *nwritten_out = u32::try_from(needed).unwrap_or(0);
+                            }
+                        }
+
+                        return Err(e.into());
+                    }
+                }
+            }
+        )
+    }
+}

lib/compute-at-edge-abi/compute-at-edge.witx

@@ -2,6 +2,7 @@
 (use "cache.witx")
 (use "config-store.witx")
 (use "http-cache.witx")
+(use "shielding.witx")
 
 (module $fastly_abi
     (@interface func (export "init")

lib/compute-at-edge-abi/shielding.witx

@@ -0,0 +1,31 @@
+(typename $shield_backend_options
+  (flags (@witx repr u32)
+    $reserved
+    $use_cache_key
+  ))
+
+(typename $shield_backend_config
+  (record
+    (field $cache_key (@witx pointer (@witx char8)))
+    (field $cache_key_len u32)
+    ))
+
+(module $fastly_shielding
+
+  (@interface func (export "shield_info")
+    (param $name string)
+    (param $info_block (@witx pointer (@witx char8)))
+    (param $info_block_max_len (@witx usize))
+    (result $err (expected $num_bytes (error $fastly_status)))
+  )
+
+  (@interface func (export "backend_for_shield")
+    (param $shield_name string)
+    (param $backend_config_mask $shield_backend_options)
+    (param $backend_configuration (@witx pointer $shield_backend_config))
+    (param $backend_name_out (@witx pointer (@witx char8)))
+    (param $backend_name_max_len (@witx usize))
+    (result $err (expected $num_bytes (error $fastly_status)))
+  )
+
+)