ptth/prototypes/quic_demo/src/executable_end_server.rs

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

use crate::prelude::*;
use protocol::PeerId;

/// A partially-filled-out config that structopt can deal with
/// Try to turn this into a Config as soon as possible.
#[derive (Debug, StructOpt)]
struct Opt {
	#[structopt (long)]
	relay_addr: Option <String>,
	#[structopt (long)]
	server_id: Option <PeerId>,
	#[structopt (long)]
	debug_echo: bool,
	#[structopt (long)]
	cert_url: Option <String>,
}

pub async fn main (args: &[OsString]) -> anyhow::Result <()> {
	let opt = Opt::from_iter (args);
	let conf = opt.into_config ().await?;
	
	let mut end_server = P4EndServer::connect (conf).await?;
	
	end_server.run ().await?;
	Ok (())
}

/// A filled-out config for constructing an end server
#[derive (Clone)]
pub struct Config {
	pub debug_echo: bool,
	pub id: String,
	pub relay_addr: SocketAddr,
	pub relay_cert: Vec <u8>,
}

impl Opt {
	/// Converts self into a Config that the server can use.
	/// Performs I/O to load the relay cert from disk or from HTTP.
	/// Fails if arguments can't be parsed or if I/O fails.
	
	pub async fn into_config (self) -> anyhow::Result <Config> {
		let id = self.server_id.clone ().unwrap_or_else (|| "bogus_server".to_string ());
		
		let relay_addr: SocketAddr = self.relay_addr.clone ().unwrap_or_else (|| String::from ("127.0.0.1:30380")).parse ()?;
		
		// Do I/O after all parsing is done.
		// We don't want to waste a network request only to come back and error
		// out on like "127.oooo.1" not parsing into a relay address.
		
		let relay_cert: Vec <u8> = match self.cert_url.as_ref () {
			Some (url) => reqwest::get (url).await?.bytes ().await?.into_iter ().collect (),
			None => tokio::fs::read ("quic_server.crt").await?,
		};
		
		Ok (Config {
			debug_echo: self.debug_echo,
			id,
			relay_addr,
			relay_cert,
		})
	}
}

/// An end server that is connected to its relay server
pub struct P4EndServer {
	bi_streams: quinn::IncomingBiStreams,
	conf: Arc <Config>,
}

impl P4EndServer {
	/// Tries to connect to a relay server and return an end server.
	/// Fails if there's a network issue while making the QUIC connection
	/// to the relay server.
	
	pub async fn connect (conf: Config) -> anyhow::Result <Self> {
		trace! ("P4 end server making its QUIC endpoint");
		let endpoint = make_client_endpoint ("0.0.0.0:0".parse ()?, &[&conf.relay_cert])?;
		
		trace! ("P4 end server connecting to P3 relay server");
		let quinn::NewConnection {
			bi_streams,
			..
		} = protocol::p4_connect_to_p3 (&endpoint, &conf.relay_addr, &conf.id).await?;
		
		debug! ("Connected to relay server");
		
		Ok (P4EndServer {
			bi_streams,
			conf: Arc::new (conf),
		})
	}
	
	pub fn config (&self) -> &Config {
		&*self.conf
	}
	
	pub async fn run (&mut self) -> anyhow::Result <()> {
		trace! ("Accepting bi streams from P3");
		
		loop {
			let (relay_send, relay_recv) = self.bi_streams.next ().await.ok_or_else (|| anyhow::anyhow! ("Relay server didn't open a bi stream"))??;
			
			tokio::spawn (handle_bi_stream (Arc::clone (&self.conf), relay_send, relay_recv));
		}
	}
}

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

async fn handle_new_ptth_connection (
	conf: Arc <Config>,
	mut relay_send: quinn::SendStream,
	mut relay_recv: quinn::RecvStream,
	_client_id: String,
) -> anyhow::Result <()>
{
	// TODO: Check authorization for P2 --> P4
	
	protocol::p4_authorize_p2_connection (&mut relay_send).await?;
	let p4_to_p5_req = protocol::p4_expect_p5_request (&mut relay_recv).await?;
	
	// TODO: Check authorization for P1 --> P5
	
	protocol::p4_authorize_p1_connection (&mut relay_send).await?;
	
	debug! ("Started PTTH connection");
	
	if conf.debug_echo {
		relay_send.write (b"Connected to P4=P5 debug echo server\n").await?;
		debug! ("Relaying bytes using internal debug echo server (P4=P5)");
		tokio::io::copy (&mut relay_recv, &mut relay_send).await?;
	}
	else {
		let stream = TcpStream::connect (("127.0.0.1", p4_to_p5_req.port)).await?;
		let (local_recv, local_send) = stream.into_split ();
		
		trace! ("Relaying bytes...");
		
		let ptth_conn = crate::connection::NewConnection {
			local_send,
			local_recv,
			relay_send,
			relay_recv,
		}.build ();
		
		ptth_conn.wait_for_close ().await?;
	}
	
	Ok (())
}