This repository has no description
1#![allow(dead_code)]
2
3use std::net::SocketAddr;
4use std::num::{NonZeroU32, NonZeroU64};
5use std::path::{Path, PathBuf};
6use std::sync::Arc;
7use std::time::Duration;
8
9use bytes::{Buf, Bytes};
10use http::{Method, StatusCode};
11use knot_edge::{
12 BodyInactivityTimeout, BurstSize, CertSource, EdgeConfig, EdgeGuards, HeaderTimeout,
13 IdleTimeout, ListenLimits, MaxInflightRequests, RequestTimeout, RequestsPerSecond,
14 RequiresFullHandshake, StaticCertPaths, TlsSetup, WriteRequestTimeout, ZeroRttRoutes,
15};
16use rustls::client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier};
17use rustls::crypto::aws_lc_rs;
18use rustls::pki_types::{CertificateDer, ServerName, UnixTime};
19use rustls::{DigitallySignedStruct, SignatureScheme};
20use tokio::task::JoinHandle;
21use tokio_util::sync::CancellationToken;
22
23pub fn nz32(value: u32) -> NonZeroU32 {
24 NonZeroU32::new(value).unwrap()
25}
26
27pub fn nz64(value: u64) -> NonZeroU64 {
28 NonZeroU64::new(value).unwrap()
29}
30
31pub fn pkt(payload: &[u8]) -> Vec<u8> {
32 assert!(
33 payload.len() + 4 <= 0xFFF0,
34 "pkt-line payload exceeds the 65516-byte maximum"
35 );
36 let mut out = format!("{:04x}", payload.len() + 4).into_bytes();
37 out.extend_from_slice(payload);
38 out
39}
40
41pub fn free_port() -> u16 {
42 std::net::TcpListener::bind("127.0.0.1:0")
43 .unwrap()
44 .local_addr()
45 .unwrap()
46 .port()
47}
48
49pub fn write_self_signed(dir: &Path) -> (PathBuf, PathBuf, Vec<u8>) {
50 let generated =
51 rcgen::generate_simple_self_signed(vec!["localhost".to_string(), "127.0.0.1".to_string()])
52 .unwrap();
53 let cert_path = dir.join("cert.pem");
54 let key_path = dir.join("key.pem");
55 std::fs::write(&cert_path, generated.cert.pem()).unwrap();
56 std::fs::write(&key_path, generated.signing_key.serialize_pem()).unwrap();
57 let der = generated.cert.der().as_ref().to_vec();
58 (cert_path, key_path, der)
59}
60
61pub fn edge_config(addr: SocketAddr, cert: PathBuf, key: PathBuf) -> EdgeConfig {
62 EdgeConfig {
63 http_addr: knot_edge::PublicBind::new(addr),
64 limits: ListenLimits::new(
65 HeaderTimeout::from_millis(nz64(30_000)),
66 IdleTimeout::from_millis(nz64(120_000)),
67 nz32(1024),
68 ),
69 guards: EdgeGuards::new(
70 RequestsPerSecond::new(nz32(1_000_000)),
71 BurstSize::new(nz32(1_000_000)),
72 MaxInflightRequests::new(nz32(10_000)),
73 RequestTimeout::from_millis(nz64(120_000)),
74 BodyInactivityTimeout::from_millis(nz64(120_000)),
75 WriteRequestTimeout::from_millis(nz64(1_800_000)),
76 knot_types::ProxyTrust::default(),
77 ),
78 tls: Some(TlsSetup {
79 source: CertSource::Static(StaticCertPaths {
80 cert_path: knot_edge::CertChainPath::new(cert),
81 key_path: knot_edge::PrivateKeyPath::new(key),
82 }),
83 http3: true,
84 internal: None,
85 }),
86 }
87}
88
89pub struct Edge {
90 pub addr: SocketAddr,
91 pub shutdown: CancellationToken,
92 pub task: JoinHandle<Result<(), knot_edge::EdgeError>>,
93 pub client: quinn::Endpoint,
94}
95
96async fn probe_identity(
97 endpoint: &quinn::Endpoint,
98 addr: SocketAddr,
99 expected_cert: &[u8],
100) -> Option<bool> {
101 let connecting = endpoint.connect(addr, "localhost").ok()?;
102 let connection = tokio::time::timeout(Duration::from_millis(250), connecting)
103 .await
104 .ok()?
105 .ok()?;
106 let ours = connection
107 .peer_identity()
108 .and_then(|identity| identity.downcast::<Vec<CertificateDer<'static>>>().ok())
109 .map(|certs| {
110 certs
111 .first()
112 .is_some_and(|cert| cert.as_ref() == expected_cert)
113 })
114 .unwrap_or(false);
115 connection.close(0u32.into(), b"probe done");
116 Some(ours)
117}
118
119async fn await_ready(
120 endpoint: &quinn::Endpoint,
121 addr: SocketAddr,
122 task: &mut JoinHandle<Result<(), knot_edge::EdgeError>>,
123 expected_cert: &[u8],
124) -> bool {
125 for _ in 0..200 {
126 if task.is_finished() {
127 return false;
128 }
129 if let Some(ours) = probe_identity(endpoint, addr, expected_cert).await {
130 return ours;
131 }
132 tokio::time::sleep(Duration::from_millis(20)).await;
133 }
134 false
135}
136
137pub async fn serve_edge(
138 certdir: &Path,
139 build: impl Fn() -> (RequiresFullHandshake, ZeroRttRoutes),
140) -> Edge {
141 for _ in 0..8 {
142 let addr: SocketAddr = format!("127.0.0.1:{}", free_port()).parse().unwrap();
143 let (cert, key, cert_der) = write_self_signed(certdir);
144 let (app, advertisement) = build();
145 let shutdown = CancellationToken::new();
146 let mut task = tokio::spawn(knot_edge::serve(
147 edge_config(addr, cert, key),
148 app,
149 advertisement,
150 shutdown.clone(),
151 ));
152 let client = h3_client();
153 if await_ready(&client, addr, &mut task, &cert_der).await {
154 return Edge {
155 addr,
156 shutdown,
157 task,
158 client,
159 };
160 }
161 client.close(0u32.into(), b"stand up retry");
162 shutdown.cancel();
163 let _ = task.await;
164 }
165 panic!("couldn't bind a free TCP+UDP port for the edge after several attempts");
166}
167
168#[derive(Debug)]
169struct AcceptAnyServerCert;
170
171impl ServerCertVerifier for AcceptAnyServerCert {
172 fn verify_server_cert(
173 &self,
174 _end_entity: &CertificateDer<'_>,
175 _intermediates: &[CertificateDer<'_>],
176 _server_name: &ServerName<'_>,
177 _ocsp_response: &[u8],
178 _now: UnixTime,
179 ) -> Result<ServerCertVerified, rustls::Error> {
180 Ok(ServerCertVerified::assertion())
181 }
182
183 fn verify_tls12_signature(
184 &self,
185 message: &[u8],
186 cert: &CertificateDer<'_>,
187 dss: &DigitallySignedStruct,
188 ) -> Result<HandshakeSignatureValid, rustls::Error> {
189 rustls::crypto::verify_tls12_signature(
190 message,
191 cert,
192 dss,
193 &aws_lc_rs::default_provider().signature_verification_algorithms,
194 )
195 }
196
197 fn verify_tls13_signature(
198 &self,
199 message: &[u8],
200 cert: &CertificateDer<'_>,
201 dss: &DigitallySignedStruct,
202 ) -> Result<HandshakeSignatureValid, rustls::Error> {
203 rustls::crypto::verify_tls13_signature(
204 message,
205 cert,
206 dss,
207 &aws_lc_rs::default_provider().signature_verification_algorithms,
208 )
209 }
210
211 fn supported_verify_schemes(&self) -> Vec<SignatureScheme> {
212 aws_lc_rs::default_provider()
213 .signature_verification_algorithms
214 .supported_schemes()
215 }
216}
217
218pub fn h3_client() -> quinn::Endpoint {
219 let mut crypto =
220 rustls::ClientConfig::builder_with_provider(Arc::new(aws_lc_rs::default_provider()))
221 .with_protocol_versions(&[&rustls::version::TLS13])
222 .unwrap()
223 .dangerous()
224 .with_custom_certificate_verifier(Arc::new(AcceptAnyServerCert))
225 .with_no_client_auth();
226 crypto.alpn_protocols = vec![b"h3".to_vec()];
227 crypto.resumption = rustls::client::Resumption::disabled();
228 let quic = quinn::crypto::rustls::QuicClientConfig::try_from(crypto).unwrap();
229 let mut endpoint = quinn::Endpoint::client("127.0.0.1:0".parse().unwrap()).unwrap();
230 endpoint.set_default_client_config(quinn::ClientConfig::new(Arc::new(quic)));
231 endpoint
232}
233
234pub async fn drain(
235 stream: &mut h3::client::RequestStream<h3_quinn::BidiStream<Bytes>, Bytes>,
236) -> Vec<u8> {
237 let mut out = Vec::new();
238 while let Some(mut chunk) = stream.recv_data().await.unwrap() {
239 out.extend_from_slice(&chunk.copy_to_bytes(chunk.remaining()));
240 }
241 out
242}
243
244pub async fn finish_request(
245 stream: &mut h3::client::RequestStream<h3_quinn::BidiStream<Bytes>, Bytes>,
246) {
247 match stream.finish().await {
248 Ok(()) => (),
249 Err(h3::error::StreamError::RemoteTerminate { code, .. })
250 if code == h3::error::Code::H3_NO_ERROR => {}
251 Err(error) => panic!("finishing the request stream failed: {error}"),
252 }
253}
254
255pub async fn h3_request(
256 edge: &Edge,
257 method: Method,
258 uri: String,
259 headers: &[(&str, &str)],
260 body: Option<Bytes>,
261 warmup: Option<&str>,
262) -> (StatusCode, Vec<u8>) {
263 let connection = edge
264 .client
265 .connect(edge.addr, "localhost")
266 .unwrap()
267 .await
268 .unwrap();
269 let quic = connection.clone();
270 let (mut driver, mut sender) = h3::client::new(h3_quinn::Connection::new(connection))
271 .await
272 .unwrap();
273 let drive = tokio::spawn(async move {
274 let _ = std::future::poll_fn(|cx| driver.poll_close(cx)).await;
275 });
276
277 if let Some(warmup) = warmup {
278 let request = http::Request::get(warmup).body(()).unwrap();
279 let mut stream = sender.send_request(request).await.unwrap();
280 finish_request(&mut stream).await;
281 let _ = stream.recv_response().await.unwrap();
282 drain(&mut stream).await;
283 }
284
285 let request = headers
286 .iter()
287 .fold(
288 http::Request::builder().method(method).uri(uri),
289 |builder, (name, value)| builder.header(*name, *value),
290 )
291 .body(())
292 .unwrap();
293 let mut stream = sender.send_request(request).await.unwrap();
294 if let Some(body) = body {
295 stream.send_data(body).await.unwrap();
296 }
297 finish_request(&mut stream).await;
298 let status = stream.recv_response().await.unwrap().status();
299 let out = drain(&mut stream).await;
300 quic.close(0u32.into(), b"done");
301 drive.abort();
302 (status, out)
303}