♻️ refactor

main
_ 2021-07-17 07:57:55 +00:00
parent b75b7c5a74
commit 67fc89e640
1 changed files with 131 additions and 106 deletions

View File

@ -106,18 +106,15 @@ impl PtthNewConnection {
async fn handle_quic_connection ( async fn handle_quic_connection (
relay_state: Arc <Mutex <RelayState>>, relay_state: Arc <Mutex <RelayState>>,
conn: quinn::Connecting, conn: quinn::Connecting,
) -> anyhow::Result <()> { ) -> anyhow::Result <()>
let quinn::NewConnection { {
connection, let mut conn = conn.await?;
mut bi_streams,
..
} = conn.await?;
// Everyone who connects must identify themselves with the first // Everyone who connects must identify themselves with the first
// bi stream // bi stream
// TODO: Timeout // TODO: Timeout
let (mut send, mut recv) = bi_streams.next ().await.ok_or_else (|| anyhow::anyhow! ("QUIC client didn't identify itself"))??; let (mut send, mut recv) = conn.bi_streams.next ().await.ok_or_else (|| anyhow::anyhow! ("QUIC client didn't identify itself"))??;
let mut req_buf = [0u8; 4]; let mut req_buf = [0u8; 4];
recv.read_exact (&mut req_buf).await?; recv.read_exact (&mut req_buf).await?;
@ -136,7 +133,29 @@ async fn handle_quic_connection (
match peer_type { match peer_type {
2 => { 2 => {
let client_id = peer_id; handle_p2_connection (relay_state, peer_id, conn).await?;
},
4 => {
handle_p4_connection (relay_state, peer_id, conn).await?;
},
_ => bail! ("Unknown QUIC client type"),
}
debug! ("Peer {} disconnected", peer_id);
Ok::<_, anyhow::Error> (())
}
async fn handle_p2_connection (
relay_state: Arc <Mutex <RelayState>>,
client_id: u8,
conn: quinn::NewConnection,
) -> anyhow::Result <()>
{
let quinn::NewConnection {
mut bi_streams,
..
} = conn;
while let Some (bi_stream) = bi_streams.next ().await { while let Some (bi_stream) = bi_streams.next ().await {
let (mut client_send, mut client_recv) = bi_stream?; let (mut client_send, mut client_recv) = bi_stream?;
let relay_state = Arc::clone (&relay_state); let relay_state = Arc::clone (&relay_state);
@ -150,7 +169,7 @@ async fn handle_quic_connection (
1 => { 1 => {
let server_id = req_buf [1]; let server_id = req_buf [1];
debug! ("P2 {} wants to connect to P4 {}", peer_id, server_id); debug! ("P2 {} wants to connect to P4 {}", client_id, server_id);
// TODO: Auth checks // TODO: Auth checks
@ -180,9 +199,20 @@ async fn handle_quic_connection (
}); });
} }
debug! ("P2 {} disconnected", peer_id); debug! ("P2 {} disconnected", client_id);
}, Ok (())
4 => { }
async fn handle_p4_connection (
relay_state: Arc <Mutex <RelayState>>,
server_id: u8,
conn: quinn::NewConnection,
) -> anyhow::Result <()>
{
let quinn::NewConnection {
connection,
..
} = conn;
let (tx, mut rx) = mpsc::channel (2); let (tx, mut rx) = mpsc::channel (2);
let p4_state = P4State { let p4_state = P4State {
@ -191,7 +221,7 @@ async fn handle_quic_connection (
{ {
let mut relay_state = relay_state.lock ().await; let mut relay_state = relay_state.lock ().await;
relay_state.p4_server_proxies.insert (peer_id, p4_state); relay_state.p4_server_proxies.insert (server_id, p4_state);
} }
while let Some (req) = rx.recv ().await { while let Some (req) = rx.recv ().await {
@ -204,7 +234,7 @@ async fn handle_quic_connection (
client_id, client_id,
} = req; } = req;
debug! ("P4 {} got a request from P2 {}", peer_id, req.client_id); debug! ("P4 {} got a request from P2 {}", server_id, req.client_id);
let (mut server_send, mut server_recv) = connection.open_bi ().await?; let (mut server_send, mut server_recv) = connection.open_bi ().await?;
@ -236,11 +266,6 @@ async fn handle_quic_connection (
}); });
} }
debug! ("P4 {} disconnected", peer_id); debug! ("P4 {} disconnected", server_id);
}, Ok (())
_ => bail! ("Unknown QUIC client type"),
}
debug! ("Peer {} disconnected", peer_id);
Ok::<_, anyhow::Error> (())
} }