This repository has no description
0

Configure Feed

Select the types of activity you want to include in your feed.

core / knot2 / crates / knot-sim / tests / common / mod.rs
9.7 kB 303 lines
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}