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 (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 {
			handle_quic_connection (relay_state, conn).await
		});
	}
	
	Ok (())
}

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

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

impl RelayState {
	
}

struct RequestP2ToP4 {
	client_send: quinn::SendStream,
	client_recv: quinn::RecvStream,
	client_id: u8,
}

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? {
				if bytes_read == 0 {
					break;
				}
				let buf_slice = &buf [0..bytes_read];
				trace! ("Uplink relaying {} bytes", bytes_read);
				server_send.write_all (buf_slice).await?;
			}
			
			trace! ("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];
				trace! ("Downlink relaying {} bytes", bytes_read);
				client_send.write_all (buf_slice).await?;
			}
			
			trace! ("Downlink closed");
			
			Ok::<_, anyhow::Error> (())
		});
		
		PtthConnection {
			uplink_task,
			downlink_task,
		}
	}
}

async fn handle_quic_connection (
	relay_state: Arc <RelayState>,
	conn: quinn::Connecting,
) -> anyhow::Result <()> 
{
	let mut conn = conn.await?;
	
	// Everyone who connects must identify themselves with the first
	// bi stream
	// TODO: Timeout
	
	let (mut send, mut recv) = conn.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 => {
			handle_p2_connection (relay_state, peer_id, conn).await?;
		},
		4 => {
			handle_p4_connection (relay_state, peer_id, conn).await?;
		},
		_ => bail! ("Unknown QUIC client type"),
	}
	
	debug! ("Peer {} disconnected", peer_id);
	Ok::<_, anyhow::Error> (())
}

async fn handle_p2_connection (
	relay_state: Arc <RelayState>,
	client_id: u8,
	conn: quinn::NewConnection,
) -> anyhow::Result <()>
{
	let quinn::NewConnection {
		mut bi_streams,
		..
	} = conn;
	
	while let Some (bi_stream) = bi_streams.next ().await {
		let (client_send, mut client_recv) = bi_stream?;
		let relay_state = Arc::clone (&relay_state);
		
		tokio::spawn (async move {
			debug! ("Request started for P2");
			
			let mut req_buf = [0u8; 4];
			client_recv.read_exact (&mut req_buf).await?;
			
			let cmd_type = req_buf [0];
			match cmd_type {
				1 => {
					let server_id = req_buf [1];
					handle_request_p2_to_p4 (relay_state, client_id, server_id, client_send, client_recv).await?;
				},
				_ => bail! ("Unknown command type from P2"),
			}
			
			debug! ("Request ended for P2");
			
			Ok::<_, anyhow::Error> (())
		});
	}
	
	debug! ("P2 {} disconnected", client_id);
	Ok (())
}

async fn handle_request_p2_to_p4 (
	relay_state: Arc <RelayState>,
	client_id: u8,
	server_id: u8,
	mut client_send: quinn::SendStream,
	client_recv: quinn::RecvStream,
) -> anyhow::Result <()>
{
	trace! ("P2 {} wants to connect to P4 {}", client_id, server_id);
	
	// TODO: Auth checks
	
	let resp_buf = [0, 0, 0, 0];
	client_send.write_all (&resp_buf).await?;
	
	{
		let p4_server_proxies = relay_state.p4_server_proxies.lock ().await;
		match p4_server_proxies.get (&server_id) {
			Some (p4_state) => {
				p4_state.req_channel.send (RequestP2ToP4 {
					client_send,
					client_recv,
					client_id,
				}).await.map_err (|_| anyhow::anyhow! ("Can't send request to P4 server"))?;
			},
			None => warn! ("That server isn't connected"),
		}
	}
				
	Ok (())
}

async fn handle_p4_connection (
	relay_state: Arc <RelayState>,
	server_id: u8,
	conn: quinn::NewConnection,
) -> anyhow::Result <()>
{
	let quinn::NewConnection {
		connection,
		..
	} = conn;
	let (tx, mut rx) = mpsc::channel (2);
	
	let p4_state = P4State {
		req_channel: tx,
	};
	
	{
		let mut p4_server_proxies = relay_state.p4_server_proxies.lock ().await;
		p4_server_proxies.insert (server_id, p4_state);
	}
	
	while let Some (req) = rx.recv ().await {
		let connection = connection.clone ();
		
		tokio::spawn (async move {
			let RequestP2ToP4 {
				client_send,
				client_recv,
				client_id,
			} = req;
			
			debug! ("P4 {} got a request from P2 {}", server_id, req.client_id);
			
			let (mut server_send, mut server_recv) = connection.open_bi ().await?;
			
			let req_buf = [2u8, client_id, 0, 0];
			server_send.write_all (&req_buf).await?;
			
			let mut resp_buf = [0u8, 0, 0, 0];
			server_recv.read_exact (&mut resp_buf).await?;
			
			let status_code = resp_buf [0];
			if status_code != 20 {
				bail! ("P4 rejected request from {}", client_id);
			}
			
			trace! ("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??;
			
			debug! ("Request ended for P4");
			Ok::<_, anyhow::Error> (())
		});
	}
	
	debug! ("P4 {} disconnected", server_id);
	Ok (())
}