1 use anyhow::{Context, Result};
2 use test_programs::wasi::http::types::{
3     Headers, IncomingRequest, Method, OutgoingBody, OutgoingResponse, ResponseOutparam,
4 };
5 
6 struct T;
7 
8 test_programs::proxy::export!(T);
9 
10 impl test_programs::proxy::exports::wasi::http::incoming_handler::Guest for T {
11     fn handle(request: IncomingRequest, outparam: ResponseOutparam) {
12         assert!(request.scheme().is_some());
13         assert!(request.authority().is_some());
14         assert!(request.path_with_query().is_some());
15 
16         test_filesystem();
17 
18         match (request.method(), request.path_with_query().as_deref()) {
19             (Method::Get, Some("/early_drop")) => {
20                 // Ignore all the errors for this endpoint.
21                 let resp = OutgoingResponse::new(Headers::new());
22                 let body = resp.body().expect("outgoing response");
23                 ResponseOutparam::set(outparam, Ok(resp));
24                 let _ = body.write().and_then(|out| {
25                     let _ = out.blocking_write_and_flush(b"hello, world!");
26                     drop(out);
27                     Ok(())
28                 });
29                 let _ = OutgoingBody::finish(body, None);
30 
31                 return;
32             }
33             (Method::Get, Some(p)) if p.starts_with("/modify_fields/") => {
34                 let r = modify_fields_handler(request);
35                 response_for(r, outparam);
36                 return;
37             }
38             (Method::Get, Some(p)) if p.starts_with("/new_fields/") => {
39                 let r = new_fields_handler(request);
40                 response_for(r, outparam);
41                 return;
42             }
43 
44             _ => {}
45         }
46 
47         let header = String::from("custom-forbidden-header");
48         let req_hdrs = request.headers();
49 
50         assert!(
51             !req_hdrs.has(&header),
52             "forbidden `custom-forbidden-header` found in request"
53         );
54 
55         assert!(req_hdrs.delete(&header).is_err());
56         assert!(req_hdrs.append(&header, b"no".as_ref()).is_err());
57 
58         assert!(
59             !req_hdrs.has(&header),
60             "append of forbidden header succeeded"
61         );
62 
63         assert!(
64             !req_hdrs.has("host"),
65             "forbidden host header present in incoming request"
66         );
67 
68         let hdrs = Headers::new();
69         let resp = OutgoingResponse::new(hdrs);
70         let body = resp.body().expect("outgoing response");
71 
72         ResponseOutparam::set(outparam, Ok(resp));
73 
74         let out = body.write().expect("outgoing stream");
75         out.blocking_write_and_flush(b"hello, world!")
76             .expect("writing response");
77 
78         drop(out);
79         OutgoingBody::finish(body, None).expect("outgoing-body.finish");
80     }
81 }
82 
83 fn response_for(r: Result<()>, outparam: ResponseOutparam) {
84     let resp = OutgoingResponse::new(Headers::new());
85     resp.set_status_code(if r.is_ok() { 200 } else { 500 })
86         .unwrap();
87     let body = resp.body().expect("outgoing response");
88     ResponseOutparam::set(outparam, Ok(resp));
89     let _ = body.write().and_then(|out| {
90         let _ = out.blocking_write_and_flush(format!("{r:?}").as_bytes());
91         drop(out);
92         Ok(())
93     });
94     let _ = OutgoingBody::finish(body, None);
95 }
96 
97 // Technically this should not be here for a proxy, but given the current
98 // framework for tests it's required since this file is built as a `bin`
99 fn main() {}
100 
101 fn test_filesystem() {
102     assert!(std::fs::File::open(".").is_err());
103 }
104 
105 fn add_bytes_to_headers(headers: Headers, size: usize) {
106     if size == 0 {
107         return;
108     } else if size < 10 {
109         headers.append("k", &b"abcdefghi"[0..size - 1]).unwrap()
110     } else {
111         for chunk in 0..(size / 10) {
112             let k = format!("g{chunk:04}");
113             let mut v = format!("h{chunk:04}");
114             if chunk == 0 {
115                 for _ in 0..(size % 10) {
116                     v.push('#');
117                 }
118             }
119             headers.append(k.as_str(), v.as_bytes()).unwrap()
120         }
121     }
122 }
123 
124 fn modify_fields_handler(request: IncomingRequest) -> Result<()> {
125     let path = request.path_with_query().unwrap();
126     let rest = path.trim_start_matches("/modify_fields/");
127     let added_field_bytes: usize = rest
128         .parse()
129         .context("expect remainder of url to parse as number")?;
130     add_bytes_to_headers(request.headers().clone(), added_field_bytes);
131 
132     Ok(())
133 }
134 fn new_fields_handler(request: IncomingRequest) -> Result<()> {
135     let path = request.path_with_query().unwrap();
136     let rest = path.trim_start_matches("/new_fields/");
137     let added_field_bytes: usize = rest
138         .parse()
139         .context("expect remainder of url to parse as number")?;
140     add_bytes_to_headers(Headers::new(), added_field_bytes);
141 
142     Ok(())
143 }
144