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

use tokio::task::JoinHandle;

use quic_demo::prelude::*;

#[tokio::main]
async fn main () -> anyhow::Result <()> {
	tracing_subscriber::fmt::init ();
	
	let server_addr = "0.0.0.0:30380".parse ()?;
	let (mut incoming, server_cert) = make_server_endpoint (server_addr)?;
	tokio::fs::write ("quic_server.crt", &server_cert).await?;
	
	debug! ("Waiting for end server to connect");
	
	let end_server_conn = incoming.next ().await.ok_or_else (|| anyhow::anyhow! ("No end server connection"))?;
	
	let end_server_conn = end_server_conn.await?;
	
	let quinn::NewConnection {
		connection: end_server_conn,
		..
	} = end_server_conn;
	
	debug! ("Waiting for client to connect");
	
	let client_conn = incoming.next ().await.ok_or_else (|| anyhow::anyhow! ("No client connection"))?;
	
	let client_conn = client_conn.await?;
	
	let quinn::NewConnection {
		connection: _client_conn,
		bi_streams: mut client_incoming_bi_streams,
		..
	} = client_conn;
	
	debug! ("Waiting for client to open bi stream");
	
	let (client_send, client_recv) = client_incoming_bi_streams.next ().await.ok_or_else (|| anyhow::anyhow! ("Client didn't open a bi stream"))??;
	
	debug! ("Opening bi stream to the end server");
	
	let (server_send, server_recv) = end_server_conn.open_bi ().await?;
	
	debug! ("Relaying bytes...");
	
	let ptth_conn = PtthNewConnection {
		client_send,
		client_recv,
		server_send,
		server_recv,
	}.build ();
	
	ptth_conn.uplink_task.await??;
	ptth_conn.downlink_task.await??;
	
	Ok (())
}

struct PtthNewConnection {
	client_send: quinn::SendStream,
	client_recv: quinn::RecvStream,
	server_send: quinn::SendStream,
	server_recv: quinn::RecvStream,
}

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

impl PtthNewConnection {
	fn build (self) -> PtthConnection {
		let Self {
			mut client_send,
			mut client_recv,
			mut server_send,
			mut server_recv,
		} = self;
		
		let uplink_task = tokio::spawn (async move {
			// Uplink - Client to end server
			
			let mut buf = vec! [0u8; 65_536];
			while let Some (bytes_read) = client_recv.read (&mut buf).await? {
				let buf_slice = &buf [0..bytes_read];
				server_send.write_all (buf_slice).await?;
			}
			
			debug! ("Uplink closed");
			
			Ok::<_, anyhow::Error> (())
		});
		
		let downlink_task = tokio::spawn (async move {
			// Downlink - End server to client
			
			let mut buf = vec! [0u8; 65_536];
			while let Some (bytes_read) = server_recv.read (&mut buf).await? {
				let buf_slice = &buf [0..bytes_read];
				client_send.write_all (buf_slice).await?;
			}
			
			debug! ("Downlink closed");
			
			Ok::<_, anyhow::Error> (())
		});
		
		PtthConnection {
			uplink_task,
			downlink_task,
		}
	}
}