tokio_quicker/backend/
client.rs

1use std::{
2    sync::Arc,
3    task::{ready, Poll},
4};
5
6use quiche::Connection;
7use tokio::{io::ReadBuf, net::UdpSocket};
8
9use crate::backend::{to_io_error, to_wire, IoHandler};
10
11use crate::error::Result;
12
13use super::timer::Timer;
14
15pub(crate) struct Inner {
16    pub io: Arc<UdpSocket>,
17    pub connection: Connection,
18    pub send_flush: bool,
19    pub send_end: usize,
20    pub send_pos: usize,
21    pub recv_buf: Vec<u8>,
22    pub send_buf: Vec<u8>,
23    pub timer: Timer,
24}
25
26impl IoHandler for Inner {
27    fn timer(&mut self) -> &mut Timer {
28        &mut self.timer
29    }
30
31    fn connection(&mut self) -> &mut Connection {
32        &mut self.connection
33    }
34
35    fn poll_send(&mut self, cx: &mut std::task::Context<'_>) -> Poll<Result<()>> {
36        if self.send_flush {
37            while self.send_pos != self.send_end {
38                let n = ready!(self.io.poll_send(cx, &self.send_buf[self.send_pos..]))?;
39                self.send_pos += n;
40            }
41
42            self.send_pos = 0;
43            self.send_end = 0;
44            self.send_flush = false;
45        }
46
47        match self.connection.send(&mut self.send_buf[self.send_end..]) {
48            Ok((n, _info)) => {
49                self.send_end += n;
50                self.send_flush = self.send_end == self.send_buf.len();
51            }
52            Err(quiche::Error::Done) if self.send_pos != self.send_end => (),
53            Err(quiche::Error::Done) => return Poll::Pending,
54            Err(quiche::Error::BufferTooShort) => {
55                self.send_flush = true;
56                return Poll::Ready(Ok(()));
57            }
58            Err(err) => {
59                self.connection
60                    .close(false, to_wire(err), b"fail")
61                    .map_err(to_io_error)?;
62                return Poll::Pending;
63            }
64        }
65
66        let n = ready!(self
67            .io
68            .poll_send(cx, &self.send_buf[self.send_pos..self.send_end]))?;
69        self.send_pos += n;
70
71        Poll::Ready(Ok(()))
72    }
73
74    fn poll_recv(&mut self, cx: &mut std::task::Context<'_>) -> Poll<Result<()>> {
75        let buf = &mut ReadBuf::new(&mut self.recv_buf);
76        let from = ready!(self.io.poll_recv_from(cx, buf))?;
77        let info = quiche::RecvInfo {
78            from,
79            to: self.io.local_addr()?,
80        };
81        match self.connection.recv(buf.filled_mut(), info) {
82            Ok(_) => Poll::Ready(Ok(())),
83            Err(quiche::Error::Done) => Poll::Ready(Ok(())),
84            Err(err) => {
85                self.connection
86                    .close(false, to_wire(err), b"fail")
87                    .map_err(to_io_error)?;
88                Poll::Pending
89            }
90        }
91    }
92}