From fed401f622aad331477fb279272bc1edafdff9ea Mon Sep 17 00:00:00 2001
From: _ <>
Date: Sun, 18 Jul 2021 20:23:20 +0000
Subject: [PATCH] :recycle: refactor: all protocol code is in the protocol mod
 now

---
 .../src/bin/quic_demo_relay_server.rs         | 27 ++++-----------
 prototypes/quic_demo/src/protocol.rs          | 33 ++++++++++++++++++-
 2 files changed, 39 insertions(+), 21 deletions(-)

diff --git a/prototypes/quic_demo/src/bin/quic_demo_relay_server.rs b/prototypes/quic_demo/src/bin/quic_demo_relay_server.rs
index 1924efe..00d174f 100644
--- a/prototypes/quic_demo/src/bin/quic_demo_relay_server.rs
+++ b/prototypes/quic_demo/src/bin/quic_demo_relay_server.rs
@@ -1,7 +1,6 @@
 use structopt::StructOpt;
 
 use quic_demo::prelude::*;
-use protocol::Command;
 
 #[derive (Debug, StructOpt)]
 struct Opt {
@@ -204,22 +203,16 @@ async fn handle_p2_connection (
 	} = conn;
 	
 	while let Some (bi_stream) = bi_streams.next ().await {
-		let (client_send, mut client_recv) = bi_stream?;
+		let (send, mut recv) = bi_stream?;
 		let relay_state = Arc::clone (&relay_state);
 		
 		tokio::spawn (async move {
 			debug! ("Request started for P2");
 			
-			let mut req_buf = [0, 0, 0, 0];
-			client_recv.read_exact (&mut req_buf).await?;
-			
-			let cmd_type = req_buf [0];
-			match Command (cmd_type) {
-				Command::CONNECT_P2_TO_P4 => {
-					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"),
+			match protocol::p3_accept_p2_stream (&mut recv).await? {
+				protocol::P2ToP3Stream::ConnectP2ToP4 {
+					server_id,
+				} => handle_request_p2_to_p4 (relay_state, client_id, server_id, send, recv).await?,
 			}
 			
 			debug! ("Request ended for P2");
@@ -242,15 +235,9 @@ async fn handle_request_p2_to_p4 (
 {
 	trace! ("P2 {} wants to connect to P4 {}", client_id, server_id);
 	
-	// TODO: Auth checks
+	// TODO: Check authorization for P2 to connect to P4
 	
-	let resp_buf = [
-		Command::OKAY.0,
-		Command::CONNECT_P2_TO_P4.0,
-		0,
-		0,
-	];
-	client_send.write_all (&resp_buf).await?;
+	protocol::p3_authorize_p2_to_p4_connection (&mut client_send).await?;
 	
 	{
 		let p4_server_proxies = relay_state.p4_server_proxies.lock ().await;
diff --git a/prototypes/quic_demo/src/protocol.rs b/prototypes/quic_demo/src/protocol.rs
index 88befd1..1d79481 100644
--- a/prototypes/quic_demo/src/protocol.rs
+++ b/prototypes/quic_demo/src/protocol.rs
@@ -143,6 +143,37 @@ pub async fn p3_connect_p2_to_p4 (
 	Ok ((send, recv))
 }
 
+pub enum P2ToP3Stream {
+	ConnectP2ToP4 {
+		server_id: u8,
+	},
+}
+
+pub async fn p3_accept_p2_stream (
+	recv: &mut RecvStream,
+) -> Result <P2ToP3Stream>
+{
+	let mut buf = [0, 0, 0, 0];
+	recv.read_exact (&mut buf).await?;
+	
+	let cmd_type = buf [0];
+	
+	Ok (match Command (cmd_type) {
+		Command::CONNECT_P2_TO_P4 => P2ToP3Stream::ConnectP2ToP4 {
+			server_id: buf [1],
+		},
+		_ => bail! ("Invalid command type while P3 was accepting a new bi stream from P2"),
+	})
+}
+
+pub async fn p3_authorize_p2_to_p4_connection (
+	send: &mut SendStream,
+) -> Result <()>
+{
+	send.write_all (&[Command::OKAY.0, Command::CONNECT_P2_TO_P4.0, 0, 0]).await?;
+	Ok (())
+}
+
 pub async fn p4_connect_to_p3 (
 	endpoint: &quinn::Endpoint,
 	relay_addr: &std::net::SocketAddr,
@@ -180,7 +211,7 @@ pub async fn p4_accept_p3_stream (
 		Command::CONNECT_P2_TO_P4_STEP_2 => P3ToP4Stream::NewPtthConnection {
 			client_id: buf [1],
 		},
-		_ => bail! ("Invalid command type while P2 was accepting a new bi stream from P3"),
+		_ => bail! ("Invalid command type while P4 was accepting a new bi stream from P3"),
 	})
 }