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

use quic_demo::prelude::*;

#[tokio::main]
async fn main () -> anyhow::Result <()> {
	tracing_subscriber::fmt::init ();
	
	let server_addr = "0.0.0.0:30380".parse ()?;
	let (mut incoming, server_cert) = make_server_endpoint (server_addr)?;
	tokio::fs::write ("quic_server.crt", &server_cert).await?;
	
	let relay_state = RelayState::default ();
	let relay_state = Arc::new (Mutex::new (relay_state));
	
	while let Some (conn) = incoming.next ().await {
		let relay_state = Arc::clone (&relay_state);
		
		// Each new peer QUIC connection gets its own task
		tokio::spawn (async move {
			let quinn::NewConnection {
				connection,
				mut bi_streams,
				..
			} = conn.await?;
			
			// Everyone who connects must identify themselves with the first
			// bi stream
			// TODO: Timeout
			
			let (mut send, mut recv) = bi_streams.next ().await.ok_or_else (|| anyhow::anyhow! ("QUIC client didn't identify itself"))??;
			
			let mut req_buf = [0u8; 4];
			recv.read_exact (&mut req_buf).await?;
			
			let peer_type = req_buf [0];
			let peer_id = req_buf [1];
			
			match peer_type {
				4 => debug! ("Server-side proxy (P4) connected, ID {}", peer_id),
				2 => debug! ("Client-side proxy (P2) connected, ID {}", peer_id),
				_ => bail! ("Unknown QUIC client type"),
			}
			
			let resp_buf = [20u8, 0, 0, 0];
			send.write_all (&resp_buf).await?;
			
			match peer_type {
				2 => {
					while let Some (bi_stream) = bi_streams.next ().await {
						let (mut send, mut recv) = bi_stream?;
						
						tokio::spawn (async move {
							let mut req_buf = [0u8; 4];
							recv.read_exact (&mut req_buf).await?;
							
							let cmd_type = req_buf [0];
							match cmd_type {
								1 => {
									let server_id = req_buf [1];
									
									debug! ("P2 {} wants to connect to P4 {}", peer_id, server_id);
								},
								_ => bail! ("Unknown command type from P2"),
							}
							
							Ok::<_, anyhow::Error> (())
						});
					}
					
					debug! ("P2 {} disconnected", peer_id);
				},
				4 => {
					let (tx, mut rx) = mpsc::channel (2);
					
					let p4_state = P4State {
						req_channel: tx,
					};
					
					{
						let mut relay_state = relay_state.lock ().await;
						relay_state.p4_server_proxies.insert (peer_id, p4_state);
					}
					
					while let Some (req) = rx.recv ().await {
						debug! ("P4 {} got a request", peer_id);
					}
					
					debug! ("P4 {} disconnected", peer_id);
				},
				_ => bail! ("Unknown QUIC client type"),
			}
			
			debug! ("Peer {} disconnected", peer_id);
			Ok::<_, anyhow::Error> (())
		});
	}
	
	if false {
		debug! ("Waiting for end server to connect");
		
		let end_server_conn = incoming.next ().await.ok_or_else (|| anyhow::anyhow! ("No end server connection"))?;
		
		let end_server_conn = end_server_conn.await?;
		
		let quinn::NewConnection {
			connection: end_server_conn,
			..
		} = end_server_conn;
		
		debug! ("Waiting for client to connect");
		
		let client_conn = incoming.next ().await.ok_or_else (|| anyhow::anyhow! ("No client connection"))?;
		
		let client_conn = client_conn.await?;
		
		let quinn::NewConnection {
			connection: _client_conn,
			bi_streams: mut client_incoming_bi_streams,
			..
		} = client_conn;
		
		debug! ("Waiting for client to open bi stream");
		
		let (client_send, client_recv) = client_incoming_bi_streams.next ().await.ok_or_else (|| anyhow::anyhow! ("Client didn't open a bi stream"))??;
		
		debug! ("Opening bi stream to the end server");
		
		let (server_send, server_recv) = end_server_conn.open_bi ().await?;
		
		debug! ("Relaying bytes...");
		
		let ptth_conn = PtthNewConnection {
			client_send,
			client_recv,
			server_send,
			server_recv,
		}.build ();
		
		ptth_conn.uplink_task.await??;
		ptth_conn.downlink_task.await??;
	}
	
	Ok (())
}

#[derive (Default)]
struct RelayState {
	p4_server_proxies: HashMap <u8, P4State>,
}

struct P4State {
	req_channel: mpsc::Sender <RequestP2ToP4>,
}

impl RelayState {
	
}

enum RequestP2ToP4 {
	Connect {
		
	},
}

struct PtthNewConnection {
	client_send: quinn::SendStream,
	client_recv: quinn::RecvStream,
	server_send: quinn::SendStream,
	server_recv: quinn::RecvStream,
}

struct PtthConnection {
	uplink_task: JoinHandle <anyhow::Result <()>>,
	downlink_task: JoinHandle <anyhow::Result <()>>,
}

impl PtthNewConnection {
	fn build (self) -> PtthConnection {
		let Self {
			mut client_send,
			mut client_recv,
			mut server_send,
			mut server_recv,
		} = self;
		
		let uplink_task = tokio::spawn (async move {
			// Uplink - Client to end server
			
			let mut buf = vec! [0u8; 65_536];
			while let Some (bytes_read) = client_recv.read (&mut buf).await? {
				let buf_slice = &buf [0..bytes_read];
				server_send.write_all (buf_slice).await?;
			}
			
			debug! ("Uplink closed");
			
			Ok::<_, anyhow::Error> (())
		});
		
		let downlink_task = tokio::spawn (async move {
			// Downlink - End server to client
			
			let mut buf = vec! [0u8; 65_536];
			while let Some (bytes_read) = server_recv.read (&mut buf).await? {
				let buf_slice = &buf [0..bytes_read];
				client_send.write_all (buf_slice).await?;
			}
			
			debug! ("Downlink closed");
			
			Ok::<_, anyhow::Error> (())
		});
		
		PtthConnection {
			uplink_task,
			downlink_task,
		}
	}
}