Fixes

src/client.rs

@@ -3,11 +3,15 @@ use crate::{
 	record::{Direction, Records},
 };
 
-use std::{net::ToSocketAddrs, sync::Arc};
+use futures_util::StreamExt;
+use std::{
+	net::ToSocketAddrs,
+	sync::{Arc, atomic::AtomicU32},
+};
 use tokio::{
-	io::{AsyncReadExt, AsyncWriteExt},
+	io::AsyncWriteExt,
 	net::TcpStream,
-	sync::oneshot,
+	sync::{Semaphore, oneshot},
 };
 use tokio_rustls::{
 	TlsConnector,
@@ -17,6 +21,7 @@ use tokio_rustls::{
 		pki_types::ServerName,
 	},
 };
+use tokio_util::codec::Framed;
 
 #[derive(Debug)]
 struct DummyCertVerifier;
@@ -82,6 +87,11 @@ pub async fn play(
 	repeat: u32,
 ) {
 	sync_receiver.await.unwrap();
+	// Semaphore used to limit the number of concurrent clients.
+	// Its handle is released when the task panics.
+	let limiter = Arc::new(Semaphore::new(32));
+	let counter = Arc::new(AtomicU32::new(0));
+	let total = records.len() * repeat as usize;
 	let mut handles = Vec::new();
 	let connect_to = connect_to.to_socket_addrs().unwrap().next().unwrap();
 	match tls_mode {
@@ -92,58 +102,159 @@ pub async fn play(
 					.with_custom_certificate_verifier(Arc::new(DummyCertVerifier))
 					.with_no_client_auth(),
 			);
-			for (_id, (server_name, records)) in records.iter() {
+			for (id, (server_name, records)) in records.iter() {
 				let connector = TlsConnector::from(config.clone());
+				let counter = counter.clone();
+				let limiter = limiter.clone();
 				handles.push(tokio::spawn(async move {
+					let limiter = limiter.acquire().await.unwrap();
 					let server_name =
 						ServerName::try_from(String::from_utf8(server_name.clone()).unwrap())
 							.unwrap();
 					for _i in 0..repeat {
 						let stream = TcpStream::connect(connect_to).await.unwrap();
-						let mut stream = connector
+						let stream = connector
 							.connect(server_name.clone(), stream)
 							.await
 							.unwrap();
+						let mut stream = Framed::new(stream, crate::http::HttpCodec {});
 						for (direction, data) in records {
 							match direction {
 								Direction::ClientToServer => {
-									stream.write_all(data).await.unwrap();
+									println!("[CLT] ({id}) >> {}", data.len());
+									stream.get_mut().write_all(data).await.unwrap();
 								}
 								Direction::ServerToClient => {
-									let mut buf = Vec::new();
-									stream.read_buf(&mut buf).await.ok();
+									println!("[CLT] ({id}) << {}", data.len());
+									// let mut buf = Vec::new();
+									// stream.read_buf(&mut buf).await.ok();
+									//let mut buf = vec![0; data.len().saturating_sub(50).max(1)];
+									//let resp = stream.next().await.unwrap().unwrap();
+									let resp = tokio::time::timeout(
+										std::time::Duration::from_millis(500),
+										stream.next(),
+									)
+									.await
+									.unwrap()
+									.unwrap()
+									.unwrap();
+									dbg!(resp.len());
+									//crate::http::decode_http(&mut buf, &mut stream).await;
 								}
 							}
 						}
-						stream.shutdown().await.unwrap();
+						stream.get_mut().shutdown().await.unwrap();
+						let cnt = counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+						println!("Client: {} / {}", cnt + 1, total);
 					}
+					drop(limiter);
 				}));
+				//tokio::time::sleep(std::time::Duration::from_millis(500)).await;
 			}
 		}
 		TlsMode::None | TlsMode::Server => {
-			for (_id, (_server_name, records)) in records.iter() {
+			for (id, (_server_name, records)) in records.iter() {
+				/*if *id != 33 {
+					continue
+				}*/
+				let counter = counter.clone();
+				let limiter = limiter.clone();
 				handles.push(tokio::spawn(async move {
+					dbg!(limiter.available_permits());
+					let limiter = limiter.acquire().await.unwrap();
+					//let mut buf = Vec::new();
 					for _i in 0..repeat {
-						let mut stream = TcpStream::connect(connect_to).await.unwrap();
+						dbg!();
+						let stream = TcpStream::connect(connect_to).await.unwrap();
+						let mut stream = Framed::new(stream, crate::http::HttpCodec {});
+						/*let mut skip_recv = false;
 						for (direction, data) in records {
 							match direction {
 								Direction::ClientToServer => {
+									skip_recv = false;
+									println!("[CLT] ({id}) >> {}", data.len());
 									stream.write_all(data).await.unwrap();
 								}
 								Direction::ServerToClient => {
-									let mut buf = Vec::new();
-									stream.read_buf(&mut buf).await.ok();
+									if skip_recv {
+										continue;
+									}
+									println!("[CLT] ({id}) << {}", data.len());
+									//let mut buf = Vec::new();
+									//stream.read_buf(&mut buf).await.ok();
+									//let mut buf = vec![0; data.len().saturating_sub(50).max(1)];
+									let mut buf = vec![0; data.len()];
+									match tokio::time::timeout(
+										std::time::Duration::from_millis(500),
+										stream.readable(),
+									)
+									.await
+									{
+										Ok(r) => {
+											r.unwrap();
+										}
+										Err(_) => {
+											println!("[CLT] timeout recv ({id})");
+											break;
+										}
+									}
+									// TODO utiliser crate::http ici
+									match tokio::time::timeout(
+										std::time::Duration::from_millis(500),
+										stream.read_exact(&mut buf),
+									)
+									.await
+									{
+										Ok(r) => {
+											r.unwrap();
+										}
+										Err(_) => {
+											println!("[CLT] skip recv ({id})");
+											skip_recv = true;
+										}
+									}
+								}
+							}
+						}*/
+						for (direction, data) in records {
+							match direction {
+								Direction::ClientToServer => {
+									println!("[CLT] ({id}) >> {}", data.len());
+									stream.get_mut().write_all(data).await.unwrap();
+								}
+								Direction::ServerToClient => {
+									println!("[CLT] ({id}) << {}", data.len());
+									//let mut buf = Vec::new();
+									//stream.read_buf(&mut buf).await.ok();
+									//let mut buf = vec![0; data.len().saturating_sub(50).max(1)];
+									let resp = tokio::time::timeout(
+										std::time::Duration::from_millis(500),
+										stream.next(),
+									)
+									.await
+									.unwrap()
+									.unwrap()
+									.unwrap();
+									//let resp = stream.next().await.unwrap().unwrap();
+									dbg!(resp.len());
+									//crate::http::decode_http(&mut buf, &mut stream).await;
+									//buf.clear();
 								}
 							}
 						}
-						stream.shutdown().await.unwrap();
+						dbg!();
+						stream.get_mut().shutdown().await.unwrap();
+						let cnt = counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+						println!("Client: {} / {}", cnt + 1, total);
 					}
+					drop(limiter);
 				}));
+				//tokio::time::sleep(std::time::Duration::from_millis(500)).await;
 			}
 		}
 	}
 	for handle in handles {
 		handle.await.unwrap();
 	}
-	//std::process::exit(0);
+	std::process::exit(0);
 }

src/http.rs

@@ -0,0 +1,120 @@
+use regex::bytes::Regex;
+use std::sync::LazyLock;
+use tokio::io::AsyncReadExt;
+use tokio_util::codec::{Decoder, Encoder};
+
+static REGEX_CONTENT_LENGTH: LazyLock<Regex> =
+	LazyLock::new(|| Regex::new(r#"[cC]ontent-[lL]ength: *(\d+)\r\n"#).unwrap());
+
+pub async fn _decode_http<R: AsyncReadExt + Unpin>(buf: &mut Vec<u8>, stream: &mut R) {
+	dbg!();
+	loop {
+		dbg!();
+		if let Some(mut end_index) = memchr::memmem::find(buf, b"\r\n\r\n") {
+			end_index += 4;
+			if let Some(captures) = REGEX_CONTENT_LENGTH.captures(buf) {
+				if let Some(content_length) = captures.get(1) {
+					// Read body
+					let content_length: usize = str::from_utf8(content_length.as_bytes())
+						.unwrap()
+						.parse()
+						.unwrap();
+					while buf.len() < end_index + content_length {
+						dbg!();
+						match tokio::time::timeout(
+							std::time::Duration::from_millis(500),
+							stream.read_buf(buf),
+						)
+						.await
+						{
+							Ok(Ok(_n)) => {}
+							Ok(Err(e)) => {
+								println!("[http] error reading: {e:?}");
+								break;
+							}
+							Err(_e) => {
+								// timeout
+								break;
+							}
+						}
+					}
+					break;
+				} else {
+					// Erroneous Content-Type
+					break;
+				}
+			} else {
+				// Header ended without Content-Type => no body
+				break;
+			}
+		}
+		match tokio::time::timeout(std::time::Duration::from_millis(500), stream.read_buf(buf))
+			.await
+		{
+			Ok(Ok(n)) => {
+				println!("[http] read {n}");
+			}
+			Ok(Err(e)) => {
+				println!("[http] error reading: {e:?}");
+				break;
+			}
+			Err(_e) => {
+				// timeout
+				break;
+			}
+		}
+	}
+}
+
+pub struct HttpCodec {}
+
+impl Decoder for HttpCodec {
+	type Item = Vec<u8>;
+	type Error = std::io::Error;
+	fn decode(
+		&mut self,
+		src: &mut tokio_util::bytes::BytesMut,
+	) -> Result<Option<Self::Item>, Self::Error> {
+		if let Some(mut end_index) = memchr::memmem::find(src, b"\r\n\r\n") {
+			end_index += 4;
+			if let Some(captures) = REGEX_CONTENT_LENGTH.captures(src) {
+				if let Some(content_length) = captures.get(1) {
+					// Read body
+					let content_length: usize = str::from_utf8(content_length.as_bytes())
+						.unwrap()
+						.parse()
+						.unwrap();
+					if src.len() >= end_index + content_length {
+						//dbg!(content_length);
+						let out = src.to_vec();
+						src.clear();
+						Ok(Some(out))
+					} else {
+						Ok(None)
+					}
+				} else {
+					// Invalid Content-Length
+					Err(std::io::ErrorKind::InvalidData.into())
+				}
+			} else {
+				// Header ended without Content-Type => no body
+				let out = src.to_vec();
+				src.clear();
+				Ok(Some(out))
+			}
+		} else {
+			Ok(None)
+		}
+	}
+}
+
+impl Encoder<Vec<u8>> for HttpCodec {
+	type Error = std::io::Error;
+	fn encode(
+		&mut self,
+		_item: Vec<u8>,
+		_dst: &mut tokio_util::bytes::BytesMut,
+	) -> Result<(), Self::Error> {
+		Ok(())
+	}
+}

src/main.rs

@@ -1,6 +1,7 @@
 #![feature(let_chains)]
 
 mod client;
+mod http;
 mod record;
 mod server;
 
@@ -53,6 +54,9 @@ struct OptPlay {
 	/// Repeat N times
 	#[argp(option, short = 'r', default = "1")]
 	repeat: u32,
+	/// Only play this record
+	#[argp(option)]
+	record: Option<u64>,
 }
 
 /// Print records
@@ -93,7 +97,13 @@ async fn main() {
 				_ => panic!("TLS mode must be one of none,client,server,both."),
 			};
 			let records = RECORDS.init(record::read_record_file(&opt.record_file));
+
+			if let Some(only_record) = subopt.record {
+				records.retain(|id, _| *id == only_record);
+			}
+
 			let (sync_sender, sync_receiver) = oneshot::channel();
+			console_subscriber::init();
 			let client = tokio::spawn(client::play(
 				records,
 				tls_mode,

src/record.rs

@@ -212,7 +212,7 @@ pub fn print_records(records: &Records, print_packets: bool) {
 				}
 			}
 			if print_packets {
-				let data = if data.len() >= 256 {
+				let data = if data.len() >= 256 && *direction == Direction::ServerToClient {
 					&data[0..256]
 				} else {
 					data.as_slice()

src/server.rs

@@ -3,12 +3,9 @@ use crate::{
 	record::{Direction, Records},
 };
 
+use futures_util::stream::StreamExt;
 use std::{collections::HashMap, sync::Arc};
-use tokio::{
-	io::{AsyncReadExt, AsyncWriteExt},
-	net::TcpListener,
-	sync::oneshot,
-};
+use tokio::{io::AsyncWriteExt, net::TcpListener, sync::oneshot};
 use tokio_rustls::rustls::{
 	pki_types::{
 		CertificateDer, PrivateKeyDer,
@@ -17,6 +14,7 @@ use tokio_rustls::rustls::{
 	server::ResolvesServerCertUsingSni,
 	sign::CertifiedKey,
 };
+use tokio_util::codec::Framed;
 use x509_parser::prelude::GeneralName;
 
 pub async fn play(
@@ -27,7 +25,7 @@ pub async fn play(
 	sync_sender: oneshot::Sender<()>,
 ) {
 	let mut response_map = HashMap::new();
-	for (_id, (server_name, records)) in records.iter() {
+	for (id, (server_name, records)) in records.iter() {
 		let mut hash = None;
 		let mut responses = Vec::new();
 		for (direction, data) in records {
@@ -36,7 +34,7 @@ pub async fn play(
 					if let Some(hash) = hash
 						&& !responses.is_empty()
 					{
-						response_map.insert((server_name.to_vec(), hash), responses);
+						response_map.insert((server_name.to_vec(), hash), (id, responses));
 						responses = Vec::new();
 					}
 					hash = Some(
@@ -52,17 +50,21 @@ pub async fn play(
 		if let Some(hash) = hash
 			&& !responses.is_empty()
 		{
-			response_map.insert((server_name.to_vec(), hash), responses);
+			response_map.insert((server_name.to_vec(), hash), (id, responses));
 		}
 	}
 
+	let response_map = Arc::new(response_map);
+
 	match tls_mode {
 		TlsMode::Both | TlsMode::Server => {
 			let mut resolver = ResolvesServerCertUsingSni::new();
 			let config = tokio_rustls::rustls::ServerConfig::builder()
 				.with_no_client_auth()
 				.with_cert_resolver(Arc::new(ResolvesServerCertUsingSni::new()));
-			for file in std::fs::read_dir(cert_path).unwrap() {
+			for file in std::fs::read_dir(cert_path).unwrap_or_else(|e| {
+				panic!("Cannot read certificate directory `{cert_path}`: {e:?}")
+			}) {
 				match file {
 					Ok(file) => {
 						if file.file_name().as_encoded_bytes().ends_with(b".crt") {
@@ -77,7 +79,7 @@ pub async fn play(
 									let (_rem, cert) =
 										x509_parser::parse_x509_certificate(&data).unwrap();
 									if !cert.is_ca() {
-										println!("File: {:?}", file.file_name());
+										//println!("File: {:?}", file.file_name());
 										let mut key_path = file.path().to_path_buf();
 										key_path.pop();
 										let file_name =
@@ -110,7 +112,7 @@ pub async fn play(
 											.iter()
 										{
 											if let GeneralName::DNSName(name) = name {
-												resolver.add(dbg!(name), cert_key.clone()).unwrap();
+												resolver.add(name, cert_key.clone()).unwrap();
 											}
 										}
 									}
@@ -130,9 +132,7 @@ pub async fn play(
 					.with_cert_resolver(Arc::new(resolver)),
 			);
 
-			//let acceptor = tokio_rustls::TlsAcceptor::from(Arc::new(config));
 			let listener = TcpListener::bind(listen_addr).await.unwrap();
-			let response_map = Arc::new(response_map);
 			sync_sender.send(()).unwrap();
 			loop {
 				let config = config.clone();
@@ -143,28 +143,12 @@ pub async fn play(
 				);
 				//let acceptor = acceptor.clone();
 				let response_map = response_map.clone();
-				let fut = async move {
-					/*let mut server_name = None;
-					let mut stream = acceptor
-						.accept_with(stream, |conn| {
-							server_name = conn.server_name().map(String::from)
-						})
-						.await
-						.unwrap();
-					let server_name = server_name.unwrap();*/
+				/*let fut = async move {
 					let accepted = acceptor.await.unwrap();
 					let server_name = accepted.client_hello().server_name().unwrap().to_string();
 					let mut stream = accepted.into_stream(config).await.unwrap();
 					let mut req = Vec::new();
-					// TODO if there is a body
-					while !req.ends_with(b"\r\n\r\n") {
-						if stream.read_buf(&mut req).await.unwrap() == 0 {
-							break;
-						}
-					}
-					if let Ok(req) = str::from_utf8(&req) {
-						println!("{req}");
-					}
+					http::decode_http(&mut req, &mut stream).await;
 					let req_hash = tlsh::hash_buf(&req)
 						.map_or_else(|_| req.clone(), |h| h.to_string().into_bytes());
 					let mut best = None;
@@ -183,15 +167,54 @@ pub async fn play(
 						}
 					}
 					if let Some((hash, _diff)) = best {
-						let responses = response_map
+						let (id, responses) = response_map
 							.get(&(server_name.as_bytes().to_vec(), hash.clone()))
 							.unwrap();
 						for &res in responses {
+							println!("[SRV] response for ({}): {} bytes", id, res.len());
 							stream.write_all(res).await.unwrap();
 							stream.flush().await.unwrap();
 						}
 					} else {
-						eprintln!("No response found for SNI=`{server_name}`");
+						println!("No response found for SNI=`{server_name}`");
+					}
+					stream.shutdown().await.unwrap();
+				};*/
+				let fut = async move {
+					let accepted = acceptor.await.unwrap();
+					let server_name = accepted.client_hello().server_name().unwrap().to_string();
+					let stream = accepted.into_stream(config).await.unwrap();
+					let mut stream = Framed::new(stream, crate::http::HttpCodec {});
+					let req = stream.next().await.unwrap().unwrap();
+					let req_hash = tlsh::hash_buf(&req)
+						.map_or_else(|_| req.clone(), |h| h.to_string().into_bytes());
+					let mut best = None;
+					for (i_server_name, hash) in response_map.keys() {
+						if i_server_name != server_name.as_bytes() {
+							continue;
+						}
+						let diff = compare(&req_hash, hash);
+						if let Some((best_hash, best_diff)) = &mut best {
+							if diff < *best_diff {
+								*best_hash = hash;
+								*best_diff = diff;
+							}
+						} else {
+							best = Some((hash, diff));
+						}
+					}
+					let stream = stream.get_mut();
+					if let Some((hash, _diff)) = best {
+						let (id, responses) = response_map
+							.get(&(server_name.as_bytes().to_vec(), hash.clone()))
+							.unwrap();
+						for &res in responses {
+							println!("[SRV] response for ({}): {} bytes", id, res.len());
+							stream.write_all(res).await.unwrap();
+							stream.flush().await.unwrap();
+						}
+					} else {
+						println!("No response found for SNI=`{server_name}`");
 					}
 					stream.shutdown().await.unwrap();
 				};
@@ -201,7 +224,88 @@ pub async fn play(
 			}
 		}
 		TlsMode::None | TlsMode::Client => {
-			// TODO
+			let listener = TcpListener::bind(listen_addr).await.unwrap_or_else(|e| {
+				println!("Server: Cannot listen: {e:?}");
+				std::process::exit(1)
+			});
+			sync_sender.send(()).unwrap();
+			loop {
+				let (stream, _peer_addr) = listener.accept().await.unwrap();
+				let response_map = response_map.clone();
+				/*let fut = async move {
+					println!("[SRV] New task");
+					let mut req = Vec::new();
+					http::decode_http(&mut req, &mut stream).await;
+					let req_hash = tlsh::hash_buf(&req)
+						.map_or_else(|_| req.clone(), |h| h.to_string().into_bytes());
+					let mut best = None;
+					for (i_server_name, hash) in response_map.keys() {
+						let diff = compare(&req_hash, hash);
+						if let Some((best_server_name, best_hash, best_diff)) = &mut best {
+							if diff < *best_diff {
+								*best_server_name = i_server_name;
+								*best_hash = hash;
+								*best_diff = diff;
+							}
+						} else {
+							best = Some((i_server_name, hash, diff));
+						}
+					}
+					if let Some((server_name, hash, _diff)) = best {
+						let (id, responses) = response_map
+							.get(&(server_name.clone(), hash.clone()))
+							.unwrap();
+						for &res in responses {
+							println!("[SRV] response for ({}): {} bytes", id, res.len());
+							stream.write_all(res).await.unwrap();
+							stream.flush().await.unwrap();
+						}
+					} else {
+						println!("[SRV] No response found");
+					}
+					//println!("Server shutdown");
+					stream.shutdown().await.unwrap();
+				};*/
+				let fut = async move {
+					println!("[SRV] New task");
+					let mut stream = Framed::new(stream, crate::http::HttpCodec {});
+					let req = stream.next().await.unwrap().unwrap();
+					let req_hash = tlsh::hash_buf(&req)
+						.map_or_else(|_| req.clone(), |h| h.to_string().into_bytes());
+					let mut best = None;
+					for (i_server_name, hash) in response_map.keys() {
+						let diff = compare(&req_hash, hash);
+						if let Some((best_server_name, best_hash, best_diff)) = &mut best {
+							if diff < *best_diff {
+								*best_server_name = i_server_name;
+								*best_hash = hash;
+								*best_diff = diff;
+							}
+						} else {
+							best = Some((i_server_name, hash, diff));
+						}
+					}
+					let stream = stream.get_mut();
+					if let Some((server_name, hash, _diff)) = best {
+						let (id, responses) = response_map
+							.get(&(server_name.clone(), hash.clone()))
+							.unwrap();
+						for &res in responses {
+							println!("[SRV] response for ({}): {} bytes", id, res.len());
+							stream.write_all(res).await.unwrap();
+							stream.flush().await.unwrap();
+						}
+					} else {
+						println!("[SRV] No response found");
+					}
+					//println!("Server shutdown");
+					stream.shutdown().await.unwrap();
+				};
+				// Using a variable for the future allows it to be detected by tokio-console
+				tokio::spawn(async move {
+					fut.await;
+				});
+			}
 		}
 	}
 }