ptth/prototypes/quic_demo/src/protocol.rs

139 lines
2.9 KiB
Rust

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 req_buf = [
Command::CONNECT_P2_TO_P3.0,
client_id,
0,
0,
];
send.write_all (&req_buf).await?;
let mut resp_buf = [0, 0, 0, 0];
recv.read_exact (&mut resp_buf).await?;
let expected = [
Command::OKAY.0,
Command::CONNECT_P2_TO_P3.0,
0,
0,
];
if resp_buf != expected {
bail! ("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 relay_send, mut relay_recv) = connection.open_bi ().await?;
// Ask P3 if we can connect to P4
let req_buf = [
Command::CONNECT_P2_TO_P4.0,
server_id,
0,
0,
];
relay_send.write_all (&req_buf).await?;
let mut resp_buf = [0; 4];
relay_recv.read_exact (&mut resp_buf).await?;
let expected = [
Command::OKAY.0,
Command::CONNECT_P2_TO_P4.0,
0,
0,
];
if resp_buf != expected {
bail! ("P2 didn't get OK response when asking P3 to connect P2 to P4");
}
// Ask P4 if we can connect to P5
let req_buf = [
Command::CONNECT_P2_TO_P5.0,
0,
0,
0,
];
relay_send.write_all (&req_buf).await?;
let mut resp_buf = [0; 4];
relay_recv.read_exact (&mut resp_buf).await?;
let expected = [
Command::OKAY.0,
Command::CONNECT_P2_TO_P5.0,
0,
0,
];
if resp_buf != expected {
bail! ("P2 didn't get OK response when asking P4 to connect P2 to P5");
}
Ok ((relay_send, relay_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 req_buf = [
Command::CONNECT_P4_TO_P3.0,
server_id,
0,
0,
];
send.write_all (&req_buf).await?;
let mut resp_buf = [0, 0, 0, 0];
recv.read_exact (&mut resp_buf).await?;
let expected = [
Command::OKAY.0,
Command::CONNECT_P4_TO_P3.0,
0,
0,
];
if resp_buf != expected {
bail! ("P4 didn't get OK response when connecting to P3");
}
Ok (new_conn)
}