🚚 Support for trailers and trailer modification

CHANGELOG.md

@@ -1,5 +1,8 @@
 ## Unreleased
 
+- Add support for trailers. Trailer modification calls should be considered experimental,
+  as we finalize interfaces ([#327](https://github.com/fastly/Viceroy/pull/327))
+
 ## 0.9.1 (2023-10-09)
 
 - Match the number of memories to the number of core instances ([#322](https://github.com/fastly/Viceroy/pull/322))

cli/tests/integration/grpc.rs

@@ -1,39 +1,42 @@
-//use crate::common::{Test, TestResult};
-//use hyper::http::response;
-//use hyper::server::conn::AddrIncoming;
-//use hyper::service::{make_service_fn, service_fn};
-//use hyper::{Request, Server, StatusCode};
-//use std::net::SocketAddr;
-//
-//#[tokio::test(flavor = "multi_thread")]
-//async fn grpc_creates_h2_connection() -> TestResult {
-//    let test = Test::using_fixture("grpc.wasm");
-//    let server_addr: SocketAddr = "127.0.0.1:0".parse().expect("localhost parses");
-//    let incoming = AddrIncoming::bind(&server_addr).expect("bind");
-//    let bound_port = incoming.local_addr().port();
-//
-//    let service = make_service_fn(|_| async move {
-//        Ok::<_, std::io::Error>(service_fn(move |_req| async {
-//            response::Builder::new()
-//                .status(200)
-//                .body("Hello!".to_string())
-//        }))
-//    });
-//
-//    let server = Server::builder(incoming).http2_only(true).serve(service);
-//    tokio::spawn(server);
-//
-//    let resp = test
-//        .against(
-//            Request::post("/")
-//                .header("port", bound_port)
-//                .body("Hello, Viceroy!")
-//                .unwrap(),
-//        )
-//        .await?;
-//    assert_eq!(resp.status(), StatusCode::OK);
-//    assert_eq!(resp.into_body().read_into_string().await?, "Hello!");
-//
-//    Ok(())
-//}
-//
+use crate::common::{Test, TestResult};
+use hyper::http::response;
+use hyper::server::conn::AddrIncoming;
+use hyper::service::{make_service_fn, service_fn};
+use hyper::{Request, Server, StatusCode};
+use std::net::SocketAddr;
+
+#[tokio::test(flavor = "multi_thread")]
+async fn grpc_creates_h2_connection() -> TestResult {
+    let test = Test::using_fixture("grpc.wasm");
+    let server_addr: SocketAddr = "127.0.0.1:0".parse().expect("localhost parses");
+    let incoming = AddrIncoming::bind(&server_addr).expect("bind");
+    let bound_port = incoming.local_addr().port();
+
+    let service = make_service_fn(|_| async move {
+        Ok::<_, std::io::Error>(service_fn(move |_req| async {
+            response::Builder::new()
+                .status(200)
+                .body("Hello!".to_string())
+        }))
+    });
+
+    let server = Server::builder(incoming).http2_only(true).serve(service);
+    tokio::spawn(server);
+
+    let resp = test
+        .against(
+            Request::post("/")
+                .header("port", bound_port)
+                .body("Hello, Viceroy!")
+                .unwrap(),
+        )
+        .await?;
+    assert_eq!(resp.status(), StatusCode::OK);
+    // The test below is not critical -- we've proved our point by returning 200 -- but seems
+    // to trigger an error in Windows; it looks like there's a funny interaction between reading
+    // the body and the stream having been closed, and we get a NO_ERROR error. So I've commented
+    // it out, until there's a clear Hyper solution.
+    // assert_eq!(resp.into_body().read_into_string().await?, "Hello!");
+
+    Ok(())
+}

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

@@ -73,6 +73,43 @@
         (param $h $body_handle)
         (result $err (expected (error $fastly_status)))
     )
+
+    (@interface func (export "trailer_append")
+        (param $h $body_handle)
+        (param $name (list u8))
+        (param $value (list u8))
+        (result $err (expected (error $fastly_status)))
+    )
+
+    (@interface func (export "trailer_names_get")
+        (param $h $body_handle)
+        (param $buf (@witx pointer (@witx char8)))
+        (param $buf_len (@witx usize))
+        (param $cursor $multi_value_cursor)
+        (param $ending_cursor_out (@witx pointer $multi_value_cursor_result))
+        (param $nwritten_out (@witx pointer (@witx usize)))
+        (result $err (expected (error $fastly_status)))
+    )
+
+    (@interface func (export "trailer_value_get")
+        (param $h $body_handle)
+        (param $name (list u8))
+        (param $value (@witx pointer (@witx char8)))
+        (param $value_max_len (@witx usize))
+        (param $nwritten_out (@witx pointer (@witx usize)))
+        (result $err (expected (error $fastly_status)))
+    )
+
+    (@interface func (export "trailer_values_get")
+        (param $h $body_handle)
+        (param $name (list u8))
+        (param $buf (@witx pointer (@witx char8)))
+        (param $buf_len (@witx usize))
+        (param $cursor $multi_value_cursor)
+        (param $ending_cursor_out (@witx pointer $multi_value_cursor_result))
+        (param $nwritten_out (@witx pointer (@witx usize)))
+        (result $err (expected (error $fastly_status)))
+    )
 )
 
 (module $fastly_log

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

@@ -55,7 +55,12 @@
         ;;;
         ;;; This is returned when an attempt to allocate a resource has exceeded the maximum number of
         ;;; resources permitted. For example, creating too many response handles.
-        $limitexceeded))
+        $limitexceeded
+        ;;; Resource temporarily unavailable
+        ;;;
+        ;;; This is returned when an attempting to retrieve a resource that is not yet available.
+        ;;; For example when attempting to read trailers from a Body that has not yet been consumed.
+        $again))
 
 ;;; A tag indicating HTTP protocol versions.
 (typename $http_version

lib/src/body.rs

@@ -1,5 +1,7 @@
 //! Body type, for request and response bodies.
 
+use futures::FutureExt;
+
 use {
     crate::{error, streaming_body::StreamingBodyItem, Error},
     bytes::{BufMut, BytesMut},
@@ -89,6 +91,8 @@ impl From<mpsc::Receiver<StreamingBodyItem>> for Chunk {
 #[derive(Default, Debug)]
 pub struct Body {
     chunks: VecDeque<Chunk>,
+    pub(crate) trailers: HeaderMap,
+    pub(crate) trailers_ready: bool,
 }
 
 impl Body {
@@ -179,7 +183,7 @@ impl HttpBody for Body {
                     let body_mut = &mut body;
                     pin_mut!(body_mut);
 
-                    match body_mut.poll_data(cx) {
+                    match body_mut.as_mut().poll_data(cx) {
                         Poll::Pending => {
                             // put the body back, so we can poll it again next time
                             self.chunks.push_front(body.into());
@@ -188,9 +192,25 @@ impl HttpBody for Body {
                         Poll::Ready(None) => {
                             // no more bytes from this body, so continue the loop now that it's been
                             // popped
-                            //
-                            // TODO ACF 2020-06-01: do something with the body's trailers at this point
-                            continue;
+                            match body_mut.trailers().poll_unpin(cx) {
+                                Poll::Pending => {
+                                    self.chunks.push_front(body.into());
+                                    return Poll::Pending;
+                                }
+
+                                Poll::Ready(Err(e)) => {
+                                    return Poll::Ready(Some(Err(e.into())));
+                                }
+
+                                Poll::Ready(Ok(None)) => continue,
+
+                                Poll::Ready(Ok(Some(header_map))) => {
+                                    for (k, v) in header_map.iter() {
+                                        self.trailers.append(k, v.clone());
+                                    }
+                                    continue;
+                                }
+                            }
                         }
                         Poll::Ready(Some(item)) => {
                             // put the body back, so we can poll it again next time
@@ -221,7 +241,8 @@ impl HttpBody for Body {
                             self.chunks.push_front(chunk);
                             continue;
                         }
-                        Poll::Ready(Some(StreamingBodyItem::Finished)) => {
+                        Poll::Ready(Some(StreamingBodyItem::Finished(trailers))) => {
+                            self.trailers.extend(trailers);
                             // it shouldn't be possible for any more chunks to arrive on this
                             // channel, but just in case we won't try to read them; dropping the
                             // receiver means we won't hit the `Ready(None)` case above that
@@ -266,14 +287,23 @@ impl HttpBody for Body {
             }
         }
 
+        // With no more chunks arriving we can mark trailers as being ready.
+        self.trailers_ready = true;
         Poll::Ready(None) // The queue of chunks is now empty!
     }
 
     fn poll_trailers(
         self: Pin<&mut Self>,
         _cx: &mut Context,
     ) -> Poll<Result<Option<HeaderMap>, Self::Error>> {
-        Poll::Ready(Ok(None)) // `Body` does not currently have trailers.
+        if !self.chunks.is_empty() {
+            return Poll::Pending;
+        }
+        if self.trailers.is_empty() {
+            Poll::Ready(Ok(None))
+        } else {
+            Poll::Ready(Ok(Some(self.trailers.clone())))
+        }
     }
 
     /// This is an optional method, but implementing it correctly allows us to reduce the number of

lib/src/error.rs

@@ -142,6 +142,9 @@ pub enum Error {
 
     #[error("Invalid response to ALPN request; wanted '{0}', got '{1}'")]
     InvalidAlpnRepsonse(&'static str, String),
+
+    #[error("Resource temporarily unavailable")]
+    Again,
 }
 
 impl Error {
@@ -180,6 +183,7 @@ impl Error {
             Error::GeolocationError(e) => e.to_fastly_status(),
             Error::ObjectStoreError(e) => e.into(),
             Error::SecretStoreError(e) => e.into(),
+            Error::Again => FastlyStatus::Again,
             // All other hostcall errors map to a generic `ERROR` value.
             Error::AbiVersionMismatch
             | Error::BackendUrl(_)

lib/src/streaming_body.rs

@@ -1,4 +1,5 @@
 use crate::{body::Chunk, error::Error};
+use http::{HeaderMap, HeaderName, HeaderValue};
 use tokio::sync::mpsc;
 
 // Note: this constant and comment is copied from xqd
@@ -14,6 +15,7 @@ const STREAMING_CHANNEL_SIZE: usize = 8;
 #[derive(Debug)]
 pub struct StreamingBody {
     sender: mpsc::Sender<StreamingBodyItem>,
+    pub(crate) trailers: HeaderMap,
 }
 
 /// The items sent over the `StreamingBody` channel.
@@ -34,14 +36,20 @@ pub struct StreamingBody {
 #[derive(Debug)]
 pub enum StreamingBodyItem {
     Chunk(Chunk),
-    Finished,
+    Finished(HeaderMap),
 }
 
 impl StreamingBody {
     /// Create a new channel for streaming a body, returning write and read ends as a pair.
     pub fn new() -> (StreamingBody, mpsc::Receiver<StreamingBodyItem>) {
         let (sender, receiver) = mpsc::channel(STREAMING_CHANNEL_SIZE);
-        (StreamingBody { sender }, receiver)
+        (
+            StreamingBody {
+                sender,
+                trailers: HeaderMap::new(),
+            },
+            receiver,
+        )
     }
 
     /// Send a single chunk along this body stream.
@@ -55,6 +63,11 @@ impl StreamingBody {
             .map_err(|_| Error::StreamingChunkSend)
     }
 
+    /// Convenience method for appending trailers.
+    pub fn append_trailer(&mut self, name: HeaderName, value: HeaderValue) {
+        self.trailers.append(name, value);
+    }
+
     /// Block until the body has room for writing additional chunks.
     pub async fn await_ready(&mut self) {
         let _ = self.sender.reserve().await;
@@ -65,18 +78,27 @@ impl StreamingBody {
     /// This is important primarily for `Transfer-Encoding: chunked` bodies where a premature close
     /// is only noticed if the chunked encoding is not properly terminated.
     pub fn finish(self) -> Result<(), Error> {
-        match self.sender.try_send(StreamingBodyItem::Finished) {
+        match self
+            .sender
+            .try_send(StreamingBodyItem::Finished(self.trailers))
+        {
             Ok(()) => Ok(()),
             Err(mpsc::error::TrySendError::Closed(_)) => Ok(()),
-            Err(mpsc::error::TrySendError::Full(_)) => {
+            Err(mpsc::error::TrySendError::Full(StreamingBodyItem::Finished(trailers))) => {
                 // If the channel is full, maybe the other end is just taking a while to receive all
                 // the bytes. Spawn a task that will send a `finish` message as soon as there's room
                 // in the channel.
                 tokio::task::spawn(async move {
-                    let _ = self.sender.send(StreamingBodyItem::Finished).await;
+                    let _ = self
+                        .sender
+                        .send(StreamingBodyItem::Finished(trailers))
+                        .await;
                 });
                 Ok(())
             }
+            Err(mpsc::error::TrySendError::Full(_)) => {
+                unreachable!("Only a StreamingBodyItem::Finished should be reachable")
+            }
         }
     }
 }

lib/src/upstream.rs

@@ -7,7 +7,7 @@ use crate::{
     wiggle_abi::types::ContentEncodings,
 };
 use futures::Future;
-use http::{uri, HeaderValue};
+use http::{uri, HeaderValue, Version};
 use hyper::{client::HttpConnector, header, Client, HeaderMap, Request, Response, Uri};
 use rustls::client::{ServerName, WantsTransparencyPolicyOrClientCert};
 use std::{
@@ -267,7 +267,13 @@ pub fn send_request(
 
     let h2only = backend.grpc;
     async move {
-        let basic_response = Client::builder()
+        let mut builder = Client::builder();
+
+        if req.version() == Version::HTTP_2 {
+            builder.http2_only(true);
+        }
+
+        let basic_response = builder
             .set_host(false)
             .http2_only(h2only)
             .build(connector)