411 lines
9.4 KiB
Rust
411 lines
9.4 KiB
Rust
use std::{
|
|
convert::TryFrom,
|
|
};
|
|
|
|
use anyhow::Result;
|
|
use quinn::{
|
|
SendStream,
|
|
RecvStream,
|
|
};
|
|
|
|
use crate::prelude::*;
|
|
|
|
pub type PeerId = String;
|
|
const MAX_ID_LENGTH: usize = 128;
|
|
|
|
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
|
pub struct Command (pub u8);
|
|
|
|
// I can't remember how I picked the numbers. Just increment I guess,
|
|
// and then switch to a variable-length format around 200.
|
|
|
|
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 const DIREC_P2_TO_P3: Command = Command (21);
|
|
}
|
|
|
|
pub async fn p2_connect_to_p3 (
|
|
endpoint: &quinn::Endpoint,
|
|
relay_addr: std::net::SocketAddr,
|
|
client_id: &str,
|
|
) -> Result <quinn::NewConnection>
|
|
{
|
|
if client_id.as_bytes ().len () > MAX_ID_LENGTH {
|
|
bail! ("Client ID is longer than MAX_ID_LENGTH");
|
|
}
|
|
|
|
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, 0, 0, 0]).await?;
|
|
send_lv_string (&mut send, client_id).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: &str,
|
|
server_port: u16,
|
|
) -> 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, 0, 0, 0]).await?;
|
|
send_lv_string (&mut send, server_id).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?;
|
|
send.write_all (&server_port.to_le_bytes ()).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 p2_direc_to_p4 (
|
|
connection: &quinn::Connection,
|
|
server_id: &str,
|
|
) -> Result <Vec <u8>>
|
|
{
|
|
let (mut send, mut recv) = connection.open_bi ().await?;
|
|
|
|
let cmd_type = Command::DIREC_P2_TO_P3.0;
|
|
|
|
let mut cookie = vec! [0u8; 32];
|
|
rand::thread_rng ().fill_bytes (&mut cookie [..]);
|
|
let cookie = cookie;
|
|
|
|
send.write_all (&[cmd_type, 0, 0, 0]).await?;
|
|
send_lv_string (&mut send, server_id).await?;
|
|
send_lv_u16 (&mut send, &cookie).await?;
|
|
|
|
debug! ("Waiting for OK response for DIREC");
|
|
|
|
expect_exact_response (&mut recv, [Command::OKAY.0, cmd_type, 0, 0]).await?;
|
|
|
|
Ok (cookie)
|
|
}
|
|
|
|
pub enum P3Peer {
|
|
P2ClientProxy (P2ClientProxy),
|
|
P4ServerProxy (P4ServerProxy),
|
|
}
|
|
|
|
pub struct P2ClientProxy {
|
|
pub id: PeerId,
|
|
}
|
|
|
|
pub struct P4ServerProxy {
|
|
pub id: PeerId,
|
|
}
|
|
|
|
pub async fn p3_accept_peer (
|
|
recv: &mut RecvStream,
|
|
) -> Result <P3Peer>
|
|
{
|
|
let mut buf = [0, 0, 0, 0];
|
|
recv.read_exact (&mut buf).await?;
|
|
|
|
let command = Command (buf [0]);
|
|
|
|
Ok (match command {
|
|
Command::CONNECT_P2_TO_P3 => {
|
|
let id = recv_lv_string (recv, MAX_ID_LENGTH).await?;
|
|
debug! ("Client-side proxy (P2) connected, ID {}", id);
|
|
P3Peer::P2ClientProxy (P2ClientProxy {
|
|
id,
|
|
})
|
|
},
|
|
Command::CONNECT_P4_TO_P3 => {
|
|
let id = recv_lv_string (recv, MAX_ID_LENGTH).await?;
|
|
debug! ("Server-side proxy (P4) connected, ID {}", id);
|
|
P3Peer::P4ServerProxy (P4ServerProxy {
|
|
id,
|
|
})
|
|
},
|
|
_ => bail! ("Unknown QUIC client type"),
|
|
})
|
|
}
|
|
|
|
pub async fn p3_authorize_p2_peer (
|
|
send: &mut SendStream,
|
|
) -> Result <()>
|
|
{
|
|
send.write_all (&[Command::OKAY.0, Command::CONNECT_P2_TO_P3.0, 0, 0]).await?;
|
|
Ok (())
|
|
}
|
|
|
|
pub async fn p3_authorize_p4_peer (
|
|
send: &mut SendStream,
|
|
) -> Result <()>
|
|
{
|
|
send.write_all (&[Command::OKAY.0, Command::CONNECT_P4_TO_P3.0, 0, 0]).await?;
|
|
Ok (())
|
|
}
|
|
|
|
pub async fn p3_connect_p2_to_p4 (
|
|
connection: &quinn::Connection,
|
|
client_id: &str,
|
|
) -> Result <(SendStream, RecvStream)>
|
|
{
|
|
if client_id.as_bytes ().len () > MAX_ID_LENGTH {
|
|
bail! ("Client ID is longer than MAX_ID_LENGTH");
|
|
}
|
|
|
|
let (mut send, mut recv) = connection.open_bi ().await?;
|
|
|
|
let cmd_type = Command::CONNECT_P2_TO_P4_STEP_2.0;
|
|
|
|
let buf = [cmd_type, 0, 0, 0];
|
|
send.write_all (&buf).await?;
|
|
send_lv_string (&mut send, client_id).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 enum P2ToP3Stream {
|
|
ConnectP2ToP4 {
|
|
server_id: PeerId,
|
|
},
|
|
DirecP2ToP4 {
|
|
/// P2 wants a P2P connection to this P4
|
|
server_id: PeerId,
|
|
|
|
/// P2 will send this cookie over plain UDP to P3
|
|
/// P3 will learn P2's WAN address from that.
|
|
cookie: Vec <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 => {
|
|
let server_id = recv_lv_string (recv, MAX_ID_LENGTH).await?;
|
|
|
|
P2ToP3Stream::ConnectP2ToP4 {
|
|
server_id,
|
|
}
|
|
},
|
|
Command::DIREC_P2_TO_P3 => {
|
|
let server_id = recv_lv_string (recv, MAX_ID_LENGTH).await?;
|
|
let cookie = recv_lv_u16 (recv, 64).await?;
|
|
|
|
P2ToP3Stream::DirecP2ToP4 {
|
|
server_id,
|
|
cookie,
|
|
}
|
|
},
|
|
_ => 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 p3_authorize_p2_to_p4_direc (
|
|
send: &mut SendStream,
|
|
) -> Result <()>
|
|
{
|
|
send.write_all (&[Command::OKAY.0, Command::DIREC_P2_TO_P3.0, 0, 0]).await?;
|
|
Ok (())
|
|
}
|
|
|
|
pub async fn p4_connect_to_p3 (
|
|
endpoint: &quinn::Endpoint,
|
|
relay_addr: std::net::SocketAddr,
|
|
server_id: &str,
|
|
) -> Result <quinn::NewConnection>
|
|
{
|
|
if server_id.as_bytes ().len () > MAX_ID_LENGTH {
|
|
bail! ("Server ID is longer than MAX_ID_LENGTH");
|
|
}
|
|
|
|
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, 0, 0, 0]).await?;
|
|
send_lv_string (&mut send, server_id).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: PeerId,
|
|
}
|
|
}
|
|
|
|
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 => {
|
|
let client_id = recv_lv_string (recv, MAX_ID_LENGTH).await?;
|
|
|
|
P3ToP4Stream::NewPtthConnection {
|
|
client_id,
|
|
}
|
|
},
|
|
_ => bail! ("Invalid command type while P4 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 struct P2ConnectToP5Request {
|
|
pub port: u16,
|
|
}
|
|
|
|
pub async fn p4_expect_p5_request (
|
|
recv: &mut RecvStream,
|
|
) -> Result <P2ConnectToP5Request>
|
|
{
|
|
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");
|
|
}
|
|
let mut port_buf = [0, 0];
|
|
recv.read_exact (&mut port_buf).await?;
|
|
let port = u16::from_le_bytes (port_buf);
|
|
|
|
Ok (P2ConnectToP5Request {
|
|
port,
|
|
})
|
|
}
|
|
|
|
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, got {:?}", buf);
|
|
}
|
|
|
|
Ok (())
|
|
}
|
|
|
|
async fn send_lv_u16 (send: &mut SendStream, buf: &[u8]) -> Result <()>
|
|
{
|
|
if buf.len () >= 65_536 {
|
|
bail! ("Buffer is too long to send with u16 length prefix");
|
|
}
|
|
|
|
let len = u16::try_from (buf.len ())?;
|
|
send.write_all (&len.to_le_bytes ()).await?;
|
|
send.write_all (buf).await?;
|
|
|
|
Ok (())
|
|
}
|
|
|
|
async fn send_lv_string (send: &mut SendStream, s: &str) -> Result <()>
|
|
{
|
|
send_lv_u16 (send, s.as_bytes ()).await?;
|
|
|
|
Ok (())
|
|
}
|
|
|
|
async fn recv_lv_u16 (recv: &mut RecvStream, max_len: usize) -> Result <Vec <u8>>
|
|
{
|
|
let mut len_buf = [0, 0];
|
|
recv.read_exact (&mut len_buf).await?;
|
|
let len = u16::from_le_bytes (len_buf);
|
|
let len = usize::try_from (len)?;
|
|
|
|
if len > max_len {
|
|
bail! ("Buffer is longer than max_len");
|
|
}
|
|
|
|
// We could skip this allocation but who cares
|
|
let mut buf = vec! [0u8; len];
|
|
recv.read_exact (&mut buf).await?;
|
|
|
|
Ok (buf)
|
|
}
|
|
|
|
async fn recv_lv_string (recv: &mut RecvStream, max_len: usize) -> Result <String>
|
|
{
|
|
let buf = recv_lv_u16 (recv, max_len).await?;
|
|
let s = String::from_utf8 (buf)?;
|
|
|
|
Ok (s)
|
|
}
|