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 server_id = opt.server_id.unwrap_or (43);
	
	let quinn::NewConnection {
		mut bi_streams,
		..
	} = protocol::p4_connect_to_p3 (&endpoint, &relay_addr, server_id).await?;
	
	trace! ("Accepting bi streams from P3");
	
	loop {
		let (relay_send, relay_recv) = bi_streams.next ().await.ok_or_else (|| anyhow::anyhow! ("Relay server didn't open a bi stream"))??;
		
		tokio::spawn (handle_bi_stream (relay_send, relay_recv, local_tcp_port));
	}
}

async fn handle_bi_stream (
	relay_send: quinn::SendStream,
	mut relay_recv: quinn::RecvStream,
	local_tcp_port: u16,
) -> anyhow::Result <()>
{
	match protocol::p4_accept_p3_stream (&mut relay_recv).await? {
		protocol::P3ToP4Stream::NewPtthConnection {
			client_id,
			..
		} => handle_new_ptth_connection (relay_send, relay_recv, local_tcp_port).await?,
	}
	
	Ok (())
}

async fn handle_new_ptth_connection (
	mut relay_send: quinn::SendStream,
	mut relay_recv: quinn::RecvStream,
	local_tcp_port: u16,
) -> anyhow::Result <()>
{
	// 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 = quic_demo::connection::NewConnection {
		local_send,
		local_recv,
		relay_send,
		relay_recv,
	}.build ();
	
	ptth_conn.wait_for_close ().await?;
	
	Ok (())
}