ptth/prototypes/quic_demo/src/protocol.rs

use anyhow::Result;
use quinn::{
	SendStream,
	RecvStream,
};

use crate::prelude::*;

#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Command (pub u8);

impl Command {
	pub const CONNECT_P2_TO_P3: Command = Command (2);
	pub const CONNECT_P4_TO_P3: Command = Command (4);
	pub const CONNECT_P2_TO_P4: Command = Command (10);
	pub const CONNECT_P2_TO_P4_STEP_2: Command = Command (11);
	pub const CONNECT_P2_TO_P5: Command = Command (12);
	pub const OKAY: Command = Command (20);
}

pub async fn p2_connect_to_p3 (
	endpoint: &quinn::Endpoint,
	relay_addr: &std::net::SocketAddr,
	client_id: u8,
) -> Result <quinn::NewConnection> 
{
	let new_conn = endpoint.connect (relay_addr, "localhost")?.await?;
	let (mut send, mut recv) = new_conn.connection.open_bi ().await?;
	let cmd_type = Command::CONNECT_P2_TO_P3.0;
	
	send.write_all (&[cmd_type, client_id, 0, 0]).await?;
	
	expect_exact_response (&mut recv, [Command::OKAY.0, cmd_type, 0, 0]).await
	.context ("P2 didn't get OK response when connecting to P3")?;
	
	Ok (new_conn)
}

pub async fn p2_connect_to_p5 (
	connection: &quinn::Connection,
	server_id: u8,
) -> Result <(SendStream, RecvStream)>
{
	let (mut send, mut recv) = connection.open_bi ().await?;
	
	// Ask P3 if we can connect to P4
	
	let cmd_type = Command::CONNECT_P2_TO_P4.0;
	
	send.write_all (&[cmd_type, server_id, 0, 0]).await?;
	
	expect_exact_response (&mut recv, [Command::OKAY.0, cmd_type, 0, 0]).await
	.context ("P2 didn't get OK response when asking P3 to connect P2 to P4")?;
	
	// Ask P4 if we can connect to P5
	
	let cmd_type = Command::CONNECT_P2_TO_P5.0;
	
	send.write_all (&[cmd_type, 0, 0, 0]).await?;
	
	expect_exact_response (&mut recv, [Command::OKAY.0, cmd_type, 0, 0]).await
	.context ("P2 didn't get OK response when asking P4 to connect P2 to P5")?;
	
	Ok ((send, recv))
}

pub async fn p3_connect_p2_to_p4 (
	connection: &quinn::Connection,
	client_id: u8,
) -> Result <(SendStream, RecvStream)>
{
	let (mut send, mut recv) = connection.open_bi ().await?;
	
	let cmd_type = Command::CONNECT_P2_TO_P4_STEP_2.0;
	
	let buf = [
		cmd_type,
		client_id,
		0,
		0,
	];
	send.write_all (&buf).await?;
	
	expect_exact_response (&mut recv, [Command::OKAY.0, cmd_type, 0, 0]).await
	.context ("P3 didn't get OK response when asking P4 to connect P2 to P4")?;
	
	Ok ((send, recv))
}

pub async fn p4_connect_to_p3 (
	endpoint: &quinn::Endpoint,
	relay_addr: &std::net::SocketAddr,
	server_id: u8,
) -> Result <quinn::NewConnection> 
{
	let new_conn = endpoint.connect (relay_addr, "localhost")?.await?;
	let (mut send, mut recv) = new_conn.connection.open_bi ().await?;
	let cmd_type = Command::CONNECT_P4_TO_P3.0;
	
	send.write_all (&[cmd_type, server_id, 0, 0]).await?;
	
	expect_exact_response (&mut recv, [Command::OKAY.0, cmd_type, 0, 0]).await
	.context ("P4 didn't get OK response when connecting to P3")?;
	
	Ok (new_conn)
}

pub enum P3ToP4Stream {
	NewPtthConnection {
		client_id: u8,
	}
}

pub async fn p4_accept_p3_stream (
	recv: &mut RecvStream,
) -> Result <P3ToP4Stream>
{
	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_STEP_2 => P3ToP4Stream::NewPtthConnection {
			client_id: buf [1],
		},
		_ => bail! ("Invalid command type while P2 was accepting a new bi stream from P3"),
	})
}

pub async fn p4_authorize_p2_connection (
	send: &mut SendStream,
) -> Result <()>
{
	let buf = [
		Command::OKAY.0, 
		Command::CONNECT_P2_TO_P4_STEP_2.0,
		0, 
		0,
	];
	send.write_all (&buf).await?;
	
	Ok (())
}

pub async fn p4_authorize_p1_connection (
	send: &mut SendStream,
) -> Result <()>
{
	let buf = [
		Command::OKAY.0, 
		Command::CONNECT_P2_TO_P5.0,
		0, 
		0,
	];
	send.write_all (&buf).await?;
	
	Ok (())
}

pub async fn p4_expect_p5_request (
	recv: &mut RecvStream,
) -> Result <()>
{
	let mut buf = [0, 0, 0, 0];
	recv.read_exact (&mut buf).await?;
	let cmd_type = Command (buf [0]);
	if cmd_type != Command::CONNECT_P2_TO_P5 {
		bail! ("P4 expected CONNECT_P2_TO_P5 but P2 sent something different");
	}
	
	Ok (())
}

async fn expect_exact_response (
	recv: &mut RecvStream, expected: [u8; 4]
) -> Result <()>
{
	let mut buf = [0, 0, 0, 0];
	recv.read_exact (&mut buf).await?;
	if buf != expected {
		bail! ("Didn't receive exact response");
	}
	
	Ok (())
}