Skip to main content

wasmtime_wasi_tls/p2/
host.rs

1use wasmtime::Result;
2use wasmtime::component::Resource;
3use wasmtime_wasi::async_trait;
4use wasmtime_wasi::p2::Pollable;
5use wasmtime_wasi::p2::{DynInputStream, DynOutputStream, DynPollable, IoError};
6
7use crate::p2::{
8    bindings,
9    io::{
10        AsyncReadStream, AsyncWriteStream, FutureOutput, WasiFuture, WasiStreamReader,
11        WasiStreamWriter,
12    },
13};
14use crate::{TlsStream, TlsTransport, WasiTlsCtxView};
15
16impl<'a> bindings::types::Host for WasiTlsCtxView<'a> {}
17
18/// Represents the ClientHandshake which will be used to configure the handshake
19pub struct HostClientHandshake {
20    server_name: String,
21    transport: Box<dyn TlsTransport>,
22}
23
24impl<'a> bindings::types::HostClientHandshake for WasiTlsCtxView<'a> {
25    fn new(
26        &mut self,
27        server_name: String,
28        input: Resource<DynInputStream>,
29        output: Resource<DynOutputStream>,
30    ) -> wasmtime::Result<Resource<HostClientHandshake>> {
31        let input = self.table.delete(input)?;
32        let output = self.table.delete(output)?;
33
34        let reader = WasiStreamReader::new(input);
35        let writer = WasiStreamWriter::new(output);
36        let transport = tokio::io::join(reader, writer);
37
38        Ok(self.table.push(HostClientHandshake {
39            server_name,
40            transport: Box::new(transport) as Box<dyn TlsTransport>,
41        })?)
42    }
43
44    fn finish(
45        &mut self,
46        this: Resource<HostClientHandshake>,
47    ) -> wasmtime::Result<Resource<HostFutureClientStreams>> {
48        let handshake = self.table.delete(this)?;
49
50        let connect = self
51            .ctx
52            .provider
53            .connect(handshake.server_name, handshake.transport);
54
55        let future = HostFutureClientStreams(WasiFuture::spawn(async move {
56            let tls_stream = connect.await?;
57
58            let (rx, tx) = tokio::io::split(tls_stream);
59            let write_stream = AsyncWriteStream::new(tx);
60            let client = HostClientConnection(write_stream.clone());
61
62            let input = Box::new(AsyncReadStream::new(rx)) as DynInputStream;
63            let output = Box::new(write_stream) as DynOutputStream;
64
65            Ok((client, input, output))
66        }));
67
68        Ok(self.table.push(future)?)
69    }
70
71    fn drop(&mut self, this: Resource<HostClientHandshake>) -> wasmtime::Result<()> {
72        self.table.delete(this)?;
73        Ok(())
74    }
75}
76
77/// Future streams provides the tls streams after the handshake is completed
78pub struct HostFutureClientStreams(
79    WasiFuture<Result<(HostClientConnection, DynInputStream, DynOutputStream), IoError>>,
80);
81
82#[async_trait]
83impl Pollable for HostFutureClientStreams {
84    async fn ready(&mut self) {
85        self.0.ready().await
86    }
87}
88
89impl<'a> bindings::types::HostFutureClientStreams for WasiTlsCtxView<'a> {
90    fn subscribe(
91        &mut self,
92        this: Resource<HostFutureClientStreams>,
93    ) -> wasmtime::Result<Resource<DynPollable>> {
94        wasmtime_wasi::p2::subscribe(self.table, this)
95    }
96
97    fn get(
98        &mut self,
99        this: Resource<HostFutureClientStreams>,
100    ) -> wasmtime::Result<
101        Option<
102            Result<
103                Result<
104                    (
105                        Resource<HostClientConnection>,
106                        Resource<DynInputStream>,
107                        Resource<DynOutputStream>,
108                    ),
109                    Resource<IoError>,
110                >,
111                (),
112            >,
113        >,
114    > {
115        let future = self.table.get_mut(&this)?;
116
117        let result = match future.0.get() {
118            FutureOutput::Ready(Ok((client, input, output))) => {
119                let client = self.table.push(client)?;
120                let input = self.table.push_child(input, &client)?;
121                let output = self.table.push_child(output, &client)?;
122
123                Some(Ok(Ok((client, input, output))))
124            }
125            FutureOutput::Ready(Err(io_error)) => {
126                let io_error = self.table.push(io_error)?;
127
128                Some(Ok(Err(io_error)))
129            }
130            FutureOutput::Consumed => Some(Err(())),
131            FutureOutput::Pending => None,
132        };
133
134        Ok(result)
135    }
136
137    fn drop(&mut self, this: Resource<HostFutureClientStreams>) -> wasmtime::Result<()> {
138        self.table.delete(this)?;
139        Ok(())
140    }
141}
142
143/// Represents the client connection and used to shut down the tls stream
144pub struct HostClientConnection(
145    crate::p2::io::AsyncWriteStream<tokio::io::WriteHalf<Box<dyn TlsStream>>>,
146);
147
148impl<'a> bindings::types::HostClientConnection for WasiTlsCtxView<'a> {
149    fn close_output(&mut self, this: Resource<HostClientConnection>) -> wasmtime::Result<()> {
150        self.table.get_mut(&this)?.0.close()
151    }
152
153    fn drop(&mut self, this: Resource<HostClientConnection>) -> wasmtime::Result<()> {
154        self.table.delete(this)?;
155        Ok(())
156    }
157}