1use core::net::Ipv6Addr;
2
3use std::process::ExitStatus;
4
5use anyhow::Context;
6use rcgen::{generate_simple_self_signed, CertifiedKey};
7use rustls::pki_types::{CertificateDer, PrivatePkcs8KeyDer};
8use rustls::version::TLS13;
9use rustls::{ClientConfig, RootCertStore, ServerConfig};
10use tokio::net::TcpListener;
11use tokio::process::Command;
12use tokio::sync::oneshot;
13use tokio::task::JoinHandle;
14use tokio::{select, spawn};
15
16pub async fn free_port() -> anyhow::Result<u16> {
17 TcpListener::bind((Ipv6Addr::LOCALHOST, 0))
18 .await
19 .context("failed to start TCP listener")?
20 .local_addr()
21 .context("failed to query listener local address")
22 .map(|v| v.port())
23}
24
25pub async fn spawn_server(
26 cmd: &mut Command,
27) -> anyhow::Result<(JoinHandle<anyhow::Result<ExitStatus>>, oneshot::Sender<()>)> {
28 let mut child = cmd
29 .kill_on_drop(true)
30 .spawn()
31 .context("failed to spawn child")?;
32 let (stop_tx, stop_rx) = oneshot::channel();
33 let child = spawn(async move {
34 select!(
35 res = stop_rx => {
36 res.context("failed to wait for shutdown")?;
37 child.kill().await.context("failed to kill child")?;
38 child.wait().await
39 }
40 status = child.wait() => {
41 status
42 }
43 )
44 .context("failed to wait for child")
45 });
46 Ok((child, stop_tx))
47}
48
49pub fn cert_pair() -> anyhow::Result<(rustls::ServerConfig, rustls::ClientConfig)> {
50 let CertifiedKey {
51 cert: srv_crt,
52 signing_key: srv_key,
53 } = generate_simple_self_signed([
54 "127.0.0.1".to_string(),
55 "::1".to_string(),
56 "localhost".to_string(),
57 ])
58 .context("failed to generate server certificate")?;
59 let CertifiedKey {
60 cert: clt_crt,
61 signing_key: clt_key,
62 } = generate_simple_self_signed(["client.wrpc".to_string()])
63 .context("failed to generate client certificate")?;
64 let srv_crt = CertificateDer::from(srv_crt);
65
66 let mut ca = RootCertStore::empty();
67 ca.add(srv_crt.clone())?;
68 let clt_cnf = ClientConfig::builder_with_protocol_versions(&[&TLS13])
69 .with_root_certificates(ca)
70 .with_client_auth_cert(
71 vec![clt_crt.into()],
72 PrivatePkcs8KeyDer::from(clt_key.serialize_der()).into(),
73 )
74 .context("failed to create client config")?;
75 let srv_cnf = ServerConfig::builder_with_protocol_versions(&[&TLS13])
76 .with_no_client_auth() .with_single_cert(
78 vec![srv_crt],
79 PrivatePkcs8KeyDer::from(srv_key.serialize_der()).into(),
80 )
81 .context("failed to create server config")?;
82 Ok((srv_cnf, clt_cnf))
83}
84
85#[cfg(feature = "nats")]
86pub async fn start_nats() -> anyhow::Result<(
87 u16,
88 async_nats::Client,
89 JoinHandle<anyhow::Result<ExitStatus>>,
90 oneshot::Sender<()>,
91)> {
92 let nats_server_check = Command::new("nats-server").arg("--version").output().await;
94 if let Err(e) = nats_server_check {
95 let error_msg = if e.kind() == std::io::ErrorKind::NotFound {
96 "nats-server is not installed or not in PATH"
97 } else if e.kind() == std::io::ErrorKind::PermissionDenied {
98 "nats-server is not executable or permission denied"
99 } else {
100 "failed to execute nats-server"
101 };
102 anyhow::bail!(
103 "{}. Please install nats-server >= 2.10.20. \
104 See https://docs.nats.io/running-a-nats-service/introduction/installation for installation instructions. \
105 Original error: {}",
106 error_msg,
107 e
108 );
109 }
110
111 let port = free_port().await?;
112 let (server, stop_tx) =
113 spawn_server(Command::new("nats-server").args(["-T=false", "-p", &port.to_string()]))
114 .await
115 .context("failed to start NATS.io server")?;
116
117 let client = wrpc_cli::nats::connect(format!("nats://localhost:{port}"))
118 .await
119 .context("failed to connect to NATS.io server")?;
120 Ok((port, client, server, stop_tx))
121}
122
123#[cfg(feature = "nats")]
124pub async fn with_nats<T, Fut>(f: impl FnOnce(u16, async_nats::Client) -> Fut) -> anyhow::Result<T>
125where
126 Fut: core::future::Future<Output = anyhow::Result<T>>,
127{
128 let (port, nats_client, nats_server, stop_tx) = start_nats()
129 .await
130 .context("failed to start NATS.io server")?;
131 let res = f(port, nats_client).await.context("closure failed")?;
132 stop_tx.send(()).expect("failed to stop NATS.io server");
133 nats_server
134 .await
135 .context("failed to await NATS.io server stop")?
136 .context("NATS.io server failed to stop")?;
137 Ok(res)
138}
139
140#[cfg(feature = "quic")]
141pub async fn with_quic_endpoints<T, Fut>(
142 f: impl FnOnce(core::net::SocketAddr, quinn::Endpoint, quinn::Endpoint) -> Fut,
143) -> anyhow::Result<T>
144where
145 Fut: core::future::Future<Output = anyhow::Result<T>>,
146{
147 use std::sync::Arc;
148
149 use quinn::crypto::rustls::{QuicClientConfig, QuicServerConfig};
150 use quinn::{ClientConfig, ServerConfig};
151
152 let mut clt_ep = quinn::Endpoint::client((Ipv6Addr::LOCALHOST, 0).into())
153 .context("failed to create client endpoint")?;
154
155 let (srv_cnf, clt_cnf) = cert_pair().context("failed to generate certificates")?;
156
157 let clt_cnf: QuicClientConfig = clt_cnf
158 .try_into()
159 .context("failed to convert rustls client config to QUIC client config")?;
160 let srv_cnf: QuicServerConfig = srv_cnf
161 .try_into()
162 .context("failed to convert rustls server config to QUIC server config")?;
163
164 clt_ep.set_default_client_config(ClientConfig::new(Arc::new(clt_cnf)));
165 let srv_ep = quinn::Endpoint::server(
166 ServerConfig::with_crypto(Arc::new(srv_cnf)),
167 (Ipv6Addr::LOCALHOST, 0).into(),
168 )
169 .context("failed to create server endpoint")?;
170 let srv_addr = srv_ep
171 .local_addr()
172 .context("failed to query server address")?;
173
174 f(srv_addr, clt_ep, srv_ep).await.context("closure failed")
175}
176
177#[cfg(feature = "quic")]
178pub async fn with_quic<T, Fut>(
179 f: impl FnOnce(quinn::Connection, quinn::Connection) -> Fut,
180) -> anyhow::Result<T>
181where
182 Fut: core::future::Future<Output = anyhow::Result<T>>,
183{
184 with_quic_endpoints(|addr, clt, srv| async move {
185 let (clt, srv) = tokio::try_join!(
186 async move {
187 let conn = clt
188 .connect(addr, "::1")
189 .context("failed to connect to server")?;
190 conn.await.context("failed to establish client connection")
191 },
192 async move {
193 let conn = srv.accept().await.context("failed to accept connection")?;
194 conn.await.context("failed to establish server connection")
195 }
196 )?;
197 f(clt, srv).await.context("closure failed")
198 })
199 .await
200}
201
202#[cfg(feature = "web-transport")]
203pub async fn with_web_transport<T, Fut>(
204 f: impl FnOnce(wtransport::Connection, wtransport::Connection) -> Fut,
205) -> anyhow::Result<T>
206where
207 Fut: core::future::Future<Output = anyhow::Result<T>>,
208{
209 use wtransport::Endpoint;
210
211 let (srv_cnf, clt_cnf) = cert_pair().context("failed to generate certificates")?;
212
213 let srv = Endpoint::server(
214 wtransport::ServerConfig::builder()
215 .with_bind_default(0)
216 .with_custom_tls(srv_cnf)
217 .build(),
218 )
219 .context("failed to build server endpoint")?;
220 let clt = Endpoint::client(
221 wtransport::ClientConfig::builder()
222 .with_bind_default()
223 .with_custom_tls(clt_cnf)
224 .build(),
225 )
226 .context("failed to create client endpoint")?;
227 let addr = srv.local_addr().context("failed to query server address")?;
228 let (clt, srv) = tokio::try_join!(
229 async move {
230 clt.connect(format!("https://localhost:{}", addr.port()))
231 .await
232 .context("failed to connect to server")
233 },
234 async move {
235 let req = srv
236 .accept()
237 .await
238 .await
239 .context("failed to receive session request")?;
240 req.accept()
241 .await
242 .context("failed to accept client connection")
243 }
244 )?;
245 f(clt, srv).await.context("closure failed")
246}