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

use structopt::StructOpt;
use tokio::net::TcpStream;

use quic_demo::prelude::*;
use protocol::Command;

#[derive (Debug, StructOpt)]
struct Opt {
	#[structopt (long)]
	relay_addr: Option <String>,
	#[structopt (long)]
	local_tcp_port: Option <u16>,
	#[structopt (long)]
	server_id: Option <u8>,
}

#[tokio::main]
async fn main () -> anyhow::Result <()> {
	tracing_subscriber::fmt::init ();
	
	let opt = Opt::from_args ();
	let local_tcp_port = opt.local_tcp_port.unwrap_or (30382);
	
	let server_cert = tokio::fs::read ("quic_server.crt").await?;
	let relay_addr = opt.relay_addr.unwrap_or_else (|| String::from ("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,
		mut bi_streams,
		..
	} = endpoint.connect (&relay_addr, "localhost")?.await?;
	
	let (mut send, mut recv) = connection.open_bi ().await?;
	
	let our_id = opt.server_id.unwrap_or (43);
	let req_buf = [
		Command::CONNECT_P4_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_P4_TO_P3.0,
		0,
		0,
	]);
	
	trace! ("Accepting bi streams from P3");
	
	loop {
		let (mut relay_send, mut relay_recv) = bi_streams.next ().await.ok_or_else (|| anyhow::anyhow! ("Relay server didn't open a bi stream"))??;
		
		tokio::spawn (async move {
			let mut req_buf = [0, 0, 0, 0];
			relay_recv.read_exact (&mut req_buf).await?;
			assert_eq! (req_buf [0], Command::CONNECT_P2_TO_P4_STEP_2.0);
			
			// TODO: Authorize P2 to connect to us
			
			let resp_buf = [
				Command::OKAY.0, 
				Command::CONNECT_P2_TO_P4_STEP_2.0,
				0, 
				0,
			];
			relay_send.write_all (&resp_buf).await?;
			
			let mut req_buf = [0, 0, 0, 0];
			relay_recv.read_exact (&mut req_buf).await?;
			assert_eq! (req_buf [0], Command::CONNECT_P2_TO_P5.0);
			
			// TODO: Authorize P2 to connect to P5
			
			let resp_buf = [
				Command::OKAY.0, 
				Command::CONNECT_P2_TO_P5.0,
				0,
				0,
			];
			relay_send.write_all (&resp_buf).await?;
			
			debug! ("Started PTTH connection");
			
			let stream = TcpStream::connect (("127.0.0.1", local_tcp_port)).await?;
			let (local_recv, local_send) = stream.into_split ();
			
			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??;
			
			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 {
			// 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! ("Uplink relaying {} bytes", bytes_read);
				local_send.write_all (buf_slice).await?;
			}
			
			trace! ("Downlink closed");
			
			Ok::<_, anyhow::Error> (())
		});
		
		let downlink_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! ("Downlink relaying {} bytes", bytes_read);
				relay_send.write_all (buf_slice).await?;
			}
			
			trace! ("Uplink closed");
			
			Ok::<_, anyhow::Error> (())
		});
		
		PtthConnection {
			uplink_task,
			downlink_task,
		}
	}
}