ptth/prototypes/quic_demo/src/protocol.rs

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,
) -> anyhow::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,
) -> anyhow::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 p4_connect_to_p3 (
	endpoint: &quinn::Endpoint,
	relay_addr: &std::net::SocketAddr,
	server_id: u8,
) -> anyhow::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)
}

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