use structopt::StructOpt; use tokio::{ net::TcpStream, sync::watch, }; 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 , #[structopt (long)] server_id: Option , #[structopt (long)] debug_echo: bool, #[structopt (long)] cert_url: Option , } pub async fn main (args: &[OsString], shutdown_rx: Option >) -> anyhow::Result <()> { let opt = Opt::from_iter (args); let conf = opt.into_config ().await?; let end_server = Arc::new (P4EndServer::connect (conf)?); let run_task = { let end_server = Arc::clone (&end_server); tokio::spawn (async move { end_server.run ().await?; Ok::<_, anyhow::Error> (()) }) }; if let Some (mut shutdown_rx) = shutdown_rx { while ! *shutdown_rx.borrow () { shutdown_rx.changed ().await?; } end_server.shut_down ()?; } run_task.await??; trace! ("P4 end server shut down gracefully."); 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 , } 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 { 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 = match self.cert_url.as_ref () { Some (url) => reqwest::get (url).await?.bytes ().await?.into_iter ().collect (), None => tokio::fs::read ("ptth_quic_output/quic_server.crt").await?, }; Ok (Config { debug_echo: self.debug_echo, id, relay_addr, relay_cert, }) } } pub struct P4EndServer { endpoint: quinn::Endpoint, conf: Arc , shutdown_tx: watch::Sender , shutdown_rx: watch::Receiver , } impl P4EndServer { pub fn connect (conf: Config) -> anyhow::Result { trace! ("P4 end server making its QUIC endpoint"); let endpoint = make_client_endpoint ("0.0.0.0:0".parse ()?, &[&conf.relay_cert])?; let (shutdown_tx, shutdown_rx) = watch::channel (false); Ok (P4EndServer { conf: Arc::new (conf), endpoint, shutdown_tx, shutdown_rx, }) } pub fn config (&self) -> &Config { &*self.conf } pub async fn run (&self) -> anyhow::Result <()> { trace! ("P4 end server connecting to P3 relay server"); let quinn::NewConnection { mut bi_streams, .. } = protocol::p4_connect_to_p3 ( &self.endpoint, &self.conf.relay_addr, &self.conf.id ).await?; debug! ("Connected to relay server"); trace! ("Accepting bi streams from P3"); let mut shutdown_rx = self.shutdown_rx.clone (); loop { tokio::select! { _ = shutdown_rx.changed () => { if *shutdown_rx.borrow () { trace! ("P4 incoming bi streams task caught graceful shutdown"); break; } } stream_opt = bi_streams.next () => { let (relay_send, relay_recv) = stream_opt.ok_or_else (|| anyhow::anyhow! ("P4 ran out of incoming streams. Maybe P3 shut down or disconnected?"))??; tokio::spawn (handle_bi_stream (Arc::clone (&self.conf), relay_send, relay_recv)); } }; } Ok (()) } pub fn shut_down (&self) -> anyhow::Result <()> { trace! ("P4 end server shutting down..."); Ok (self.shutdown_tx.send (true)?) } pub fn shutting_down (&self) -> bool { *self.shutdown_rx.borrow () } } async fn handle_bi_stream ( conf: Arc , 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 , 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 (()) }