wasmcp_wasi/
http.rs

1use std::collections::HashMap;
2use std::cell::RefCell;
3use std::rc::Rc;
4use std::task::Poll;
5use std::future::Future;
6use futures::{future, sink, stream, Sink, Stream, TryStreamExt, SinkExt};
7use spin_executor::CancelOnDropToken;
8use anyhow::Result;
9use crate::wit::wasi::http0_2_0::types::{
10    Headers, IncomingBody, IncomingResponse, Method as WasiMethod, OutgoingBody, OutgoingRequest,
11    Scheme, FutureIncomingResponse, ErrorCode
12};
13use crate::wit::wasi::http0_2_0::outgoing_handler;
14use wasi::io::streams::{InputStream, OutputStream, StreamError};
15
16const READ_SIZE: u64 = 16 * 1024;
17
18#[derive(Debug, Clone, PartialEq)]
19pub enum Method {
20    Get,
21    Post,
22    Put,
23    Delete,
24    Patch,
25    Head,
26    Options,
27}
28
29impl From<Method> for WasiMethod {
30    fn from(method: Method) -> Self {
31        match method {
32            Method::Get => WasiMethod::Get,
33            Method::Post => WasiMethod::Post,
34            Method::Put => WasiMethod::Put,
35            Method::Delete => WasiMethod::Delete,
36            Method::Patch => WasiMethod::Patch,
37            Method::Head => WasiMethod::Head,
38            Method::Options => WasiMethod::Options,
39        }
40    }
41}
42
43pub struct Request {
44    method: Method,
45    uri: String,
46    headers: HashMap<String, String>,
47    body: Vec<u8>,
48}
49
50impl Request {
51    pub fn new(method: Method, uri: impl Into<String>) -> Self {
52        Self {
53            method,
54            uri: uri.into(),
55            headers: HashMap::new(),
56            body: Vec::new(),
57        }
58    }
59
60    pub fn get(uri: impl Into<String>) -> Self {
61        Self::new(Method::Get, uri)
62    }
63
64    pub fn post(uri: impl Into<String>, body: impl Into<Vec<u8>>) -> Self {
65        let mut req = Self::new(Method::Post, uri);
66        req.body = body.into();
67        req
68    }
69
70    pub fn builder() -> RequestBuilder {
71        RequestBuilder::new()
72    }
73
74    pub fn header(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
75        self.headers.insert(name.into(), value.into());
76        self
77    }
78
79    pub fn body(mut self, body: impl Into<Vec<u8>>) -> Self {
80        self.body = body.into();
81        self
82    }
83
84    pub fn build(self) -> Self {
85        self
86    }
87}
88
89pub struct RequestBuilder {
90    method: Option<Method>,
91    uri: Option<String>,
92    headers: HashMap<String, String>,
93    body: Vec<u8>,
94}
95
96impl RequestBuilder {
97    pub fn new() -> Self {
98        Self {
99            method: None,
100            uri: None,
101            headers: HashMap::new(),
102            body: Vec::new(),
103        }
104    }
105
106    pub fn method(mut self, method: Method) -> Self {
107        self.method = Some(method);
108        self
109    }
110
111    pub fn uri(mut self, uri: impl Into<String>) -> Self {
112        self.uri = Some(uri.into());
113        self
114    }
115
116    pub fn header(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
117        self.headers.insert(name.into(), value.into());
118        self
119    }
120
121    pub fn body(mut self, body: impl Into<Vec<u8>>) -> Self {
122        self.body = body.into();
123        self
124    }
125
126    pub fn build(self) -> Request {
127        Request {
128            method: self.method.unwrap_or(Method::Get),
129            uri: self.uri.unwrap_or_else(|| String::from("/")),
130            headers: self.headers,
131            body: self.body,
132        }
133    }
134}
135
136pub struct Response {
137    pub status: u16,
138    pub headers: HashMap<String, String>,
139    pub body: Vec<u8>,
140}
141
142impl Response {
143    pub fn builder() -> ResponseBuilder {
144        ResponseBuilder::new()
145    }
146
147    pub fn status(&self) -> u16 {
148        self.status
149    }
150
151    pub fn headers(&self) -> &HashMap<String, String> {
152        &self.headers
153    }
154
155    pub fn body(&self) -> &[u8] {
156        &self.body
157    }
158
159    pub fn into_body(self) -> Vec<u8> {
160        self.body
161    }
162}
163
164pub struct ResponseBuilder {
165    status: u16,
166    headers: HashMap<String, String>,
167    body: Vec<u8>,
168}
169
170impl ResponseBuilder {
171    pub fn new() -> Self {
172        Self {
173            status: 200,
174            headers: HashMap::new(),
175            body: Vec::new(),
176        }
177    }
178
179    pub fn status(mut self, status: u16) -> Self {
180        self.status = status;
181        self
182    }
183
184    pub fn header(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
185        self.headers.insert(name.into(), value.into());
186        self
187    }
188
189    pub fn body(mut self, body: impl Into<Vec<u8>>) -> Self {
190        self.body = body.into();
191        self
192    }
193
194    pub fn build(self) -> Response {
195        Response {
196            status: self.status,
197            headers: self.headers,
198            body: self.body,
199        }
200    }
201}
202
203// Internal executor functions copied from Spin SDK
204fn outgoing_body(body: OutgoingBody) -> impl Sink<Vec<u8>, Error = StreamError> {
205    struct Outgoing {
206        stream_and_body: Option<(OutputStream, OutgoingBody)>,
207        cancel_token: Option<CancelOnDropToken>,
208    }
209
210    impl Drop for Outgoing {
211        fn drop(&mut self) {
212            drop(self.cancel_token.take());
213
214            if let Some((stream, body)) = self.stream_and_body.take() {
215                drop(stream);
216                _ = OutgoingBody::finish(body, None);
217            }
218        }
219    }
220
221    let stream = body.write().expect("response body should be writable");
222    let outgoing = Rc::new(RefCell::new(Outgoing {
223        stream_and_body: Some((stream, body)),
224        cancel_token: None,
225    }));
226
227    sink::unfold((), {
228        move |(), chunk: Vec<u8>| {
229            future::poll_fn({
230                let mut offset = 0;
231                let mut flushing = false;
232                let outgoing = outgoing.clone();
233
234                move |context| {
235                    let mut outgoing = outgoing.borrow_mut();
236                    let (stream, _) = &outgoing.stream_and_body.as_ref().unwrap();
237                    loop {
238                        match stream.check_write() {
239                            Ok(0) => {
240                                outgoing.cancel_token = Some(CancelOnDropToken::from(
241                                    spin_executor::push_waker_and_get_token(
242                                        stream.subscribe(),
243                                        context.waker().clone(),
244                                    ),
245                                ));
246                                break Poll::Pending;
247                            }
248                            Ok(count) => {
249                                if offset == chunk.len() {
250                                    if flushing {
251                                        break Poll::Ready(Ok(()));
252                                    } else {
253                                        match stream.flush() {
254                                            Ok(()) => flushing = true,
255                                            Err(StreamError::Closed) => break Poll::Ready(Ok(())),
256                                            Err(e) => break Poll::Ready(Err(e)),
257                                        }
258                                    }
259                                } else {
260                                    let count =
261                                        usize::try_from(count).unwrap().min(chunk.len() - offset);
262
263                                    match stream.write(&chunk[offset..][..count]) {
264                                        Ok(()) => {
265                                            offset += count;
266                                        }
267                                        Err(e) => break Poll::Ready(Err(e)),
268                                    }
269                                }
270                            }
271                            // If the stream is closed but the entire chunk was
272                            // written then we've done all we could so this
273                            // chunk is now complete.
274                            Err(StreamError::Closed) if offset == chunk.len() => {
275                                break Poll::Ready(Ok(()))
276                            }
277                            Err(e) => break Poll::Ready(Err(e)),
278                        }
279                    }
280                }
281            })
282        }
283    })
284}
285
286fn outgoing_request_send(
287    request: OutgoingRequest,
288) -> impl Future<Output = Result<IncomingResponse, ErrorCode>> {
289    struct State {
290        response: Option<Result<FutureIncomingResponse, ErrorCode>>,
291        cancel_token: Option<CancelOnDropToken>,
292    }
293
294    impl Drop for State {
295        fn drop(&mut self) {
296            drop(self.cancel_token.take());
297            drop(self.response.take());
298        }
299    }
300
301    let response = outgoing_handler::handle(request, None);
302    let mut state = State {
303        response: Some(response),
304        cancel_token: None,
305    };
306    future::poll_fn({
307        move |context| match &state.response.as_ref().unwrap() {
308            Ok(response) => {
309                if let Some(response) = response.get() {
310                    Poll::Ready(response.unwrap())
311                } else {
312                    state.cancel_token = Some(CancelOnDropToken::from(
313                        spin_executor::push_waker_and_get_token(
314                            response.subscribe(),
315                            context.waker().clone(),
316                        ),
317                    ));
318                    Poll::Pending
319                }
320            }
321            Err(error) => Poll::Ready(Err(error.clone())),
322        }
323    })
324}
325
326fn incoming_body(
327    body: IncomingBody,
328) -> impl Stream<Item = Result<Vec<u8>, wasi::io::streams::Error>> {
329    struct Incoming {
330        stream_and_body: Option<(InputStream, IncomingBody)>,
331        cancel_token: Option<CancelOnDropToken>,
332    }
333
334    impl Drop for Incoming {
335        fn drop(&mut self) {
336            drop(self.cancel_token.take());
337
338            if let Some((stream, body)) = self.stream_and_body.take() {
339                drop(stream);
340                IncomingBody::finish(body);
341            }
342        }
343    }
344
345    stream::poll_fn({
346        let stream = body.stream().expect("response body should be readable");
347        let mut incoming = Incoming {
348            stream_and_body: Some((stream, body)),
349            cancel_token: None,
350        };
351
352        move |context| {
353            if let Some((stream, _)) = &incoming.stream_and_body {
354                match stream.read(READ_SIZE) {
355                    Ok(buffer) => {
356                        if buffer.is_empty() {
357                            incoming.cancel_token = Some(CancelOnDropToken::from(
358                                spin_executor::push_waker_and_get_token(
359                                    stream.subscribe(),
360                                    context.waker().clone(),
361                                ),
362                            ));
363                            Poll::Pending
364                        } else {
365                            Poll::Ready(Some(Ok(buffer)))
366                        }
367                    }
368                    Err(StreamError::Closed) => Poll::Ready(None),
369                    Err(StreamError::LastOperationFailed(error)) => Poll::Ready(Some(Err(error))),
370                }
371            } else {
372                Poll::Ready(None)
373            }
374        }
375    })
376}
377
378pub async fn send(request: Request) -> Result<Response> {
379    // Parse the URI to extract components
380    let uri = request.uri.parse::<http::Uri>()
381        .map_err(|e| anyhow::anyhow!("Invalid URI: {}", e))?;
382    
383    let scheme = uri.scheme()
384        .ok_or_else(|| anyhow::anyhow!("URI missing scheme"))?;
385    let authority = uri.authority()
386        .ok_or_else(|| anyhow::anyhow!("URI missing authority"))?;
387    let path_and_query = uri.path_and_query()
388        .map(|pq| pq.as_str())
389        .unwrap_or("/");
390
391    // Convert headers to WASI format
392    let headers_vec: Vec<(String, Vec<u8>)> = request.headers
393        .into_iter()
394        .map(|(k, v)| (k, v.into_bytes()))
395        .collect();
396    
397    let headers = Headers::from_list(&headers_vec)
398        .map_err(|e| anyhow::anyhow!("Failed to create headers: {:?}", e))?;
399    
400    // Create the outgoing request
401    let outgoing_request = OutgoingRequest::new(headers);
402    
403    outgoing_request
404        .set_method(&WasiMethod::from(request.method))
405        .map_err(|_| anyhow::anyhow!("Failed to set method"))?;
406    
407    outgoing_request
408        .set_scheme(Some(&match scheme.as_str() {
409            "https" => Scheme::Https,
410            "http" => Scheme::Http,
411            other => Scheme::Other(other.to_string()),
412        }))
413        .map_err(|_| anyhow::anyhow!("Failed to set scheme"))?;
414    
415    outgoing_request
416        .set_authority(Some(authority.as_str()))
417        .map_err(|_| anyhow::anyhow!("Failed to set authority"))?;
418    
419    outgoing_request
420        .set_path_with_query(Some(path_and_query))
421        .map_err(|_| anyhow::anyhow!("Failed to set path"))?;
422
423    // Send the request with body if present
424    let incoming_response = if !request.body.is_empty() {
425        let body_handle = outgoing_request.body()
426            .expect("request body should be available");
427        let mut body_sink = outgoing_body(body_handle);
428        let response_future = outgoing_request_send(outgoing_request);
429        body_sink.send(request.body).await
430            .map_err(|e| anyhow::anyhow!("Failed to send request body: {:?}", e))?;
431        drop(body_sink);
432        response_future.await
433            .map_err(|e| anyhow::anyhow!("HTTP request failed: {:?}", e))?
434    } else {
435        outgoing_request_send(outgoing_request).await
436            .map_err(|e| anyhow::anyhow!("HTTP request failed: {:?}", e))?
437    };
438
439    // Get response status and headers
440    let status = incoming_response.status();
441    let response_headers = incoming_response.headers();
442    
443    // Convert headers to HashMap
444    let mut headers = HashMap::new();
445    for (name, value) in response_headers.entries() {
446        headers.insert(
447            name.to_string(),
448            String::from_utf8_lossy(&value).to_string(),
449        );
450    }
451
452    // Read the response body using the streaming approach
453    let body_stream = incoming_response.consume()
454        .expect("response body should be available");
455    
456    let mut stream = incoming_body(body_stream);
457    let mut body = Vec::new();
458    while let Some(chunk) = stream.try_next().await? {
459        body.extend(chunk);
460    }
461
462    Ok(Response {
463        status,
464        headers,
465        body,
466    })
467}
wasmcp_wasi/http.rs

wasmcp_wasi/
http.rs
