♻️ refactor: all protocol code is in the protocol mod now

main
_ 2021-07-18 20:23:20 +00:00
parent 927ddecf97
commit fed401f622
2 changed files with 39 additions and 21 deletions

View File

@ -1,7 +1,6 @@
use structopt::StructOpt; use structopt::StructOpt;
use quic_demo::prelude::*; use quic_demo::prelude::*;
use protocol::Command;
#[derive (Debug, StructOpt)] #[derive (Debug, StructOpt)]
struct Opt { struct Opt {
@ -204,22 +203,16 @@ async fn handle_p2_connection (
} = conn; } = conn;
while let Some (bi_stream) = bi_streams.next ().await { 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); let relay_state = Arc::clone (&relay_state);
tokio::spawn (async move { tokio::spawn (async move {
debug! ("Request started for P2"); debug! ("Request started for P2");
let mut req_buf = [0, 0, 0, 0]; match protocol::p3_accept_p2_stream (&mut recv).await? {
client_recv.read_exact (&mut req_buf).await?; protocol::P2ToP3Stream::ConnectP2ToP4 {
server_id,
let cmd_type = req_buf [0]; } => handle_request_p2_to_p4 (relay_state, client_id, server_id, send, recv).await?,
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"),
} }
debug! ("Request ended for P2"); 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); 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 = [ protocol::p3_authorize_p2_to_p4_connection (&mut client_send).await?;
Command::OKAY.0,
Command::CONNECT_P2_TO_P4.0,
0,
0,
];
client_send.write_all (&resp_buf).await?;
{ {
let p4_server_proxies = relay_state.p4_server_proxies.lock ().await; let p4_server_proxies = relay_state.p4_server_proxies.lock ().await;

View File

@ -143,6 +143,37 @@ pub async fn p3_connect_p2_to_p4 (
Ok ((send, recv)) 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 ( pub async fn p4_connect_to_p3 (
endpoint: &quinn::Endpoint, endpoint: &quinn::Endpoint,
relay_addr: &std::net::SocketAddr, 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 { Command::CONNECT_P2_TO_P4_STEP_2 => P3ToP4Stream::NewPtthConnection {
client_id: buf [1], 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"),
}) })
} }