1 use {
2     flate2::{
3         Compression,
4         write::{DeflateDecoder, DeflateEncoder},
5     },
6     std::{io::Write, mem},
7     test_programs::p3::{
8         proxy::exports::wasi::http::handler::Guest as Handler,
9         wasi::http::{
10             handler,
11             types::{ErrorCode, Headers, Request, Response},
12         },
13         wit_future, wit_stream,
14     },
15     wit_bindgen::StreamResult,
16 };
17 
18 struct Component;
19 
20 test_programs::p3::proxy::export!(Component);
21 
22 impl Handler for Component {
23     /// Forward the specified request to the imported `wasi:http/handler`, transparently decoding the request body
24     /// if it is `deflate`d and then encoding the response body if the client has provided an `accept-encoding:
25     /// deflate` header.
26     async fn handle(request: Request) -> Result<Response, ErrorCode> {
27         // First, extract the parts of the request and check for (and remove) headers pertaining to body encodings.
28         let method = request.get_method();
29         let scheme = request.get_scheme();
30         let path_with_query = request.get_path_with_query();
31         let authority = request.get_authority();
32         let mut accept_deflated = false;
33         let mut content_deflated = false;
34         let headers = request.get_headers();
35         let mut headers = headers.copy_all();
36         headers.retain(|(k, v)| match (k.as_str(), v.as_slice()) {
37             ("accept-encoding", value)
38                 if std::str::from_utf8(value)
39                     .map(|v| v.contains("deflate"))
40                     .unwrap_or(false) =>
41             {
42                 accept_deflated = true;
43                 false
44             }
45             ("content-encoding", b"deflate") => {
46                 content_deflated = true;
47                 false
48             }
49             _ => true,
50         });
51         let (_, result_rx) = wit_future::new(|| Ok(()));
52         let (mut body, trailers) = Request::consume_body(request, result_rx);
53 
54         let (body, trailers) = if content_deflated {
55             // Next, spawn a task to pipe and decode the original request body and trailers into a new request
56             // we'll create below.  This will run concurrently with any code in the imported `wasi:http/handler`.
57             let (trailers_tx, trailers_rx) = wit_future::new(|| todo!());
58             let (mut pipe_tx, pipe_rx) = wit_stream::new();
59 
60             wit_bindgen::spawn(async move {
61                 {
62                     let mut decoder = DeflateDecoder::new(Vec::new());
63 
64                     let (mut status, mut chunk) = body.read(Vec::with_capacity(64 * 1024)).await;
65                     while let StreamResult::Complete(_) = status {
66                         decoder.write_all(&chunk).unwrap();
67                         let remaining = pipe_tx.write_all(mem::take(decoder.get_mut())).await;
68                         assert!(remaining.is_empty());
69                         *decoder.get_mut() = remaining;
70                         chunk.clear();
71                         (status, chunk) = body.read(chunk).await;
72                     }
73 
74                     let remaining = pipe_tx.write_all(decoder.finish().unwrap()).await;
75                     assert!(remaining.is_empty());
76 
77                     drop(pipe_tx);
78                 }
79 
80                 trailers_tx.write(trailers.await).await.unwrap();
81             });
82 
83             (pipe_rx, trailers_rx)
84         } else {
85             (body, trailers)
86         };
87 
88         // While the above task (if any) is running, synthesize a request from the parts collected above and pass
89         // it to the imported `wasi:http/handler`.
90         let (my_request, _request_complete) = Request::new(
91             Headers::from_list(&headers).unwrap(),
92             Some(body),
93             trailers,
94             None,
95         );
96         my_request.set_method(&method).unwrap();
97         my_request.set_scheme(scheme.as_ref()).unwrap();
98         my_request
99             .set_path_with_query(path_with_query.as_deref())
100             .unwrap();
101         my_request.set_authority(authority.as_deref()).unwrap();
102 
103         let response = handler::handle(my_request).await?;
104 
105         // Now that we have the response, extract the parts, adding an extra header if we'll be encoding the body.
106         let status_code = response.get_status_code();
107         let mut headers = response.get_headers().copy_all();
108         if accept_deflated {
109             headers.push(("content-encoding".into(), b"deflate".into()));
110         }
111 
112         let (_, result_rx) = wit_future::new(|| Ok(()));
113         let (mut body, trailers) = Response::consume_body(response, result_rx);
114         let (body, trailers) = if accept_deflated {
115             headers.retain(|(name, _value)| name != "content-length");
116 
117             // Spawn another task; this one is to pipe and encode the original response body and trailers into a
118             // new response we'll create below.  This will run concurrently with the caller's code (i.e. it won't
119             // necessarily complete before we return a value).
120             let (trailers_tx, trailers_rx) = wit_future::new(|| todo!());
121             let (mut pipe_tx, pipe_rx) = wit_stream::new();
122 
123             wit_bindgen::spawn(async move {
124                 {
125                     let mut encoder = DeflateEncoder::new(Vec::new(), Compression::fast());
126                     let (mut status, mut chunk) = body.read(Vec::with_capacity(64 * 1024)).await;
127 
128                     while let StreamResult::Complete(_) = status {
129                         encoder.write_all(&chunk).unwrap();
130                         let remaining = pipe_tx.write_all(mem::take(encoder.get_mut())).await;
131                         assert!(remaining.is_empty());
132                         *encoder.get_mut() = remaining;
133                         chunk.clear();
134                         (status, chunk) = body.read(chunk).await;
135                     }
136 
137                     let remaining = pipe_tx.write_all(encoder.finish().unwrap()).await;
138                     assert!(remaining.is_empty());
139 
140                     drop(pipe_tx);
141                 }
142 
143                 trailers_tx.write(trailers.await).await.unwrap();
144             });
145 
146             (pipe_rx, trailers_rx)
147         } else {
148             (body, trailers)
149         };
150 
151         // While the above tasks (if any) are running, synthesize a response from the parts collected above and
152         // return it.
153         let (my_response, _response_complete) =
154             Response::new(Headers::from_list(&headers).unwrap(), Some(body), trailers);
155         my_response.set_status_code(status_code).unwrap();
156 
157         Ok(my_response)
158     }
159 }
160 
161 // Unused function; required since this file is built as a `bin`:
162 fn main() {}
163