tls2/
tokio.rs

1use std::{
2    pin::Pin,
3    sync::Arc,
4    task::{Context, Poll, ready},
5};
6
7use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, ReadBuf};
8
9use crate::{
10    ALPN_PROTOCOL_MAX_SIZE, CertificateVerifier, CipherSuite, Client, ClientApplicationDataEvent, ClientConfig,
11    ClientHandshakeEvent, CryptoProvider, Error, KeyExchangeGroup, MAX_RECORD_SIZE, SignatureScheme,
12};
13
14/// Tokio-based TLS stream wrapping the sans-IO `Client`.
15pub struct TlsClient<S, C>
16where
17    S: AsyncRead + AsyncWrite + Unpin,
18    C: CryptoProvider + Unpin,
19{
20    stream: S,
21    client: Client<Vec<u8>, C>,
22    certificate_verifier: Arc<dyn CertificateVerifier>,
23    close_notify_sent: bool,
24}
25
26impl<S: AsyncRead + AsyncWrite + Unpin, P: CryptoProvider + Unpin> TlsClient<S, P> {
27    pub fn new(config: ClientConfig<P>, certificate_verifier: Arc<dyn CertificateVerifier>, stream: S) -> Self {
28        let client = Client::new(config, vec![0u8; MAX_RECORD_SIZE], vec![0u8; MAX_RECORD_SIZE]);
29
30        Self {
31            stream,
32            client,
33            certificate_verifier,
34            close_notify_sent: false,
35        }
36    }
37
38    /// Run the full TLS 1.3 handshake.
39    pub async fn handshake(
40        &mut self,
41        server_name: Option<&str>,
42        alpn_protocols: &[&[u8]],
43    ) -> Result<HandshakeData, Error> {
44        let mut event = self.client.start_handshake(server_name, alpn_protocols)?;
45        loop {
46            match event {
47                ClientHandshakeEvent::Send => {
48                    tokio::io::AsyncWriteExt::write_all(&mut self.stream, self.client.outgoing_data())
49                        .await
50                        .map_err(|_| Error::ConnectionClosed)?;
51                }
52                ClientHandshakeEvent::Receive => {
53                    let n = self
54                        .stream
55                        .read(self.client.receive_buffer())
56                        .await
57                        .map_err(|_| Error::ConnectionClosed)?;
58                    if n == 0 {
59                        return Err(Error::ConnectionClosed);
60                    }
61                    self.client.commit_received(n);
62                }
63                ClientHandshakeEvent::VerifyServerCertificate => {
64                    {
65                        let (cert, server_name) =
66                            self.client.server_certificate().ok_or(Error::CertificateParseFailed)?;
67                        self.certificate_verifier.verify_certificate(&cert, server_name).await?;
68                    }
69                    self.client.accept_certificate(Ok(()));
70                }
71                ClientHandshakeEvent::Done {
72                    ciphersuite,
73                    tls_version,
74                    key_exchange_group,
75                    signature_scheme,
76                    alpn,
77                } => {
78                    return Ok(HandshakeData {
79                        ciphersuite,
80                        tls_version,
81                        key_exchange_group,
82                        signature_scheme,
83                        alpn: alpn.try_into().unwrap(),
84                    });
85                }
86                ClientHandshakeEvent::Closed => return Err(Error::ConnectionClosed),
87            }
88            event = self.client.continue_handshake()?;
89        }
90    }
91}
92
93// ── AsyncRead ─────────────────────────────────────────────────────────────
94
95impl<S: AsyncRead + AsyncWrite + Unpin, C: CryptoProvider + Unpin> AsyncRead for TlsClient<S, C> {
96    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
97        let this = self.get_mut();
98        loop {
99            // 1. Flush any pending KeyUpdate response before reading more.
100            if !this.client.outgoing_key_update_data().is_empty() {
101                let resp = this.client.outgoing_key_update_data();
102                match Pin::new(&mut this.stream).poll_write(cx, resp) {
103                    Poll::Ready(Ok(n)) => {
104                        this.client.commit_key_update_data(n);
105                        if !this.client.outgoing_key_update_data().is_empty() {
106                            return Poll::Pending;
107                        }
108                    }
109                    Poll::Pending => return Poll::Pending,
110                    Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
111                }
112            }
113
114            // 2. Try to decrypt any data already in the receive buffer.
115            match this.client.decrypt() {
116                Ok(ClientApplicationDataEvent::AppData) => {
117                    let data = this.client.received_app_data();
118                    let n = data.len().min(buf.remaining());
119                    buf.put_slice(&data[..n]);
120                    this.client.commit_app_data(n);
121                    return Poll::Ready(Ok(()));
122                }
123                Ok(ClientApplicationDataEvent::Ticket {
124                    ..
125                }) => continue,
126                Ok(ClientApplicationDataEvent::KeyUpdate) => continue,
127                Ok(ClientApplicationDataEvent::None) => {}
128                Err(Error::ConnectionClosed) => return Poll::Ready(Ok(())),
129                Err(e) => return Poll::Ready(Err(std::io::Error::new(std::io::ErrorKind::Other, format!("{e:?}")))),
130            }
131
132            // 3. No complete record in buffer — read more from the network.
133            let recv_buf = this.client.receive_buffer();
134            let mut rb = ReadBuf::new(recv_buf);
135            match Pin::new(&mut this.stream).poll_read(cx, &mut rb) {
136                Poll::Ready(Ok(())) => {
137                    let n = rb.filled().len();
138                    if n == 0 {
139                        return Poll::Ready(Ok(()));
140                    }
141                    this.client.commit_received(n);
142                }
143                Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
144                Poll::Pending => return Poll::Pending,
145            }
146
147            // 4. Decrypt the newly read data.
148            match this.client.decrypt() {
149                Ok(ClientApplicationDataEvent::AppData) => {
150                    let data = this.client.received_app_data();
151                    let n = data.len().min(buf.remaining());
152                    buf.put_slice(&data[..n]);
153                    this.client.commit_app_data(n);
154                    return Poll::Ready(Ok(()));
155                }
156                Ok(ClientApplicationDataEvent::Ticket {
157                    ..
158                }) => continue,
159                Ok(ClientApplicationDataEvent::KeyUpdate) => continue,
160                Ok(ClientApplicationDataEvent::None) => continue,
161                Err(Error::ConnectionClosed) => return Poll::Ready(Ok(())),
162                Err(e) => return Poll::Ready(Err(std::io::Error::new(std::io::ErrorKind::Other, format!("{e:?}")))),
163            }
164        }
165    }
166}
167
168// ── AsyncWrite ────────────────────────────────────────────────────────────
169
170impl<S: AsyncRead + AsyncWrite + Unpin, C: CryptoProvider + Unpin> AsyncWrite for TlsClient<S, C> {
171    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<std::io::Result<usize>> {
172        let this = self.get_mut();
173
174        // 0. Flush any pending KeyUpdate response before encrypting new data.
175        if !this.client.outgoing_key_update_data().is_empty() {
176            let n = ready!(Pin::new(&mut this.stream).poll_write(cx, this.client.outgoing_key_update_data()))?;
177            this.client.commit_key_update_data(n);
178            if !this.client.outgoing_key_update_data().is_empty() {
179                return Poll::Pending;
180            }
181        }
182
183        // 1. Flush any buffered encrypted data first.
184        if !this.client.outgoing_data().is_empty() {
185            let n = ready!(Pin::new(&mut this.stream).poll_write(cx, this.client.outgoing_data()))?;
186            this.client.commit_sent(n);
187            if !this.client.outgoing_data().is_empty() {
188                return Poll::Pending;
189            }
190        }
191
192        // 2. Encrypt new plaintext.
193        let n = match this.client.encrypt(buf) {
194            Ok(n) => n,
195            Err(e) => return Poll::Ready(Err(std::io::Error::new(std::io::ErrorKind::Other, format!("{e:?}")))),
196        };
197
198        // 3. Try to send the encrypted record.
199        match Pin::new(&mut this.stream).poll_write(cx, this.client.outgoing_data()) {
200            Poll::Ready(Ok(m)) => {
201                this.client.commit_sent(m);
202                Poll::Ready(Ok(n))
203            }
204            Poll::Pending => Poll::Ready(Ok(n)),
205            Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
206        }
207    }
208
209    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
210        let this = self.get_mut();
211
212        // Flush any pending KeyUpdate response first.
213        if !this.client.outgoing_key_update_data().is_empty() {
214            let n = ready!(Pin::new(&mut this.stream).poll_write(cx, this.client.outgoing_key_update_data()))?;
215            this.client.commit_key_update_data(n);
216            if !this.client.outgoing_key_update_data().is_empty() {
217                return Poll::Pending;
218            }
219        }
220
221        // Then flush any buffered outgoing data
222        if !this.client.outgoing_data().is_empty() {
223            let n = ready!(Pin::new(&mut this.stream).poll_write(cx, this.client.outgoing_data()))?;
224            this.client.commit_sent(n);
225            if !this.client.outgoing_data().is_empty() {
226                return Poll::Pending;
227            }
228        }
229
230        Pin::new(&mut this.stream).poll_flush(cx)
231    }
232
233    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
234        let this = self.get_mut();
235
236        // 1. Flush any buffered data.
237        if !this.client.outgoing_data().is_empty() {
238            let n = ready!(Pin::new(&mut this.stream).poll_write(cx, this.client.outgoing_data()))?;
239            this.client.commit_sent(n);
240            if !this.client.outgoing_data().is_empty() {
241                return Poll::Pending;
242            }
243        }
244
245        // 2. Send close_notify (only once).
246        if !this.close_notify_sent {
247            this.close_notify_sent = true;
248            match this.client.close() {
249                Ok(data) => {
250                    let n = ready!(Pin::new(&mut this.stream).poll_write(cx, data))?;
251                    if n < data.len() {
252                        this.client.commit_sent(n);
253                        return Poll::Pending;
254                    }
255                }
256                Err(_) => {}
257            }
258        }
259
260        // 3. Flush any remaining close_notify bytes.
261        if !this.client.outgoing_data().is_empty() {
262            let n = ready!(Pin::new(&mut this.stream).poll_write(cx, this.client.outgoing_data()))?;
263            this.client.commit_sent(n);
264            if !this.client.outgoing_data().is_empty() {
265                return Poll::Pending;
266            }
267        }
268
269        Pin::new(&mut this.stream).poll_shutdown(cx)
270    }
271}
272
273/// The settings negotiated during the handshake
274#[derive(Clone, Debug)]
275pub struct HandshakeData {
276    ciphersuite: CipherSuite,
277    tls_version: u16,
278    key_exchange_group: KeyExchangeGroup,
279    signature_scheme: SignatureScheme,
280    alpn: heapless::Vec<u8, ALPN_PROTOCOL_MAX_SIZE>,
281}
282
283impl HandshakeData {
284    #[inline]
285    pub fn ciphersuite(&self) -> CipherSuite {
286        return self.ciphersuite;
287    }
288
289    /// Wire-encoded protocol version (`0x0304` for TLS 1.3).
290    #[inline]
291    pub fn tls_version(&self) -> u16 {
292        self.tls_version
293    }
294
295    #[inline]
296    pub fn key_exchange_group(&self) -> KeyExchangeGroup {
297        self.key_exchange_group
298    }
299
300    /// The signature scheme used by the server's CertificateVerify.
301    #[inline]
302    pub fn signature_scheme(&self) -> SignatureScheme {
303        self.signature_scheme
304    }
305
306    #[inline]
307    pub fn alpn(&self) -> &[u8] {
308        &self.alpn
309    }
310}
tls2/tokio.rs

tls2/
tokio.rs
