♻️ refactor

main
_ 2021-07-17 20:02:16 +00:00
parent c575f5ca22
commit fe25ff3d34
2 changed files with 87 additions and 51 deletions

View File

@ -31,7 +31,7 @@ async fn main () -> anyhow::Result <()> {
let listener = TcpListener::bind ("127.0.0.1:30381").await?;
let (tcp_socket, _) = listener.accept ().await?;
let (mut local_recv, mut local_send) = tcp_socket.into_split ();
let (local_recv, local_send) = tcp_socket.into_split ();
debug! ("Connecting to end server");
@ -45,41 +45,77 @@ async fn main () -> anyhow::Result <()> {
debug! ("Relaying bytes...");
let uplink_task = tokio::spawn (async move {
// Uplink - local client to relay server
let mut buf = vec! [0u8; 65_536];
loop {
let bytes_read = local_recv.read (&mut buf).await?;
if bytes_read == 0 {
break;
}
let buf_slice = &buf [0..bytes_read];
debug! ("Uplink relaying {} bytes", bytes_read);
relay_send.write_all (buf_slice).await?;
}
debug! ("Uplink closed");
Ok::<_, anyhow::Error> (())
});
let ptth_conn = PtthNewConnection {
local_send,
local_recv,
relay_send,
relay_recv,
}.build ();
let downlink_task = tokio::spawn (async move {
// Downlink - Relay server to local client
let mut buf = vec! [0u8; 65_536];
while let Some (bytes_read) = relay_recv.read (&mut buf).await? {
let buf_slice = &buf [0..bytes_read];
debug! ("Downlink relaying {} bytes", bytes_read);
local_send.write_all (buf_slice).await?;
}
debug! ("Downlink closed");
Ok::<_, anyhow::Error> (())
});
ptth_conn.uplink_task.await??;
ptth_conn.downlink_task.await??;
uplink_task.await??;
downlink_task.await??;
Ok (())
}
struct PtthNewConnection {
local_send: tokio::net::tcp::OwnedWriteHalf,
local_recv: tokio::net::tcp::OwnedReadHalf,
relay_send: quinn::SendStream,
relay_recv: quinn::RecvStream,
}
struct PtthConnection {
uplink_task: JoinHandle <anyhow::Result <()>>,
downlink_task: JoinHandle <anyhow::Result <()>>,
}
impl PtthNewConnection {
fn build (self) -> PtthConnection {
let Self {
mut local_send,
mut local_recv,
mut relay_send,
mut relay_recv,
} = self;
let uplink_task = tokio::spawn (async move {
// Uplink - local client to relay server
let mut buf = vec! [0u8; 65_536];
loop {
let bytes_read = local_recv.read (&mut buf).await?;
if bytes_read == 0 {
break;
}
let buf_slice = &buf [0..bytes_read];
debug! ("Uplink relaying {} bytes", bytes_read);
relay_send.write_all (buf_slice).await?;
}
debug! ("Uplink closed");
Ok::<_, anyhow::Error> (())
});
let downlink_task = tokio::spawn (async move {
// Downlink - Relay server to local client
let mut buf = vec! [0u8; 65_536];
while let Some (bytes_read) = relay_recv.read (&mut buf).await? {
let buf_slice = &buf [0..bytes_read];
debug! ("Downlink relaying {} bytes", bytes_read);
local_send.write_all (buf_slice).await?;
}
debug! ("Downlink closed");
Ok::<_, anyhow::Error> (())
});
PtthConnection {
uplink_task,
downlink_task,
}
}
}

View File

@ -47,6 +47,21 @@ async fn main () -> anyhow::Result <()> {
debug! ("Relaying bytes...");
let uplink_task = tokio::spawn (async move {
// Downlink - Relay server to local client
let mut buf = vec! [0u8; 65_536];
while let Some (bytes_read) = relay_recv.read (&mut buf).await? {
let buf_slice = &buf [0..bytes_read];
debug! ("Uplink relaying {} bytes", bytes_read);
local_send.write_all (buf_slice).await?;
}
debug! ("Downlink closed");
Ok::<_, anyhow::Error> (())
});
let downlink_task = tokio::spawn (async move {
// Uplink - local client to relay server
let mut buf = vec! [0u8; 65_536];
@ -56,7 +71,7 @@ async fn main () -> anyhow::Result <()> {
break;
}
let buf_slice = &buf [0..bytes_read];
debug! ("Uplink relaying {} bytes", bytes_read);
debug! ("Downlink relaying {} bytes", bytes_read);
relay_send.write_all (buf_slice).await?;
}
@ -65,21 +80,6 @@ async fn main () -> anyhow::Result <()> {
Ok::<_, anyhow::Error> (())
});
let downlink_task = tokio::spawn (async move {
// Downlink - Relay server to local client
let mut buf = vec! [0u8; 65_536];
while let Some (bytes_read) = relay_recv.read (&mut buf).await? {
let buf_slice = &buf [0..bytes_read];
debug! ("Downlink relaying {} bytes", bytes_read);
local_send.write_all (buf_slice).await?;
}
debug! ("Downlink closed");
Ok::<_, anyhow::Error> (())
});
uplink_task.await??;
downlink_task.await??;