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

use structopt::StructOpt;
use tokio::net::{
	TcpListener,
	TcpStream,
};

use quic_demo::prelude::*;
use protocol::PeerId;

#[derive (Debug, StructOpt)]
struct Opt {
	#[structopt (long)]
	relay_addr: Option <String>,
	#[structopt (long)]
	client_id: Option <PeerId>,
	#[structopt (long)]
	server_id: Option <PeerId>,
	#[structopt (long)]
	client_tcp_port: Option <u16>,
	#[structopt (long)]
	server_tcp_port: Option <u16>,
}

#[tokio::main]
async fn main () -> anyhow::Result <()> {
	tracing_subscriber::fmt::init ();
	
	let opt = Opt::from_args ();
	let conf = opt.into_config ().await?;
	
	let client = P2Client::connect (conf)?;
	client.run ().await?;
	
	Ok (())
}

pub struct P2Client {
	endpoint: quinn::Endpoint,
	conf: Arc <Config>,
}

impl P2Client {
	pub fn connect (conf: Config) -> anyhow::Result <Self> {
		let endpoint = make_client_endpoint ("0.0.0.0:0".parse ()?, &[&conf.relay_cert])?;
		let conf = Arc::new (conf);
		
		Ok (Self {
			endpoint,
			conf,
		})
	}
	
	pub async fn run (&self) -> anyhow::Result <()> {
		debug! ("P2 client connecting to P3 relay server");
		
		let conf = Arc::clone (&self.conf);
		
		let quinn::NewConnection {
			connection,
			..
		} = protocol::p2_connect_to_p3 (&self.endpoint, &conf.relay_addr, &conf.client_id).await?;
		
		let listener = TcpListener::bind (("127.0.0.1", conf.client_tcp_port)).await?;
		
		debug! ("Accepting local TCP connections from P1 at {}", conf.client_tcp_port);
		
		// End of per-port stuff
		// Beginning of per-connection stuff
		
		let task_tcp_server = {
			let connection = connection.clone ();
			
			tokio::spawn (async move {
				Self::run_tcp_server (listener, connection, &*conf).await
			})
		};
		
		task_tcp_server.await??;
		
		Ok (())
	}
	
	/// Runs a TCP listen to forward a single TCP port
	
	async fn run_tcp_server (
		listener: TcpListener,
		connection: quinn::Connection,
		conf: &Config,
	) -> anyhow::Result <()> {
		let running = true;
		while running {
			let (tcp_stream, _) = listener.accept ().await?;
			let connection = connection.clone ();
			let server_id = conf.server_id.clone ();
			let server_tcp_port = conf.server_tcp_port;
			
			tokio::spawn (async move {
				Self::run_tcp_stream (&connection, tcp_stream, &server_id, server_tcp_port).await?;
				
				Ok::<_, anyhow::Error> (())
			});
		}
		
		Ok (())
	}
	
	/// Forwards a single TCP stream
	
	async fn run_tcp_stream (
		connection: &quinn::Connection,
		tcp_stream: TcpStream,
		server_id: &str,
		server_tcp_port: u16,
	) -> anyhow::Result <()> {
		let (local_recv, local_send) = tcp_stream.into_split ();
		
		debug! ("Starting PTTH connection");
		
		let (relay_send, relay_recv) = protocol::p2_connect_to_p5 (connection, server_id, server_tcp_port).await?;
		
		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?;
		
		debug! ("Ended PTTH connection");
		
		Ok (())
	}
}

/// A filled-out config for constructing a P2 client
#[derive (Clone)]
pub struct Config {
	client_tcp_port: u16,
	server_tcp_port: u16,
	client_id: String,
	server_id: String,
	relay_addr: SocketAddr,
	relay_cert: Vec <u8>,
}

impl Opt {
	pub async fn into_config (self) -> anyhow::Result <Config> {
		let client_tcp_port = self.client_tcp_port.unwrap_or (30381);
		let server_tcp_port = self.server_tcp_port.unwrap_or (30382);
		let client_id = self.client_id.unwrap_or_else (|| "bogus_client".to_string ());
		let server_id = self.server_id.unwrap_or_else (|| "bogus_server".to_string ());
		let relay_addr = self.relay_addr.unwrap_or_else (|| String::from ("127.0.0.1:30380")).parse ()?;
		
		// Begin I/O
		
		let relay_cert = tokio::fs::read ("ptth_quic_output/quic_server.crt").await?;
		
		Ok (Config {
			client_tcp_port,
			server_tcp_port,
			client_id, 
			server_id,
			relay_addr,
			relay_cert,
		})
	}
}