quic_socket/
lib.rs

1use anyhow::{anyhow, Result};
2use async_trait::async_trait;
3use futures_util::StreamExt;
4use quinn::{Connection, Endpoint, IncomingBiStreams, ServerConfig};
5use std::{error::Error, fs, net::SocketAddr, net::ToSocketAddrs, sync::Arc};
6use url::Url;
7
8#[async_trait]
9pub trait QuicSocket {
10    async fn new(
11        addr: Option<SocketAddr>,
12        remote_url: Option<String>,
13        host: Option<String>,
14    ) -> Self;
15    async fn send(&mut self, payload: Vec<u8>) -> Result<()>;
16    async fn recv(&mut self, buf: &mut [u8]) -> Result<usize>;
17}
18
19pub struct QuicServer {
20    pub endpoint: Endpoint,
21    pub connection: Connection,
22    pub bi_streams: IncomingBiStreams,
23}
24
25pub struct QuicClient {
26    pub endpoint: Endpoint,
27    pub connection: Connection,
28    pub bi_streams: IncomingBiStreams,
29}
30
31#[async_trait]
32impl QuicSocket for QuicServer {
33    async fn new(
34        addr: Option<SocketAddr>,
35        _remote_url: Option<String>,
36        _host: Option<String>,
37    ) -> QuicServer {
38        let server_config = configure_server().unwrap();
39        let (endpoint, mut incoming) =
40            quinn::Endpoint::server(server_config, addr.unwrap()).unwrap();
41        let new_conn = incoming.next().await.unwrap().await.unwrap();
42        let quinn::NewConnection {
43            connection: conn,
44            bi_streams,
45            ..
46        } = new_conn;
47        QuicServer {
48            endpoint,
49            connection: conn,
50            bi_streams,
51        }
52    }
53    async fn send(&mut self, payload: Vec<u8>) -> Result<()> {
54        let (mut send, _) = self
55            .connection
56            .open_bi()
57            .await
58            .map_err(|e| anyhow!("failed to open stream: {}", e))
59            .unwrap();
60        send.write_all(&payload)
61            .await
62            .map_err(|e| anyhow!("failed to send request: {}", e))?;
63        send.finish()
64            .await
65            .map_err(|e| anyhow!("failed to shutdown stream: {}", e))?;
66        Ok(())
67    }
68
69    async fn recv(&mut self, buf: &mut [u8]) -> Result<usize> {
70        let (_, mut recv) = self.bi_streams.next().await.unwrap().unwrap();
71        let len = recv
72            .read(buf)
73            .await
74            .map_err(|e| anyhow!("failed to read response: {}", e))?;
75        Ok(len.unwrap())
76    }
77}
78
79#[async_trait]
80impl QuicSocket for QuicClient {
81    async fn new(
82        _addr: Option<SocketAddr>,
83        remote_url: Option<String>,
84        host: Option<String>,
85    ) -> Self {
86        let ca = "cert.der".to_string();
87        let mut roots = rustls::RootCertStore::empty();
88        roots
89            .add(&rustls::Certificate(fs::read(&ca).unwrap()))
90            .unwrap();
91        let mut client_crypto = rustls::ClientConfig::builder()
92            .with_safe_defaults()
93            .with_root_certificates(roots)
94            .with_no_client_auth();
95        client_crypto.alpn_protocols = ALPN_QUIC_HTTP.iter().map(|&x| x.into()).collect();
96        let mut endpoint = quinn::Endpoint::client("[::]:0".parse().unwrap()).unwrap();
97        endpoint.set_default_client_config(quinn::ClientConfig::new(Arc::new(client_crypto)));
98        let remote_url = Url::parse(&remote_url.unwrap()).unwrap();
99        let remote = (
100            remote_url.host_str().unwrap(),
101            remote_url.port().unwrap_or(4433),
102        )
103            .to_socket_addrs()
104            .unwrap()
105            .next()
106            .ok_or_else(|| anyhow!("couldn't resolve to an address"))
107            .unwrap();
108
109        let new_conn = endpoint
110            .connect(remote, &host.unwrap())
111            .unwrap()
112            .await
113            .map_err(|e| anyhow!("failed to connect: {}", e))
114            .unwrap();
115        let quinn::NewConnection {
116            connection: conn,
117            bi_streams,
118            ..
119        } = new_conn;
120        QuicClient {
121            endpoint,
122            connection: conn,
123            bi_streams,
124        }
125    }
126    async fn send(&mut self, payload: Vec<u8>) -> Result<()> {
127        let (mut send, _) = self
128            .connection
129            .open_bi()
130            .await
131            .map_err(|e| anyhow!("failed to open stream: {}", e))
132            .unwrap();
133        send.write_all(&payload)
134            .await
135            .map_err(|e| anyhow!("failed to send request: {}", e))?;
136        send.finish()
137            .await
138            .map_err(|e| anyhow!("failed to shutdown stream: {}", e))?;
139        Ok(())
140    }
141    async fn recv(&mut self, buf: &mut [u8]) -> Result<usize> {
142        let (_, mut recv) = self.bi_streams.next().await.unwrap().unwrap();
143        let len = recv
144            .read(buf)
145            .await
146            .map_err(|e| anyhow!("failed to read response: {}", e))?;
147        Ok(len.unwrap())
148    }
149}
150
151#[allow(clippy::field_reassign_with_default)] // https://github.com/rust-lang/rust-clippy/issues/6527
152fn configure_server() -> Result<ServerConfig, Box<dyn Error>> {
153    let cert_chain = fs::read("./cert.der")?;
154    let key = fs::read("./key.der")?;
155    let priv_key = rustls::PrivateKey(key);
156    let cert = vec![rustls::Certificate(cert_chain.clone())];
157    let mut server_crypto = rustls::ServerConfig::builder()
158        .with_safe_defaults()
159        .with_no_client_auth()
160        .with_single_cert(cert, priv_key)?;
161    server_crypto.alpn_protocols = ALPN_QUIC_HTTP.iter().map(|&x| x.into()).collect();
162    let mut server_config = quinn::ServerConfig::with_crypto(Arc::new(server_crypto));
163    Arc::get_mut(&mut server_config.transport)
164        .unwrap()
165        .max_concurrent_uni_streams(0_u8.into());
166    Ok(server_config)
167}
168
169#[allow(unused)]
170#[allow(clippy::field_reassign_with_default)] // https://github.com/rust-lang/rust-clippy/issues/6527
171pub fn gen_certificates() -> Result<(), Box<dyn Error>> {
172    let cert = rcgen::generate_simple_self_signed(vec!["localhost".into()]).unwrap();
173    let cert_der = cert.serialize_der().unwrap();
174    fs::write("./cert.der".to_string(), &cert_der).unwrap();
175    let priv_key = cert.serialize_private_key_der();
176    fs::write("./key.der".to_string(), &priv_key).unwrap();
177    let key = rustls::PrivateKey(priv_key);
178    let cert = vec![rustls::Certificate(cert_der.clone())];
179    let mut server_crypto = rustls::ServerConfig::builder()
180        .with_safe_defaults()
181        .with_no_client_auth()
182        .with_single_cert(cert, key)?;
183    server_crypto.alpn_protocols = ALPN_QUIC_HTTP.iter().map(|&x| x.into()).collect();
184    Ok(())
185}
186
187#[allow(unused)]
188pub const ALPN_QUIC_HTTP: &[&[u8]] = &[b"hq-29"];