ptth/prototypes/quic_demo/src/bin/quic_demo_client.rs

use tokio::{
	net::TcpListener,
};

use quic_demo::prelude::*;

#[tokio::main]
async fn main () -> anyhow::Result <()> {
	tracing_subscriber::fmt::init ();
	
	let server_cert = tokio::fs::read ("quic_server.crt").await?;
	let server_addr = "127.0.0.1:30380".parse ()?;
	let endpoint = make_client_endpoint ("0.0.0.0:0".parse ()?, &[&server_cert])?;
	
	trace! ("Connecting to relay server");
	
	let quinn::NewConnection {
		connection,
		..
	} = endpoint.connect (&server_addr, "localhost")?.await?;
	
	let (mut send, mut recv) = connection.open_bi ().await?;
	
	let our_id = 42;
	
	let req_buf = [
		Command::CONNECT_P2_TO_P3.0, 
		our_id, 
		0, 
		0,
	];
	send.write_all (&req_buf).await?;
	
	let mut resp_buf = [0, 0, 0, 0];
	recv.read_exact (&mut resp_buf).await?;
	assert_eq! (resp_buf, [
		Command::OKAY.0, 
		Command::CONNECT_P2_TO_P3.0, 
		0, 
		0,
	]);
	
	let listener = TcpListener::bind ("127.0.0.1:30381").await?;
	
	trace! ("Accepting local TCP connections from P1");
	
	loop {
		let (tcp_socket, _) = listener.accept ().await?;
		let connection = connection.clone ();
		
		tokio::spawn (async move {
			let (local_recv, local_send) = tcp_socket.into_split ();
			
			debug! ("Started PTTH connection");
			
			let (mut relay_send, mut relay_recv) = connection.open_bi ().await?;
			
			// Ask P3 if we can connect to P4
			
			let server_id = 43;
			let req_buf = [
				Command::CONNECT_P2_TO_P4.0, 
				server_id, 
				0, 
				0,
			];
			relay_send.write_all (&req_buf).await?;
			
			let mut resp_buf = [0; 4];
			relay_recv.read_exact (&mut resp_buf).await?;
			assert_eq! (resp_buf, [
				Command::OKAY.0,
				Command::CONNECT_P2_TO_P4.0, 
				0, 
				0,
			]);
			
			// Ask P4 if we can connect to P5
			
			let req_buf = [
				Command::CONNECT_P2_TO_P5.0, 
				0, 
				0, 
				0,
			];
			relay_send.write_all (&req_buf).await?;
			
			let mut resp_buf = [0; 4];
			relay_recv.read_exact (&mut resp_buf).await?;
			assert_eq! (resp_buf, [
				Command::OKAY.0, 
				Command::CONNECT_P2_TO_P5.0, 
				0, 
				0,
			]);
			
			trace! ("Relaying bytes...");
			
			let ptth_conn = PtthNewConnection {
				local_send,
				local_recv,
				relay_send,
				relay_recv,
			}.build ();
			
			ptth_conn.uplink_task.await??;
			ptth_conn.downlink_task.await??;
			
			debug! ("Ended PTTH connection");
			
			Ok::<_, anyhow::Error> (())
		});
	}
}

struct PtthNewConnection {
	local_send: tokio::net::tcp::OwnedWriteHalf,
	local_recv: tokio::net::tcp::OwnedReadHalf,
	relay_send: quinn::SendStream,
	relay_recv: quinn::RecvStream,
}

struct PtthConnection {
	uplink_task: JoinHandle <anyhow::Result <()>>,
	downlink_task: JoinHandle <anyhow::Result <()>>,
}

impl PtthNewConnection {
	fn build (self) -> PtthConnection {
		let Self {
			mut local_send,
			mut local_recv,
			mut relay_send,
			mut relay_recv,
		} = self;
		
		let uplink_task = tokio::spawn (async move {
			// Uplink - local client to relay server
			
			let mut buf = vec! [0u8; 65_536];
			loop {
				let bytes_read = local_recv.read (&mut buf).await?;
				if bytes_read == 0 {
					break;
				}
				let buf_slice = &buf [0..bytes_read];
				trace! ("Uplink relaying {} bytes", bytes_read);
				relay_send.write_all (buf_slice).await?;
			}
			
			trace! ("Uplink closed");
			
			Ok::<_, anyhow::Error> (())
		});

		let downlink_task = tokio::spawn (async move {
			// Downlink - Relay server to local client
			
			let mut buf = vec! [0u8; 65_536];
			while let Some (bytes_read) = relay_recv.read (&mut buf).await? {
				let buf_slice = &buf [0..bytes_read];
				trace! ("Downlink relaying {} bytes", bytes_read);
				local_send.write_all (buf_slice).await?;
			}
			
			trace! ("Downlink closed");
			
			Ok::<_, anyhow::Error> (())
		});
		
		PtthConnection {
			uplink_task,
			downlink_task,
		}
	}
}