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 (mut body, trailers) = request.consume_body().unwrap();
52 
53         let (body, trailers) = if content_deflated {
54             // Next, spawn a task to pipe and decode the original request body and trailers into a new request
55             // we'll create below.  This will run concurrently with any code in the imported `wasi:http/handler`.
56             let (trailers_tx, trailers_rx) = wit_future::new(|| todo!());
57             let (mut pipe_tx, pipe_rx) = wit_stream::new();
58 
59             wit_bindgen::spawn(async move {
60                 {
61                     let mut decoder = DeflateDecoder::new(Vec::new());
62 
63                     let (mut status, mut chunk) = body.read(Vec::with_capacity(64 * 1024)).await;
64                     while let StreamResult::Complete(_) = status {
65                         decoder.write_all(&chunk).unwrap();
66                         let remaining = pipe_tx.write_all(mem::take(decoder.get_mut())).await;
67                         assert!(remaining.is_empty());
68                         *decoder.get_mut() = remaining;
69                         chunk.clear();
70                         (status, chunk) = body.read(chunk).await;
71                     }
72 
73                     let remaining = pipe_tx.write_all(decoder.finish().unwrap()).await;
74                     assert!(remaining.is_empty());
75 
76                     drop(pipe_tx);
77                 }
78 
79                 trailers_tx.write(trailers.await).await.unwrap();
80 
81                 drop(request);
82             });
83 
84             (pipe_rx, trailers_rx)
85         } else {
86             (body, trailers)
87         };
88 
89         // While the above task (if any) is running, synthesize a request from the parts collected above and pass
90         // it to the imported `wasi:http/handler`.
91         let (my_request, _request_complete) = Request::new(
92             Headers::from_list(&headers).unwrap(),
93             Some(body),
94             trailers,
95             None,
96         );
97         my_request.set_method(&method).unwrap();
98         my_request.set_scheme(scheme.as_ref()).unwrap();
99         my_request
100             .set_path_with_query(path_with_query.as_deref())
101             .unwrap();
102         my_request.set_authority(authority.as_deref()).unwrap();
103 
104         let response = handler::handle(my_request).await?;
105 
106         // Now that we have the response, extract the parts, adding an extra header if we'll be encoding the body.
107         let status_code = response.get_status_code();
108         let mut headers = response.get_headers().copy_all();
109         if accept_deflated {
110             headers.push(("content-encoding".into(), b"deflate".into()));
111         }
112 
113         let (mut body, trailers) = response.consume_body().unwrap();
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                 drop(response);
145             });
146 
147             (pipe_rx, trailers_rx)
148         } else {
149             (body, trailers)
150         };
151 
152         // While the above tasks (if any) are running, synthesize a response from the parts collected above and
153         // return it.
154         let (my_response, _response_complete) =
155             Response::new(Headers::from_list(&headers).unwrap(), Some(body), trailers);
156         my_response.set_status_code(status_code).unwrap();
157 
158         Ok(my_response)
159     }
160 }
161 
162 // Unused function; required since this file is built as a `bin`:
163 fn main() {}
164