|
| 1 | +// Copyright 2020 Google LLC |
| 2 | +// |
| 3 | +// Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +// you may not use this file except in compliance with the License. |
| 5 | +// You may obtain a copy of the License at |
| 6 | +// |
| 7 | +// http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +// |
| 9 | +// Unless required by applicable law or agreed to in writing, software |
| 10 | +// distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +// See the License for the specific language governing permissions and |
| 13 | +// limitations under the License. |
| 14 | + |
| 15 | +use proxy_wasm::hostcalls; |
| 16 | +use proxy_wasm::promise::Promise; |
| 17 | +use proxy_wasm::traits::*; |
| 18 | +use proxy_wasm::types::*; |
| 19 | +use std::collections::HashMap; |
| 20 | +use std::rc::Rc; |
| 21 | +use std::time::Duration; |
| 22 | + |
| 23 | +proxy_wasm::main! {{ |
| 24 | + proxy_wasm::set_log_level(LogLevel::Trace); |
| 25 | + proxy_wasm::set_http_context(|_, _| -> Box<dyn HttpContext> { Box::new(HttpParallelCall::default()) }); |
| 26 | +}} |
| 27 | + |
| 28 | +#[derive(Default)] |
| 29 | +struct HttpParallelCall { |
| 30 | + m: HashMap<u32, Rc<Promise<(u32, usize, usize, usize)>>>, |
| 31 | +} |
| 32 | + |
| 33 | +impl HttpContext for HttpParallelCall { |
| 34 | + fn on_http_request_headers(&mut self, _: usize, _: bool) -> Action { |
| 35 | + // "Hello, " |
| 36 | + let token1 = self |
| 37 | + .dispatch_http_call( |
| 38 | + "httpbin", |
| 39 | + vec![ |
| 40 | + (":method", "GET"), |
| 41 | + (":path", "/base64/SGVsbG8sIAo="), |
| 42 | + (":authority", "httpbin.org"), |
| 43 | + ], |
| 44 | + None, |
| 45 | + vec![], |
| 46 | + Duration::from_secs(1), |
| 47 | + ) |
| 48 | + .unwrap(); |
| 49 | + |
| 50 | + // "World!" |
| 51 | + let token2 = self |
| 52 | + .dispatch_http_call( |
| 53 | + "httpbin", |
| 54 | + vec![ |
| 55 | + (":method", "GET"), |
| 56 | + (":path", "/base64/V29ybGQhCg=="), |
| 57 | + (":authority", "httpbin.org"), |
| 58 | + ], |
| 59 | + None, |
| 60 | + vec![], |
| 61 | + Duration::from_secs(1), |
| 62 | + ) |
| 63 | + .unwrap(); |
| 64 | + |
| 65 | + let promise1 = Promise::new(); |
| 66 | + let promise2 = Promise::new(); |
| 67 | + self.m.insert(token1, promise1.clone()); |
| 68 | + self.m.insert(token2, promise2.clone()); |
| 69 | + |
| 70 | + Promise::all_of(vec![ |
| 71 | + promise1 |
| 72 | + .then(|(_, _, _body_size, _)| get_http_call_response_body_string(0, _body_size)) |
| 73 | + .then(|body| body.unwrap_or_else(|| "".to_string())), |
| 74 | + promise2 |
| 75 | + .then(|(_, _, _body_size, _)| get_http_call_response_body_string(0, _body_size)) |
| 76 | + .then(|body| body.unwrap_or_else(|| "".to_string())), |
| 77 | + ]) |
| 78 | + .then(|results| { |
| 79 | + send_http_response( |
| 80 | + 200, |
| 81 | + vec![], |
| 82 | + Some( |
| 83 | + format!( |
| 84 | + "{}{}\n", |
| 85 | + results[0].strip_suffix("\n").unwrap(), |
| 86 | + results[1].strip_suffix("\n").unwrap() |
| 87 | + ) |
| 88 | + .as_bytes(), |
| 89 | + ), |
| 90 | + ); |
| 91 | + }); |
| 92 | + |
| 93 | + Action::Pause |
| 94 | + } |
| 95 | + |
| 96 | + fn on_http_response_headers(&mut self, _: usize, _: bool) -> Action { |
| 97 | + self.set_http_response_header("Powered-By", Some("proxy-wasm")); |
| 98 | + Action::Continue |
| 99 | + } |
| 100 | +} |
| 101 | + |
| 102 | +impl Context for HttpParallelCall { |
| 103 | + fn on_http_call_response( |
| 104 | + &mut self, |
| 105 | + _token_id: u32, |
| 106 | + _num_headers: usize, |
| 107 | + _body_size: usize, |
| 108 | + _num_trailers: usize, |
| 109 | + ) { |
| 110 | + let promise = self.m.remove(&_token_id); |
| 111 | + promise |
| 112 | + .unwrap() |
| 113 | + .fulfill((_token_id, _num_headers, _body_size, _num_trailers)); |
| 114 | + } |
| 115 | +} |
| 116 | + |
| 117 | +fn get_http_call_response_body_string(start: usize, max_size: usize) -> Option<String> { |
| 118 | + match hostcalls::get_buffer(BufferType::HttpCallResponseBody, start, max_size).unwrap() { |
| 119 | + None => None, |
| 120 | + Some(bytes) => { |
| 121 | + let body_string = String::from_utf8(bytes.to_vec()).unwrap(); |
| 122 | + Some(body_string) |
| 123 | + } |
| 124 | + } |
| 125 | +} |
| 126 | + |
| 127 | +fn send_http_response(status_code: u32, headers: Vec<(&str, &str)>, body: Option<&[u8]>) { |
| 128 | + hostcalls::send_http_response(status_code, headers, body).unwrap() |
| 129 | +} |
0 commit comments