♻️ refactor: use the code from client_proxy instead

main
_ 2021-10-10 17:14:32 +00:00
parent c2caeb405c
commit c002665f6c
2 changed files with 29 additions and 61 deletions

View File

@ -1,10 +1,15 @@
use structopt::StructOpt; use structopt::StructOpt;
use tokio::net::{ use tokio::{
TcpListener, sync::watch,
TcpStream,
}; };
use quic_demo::prelude::*; use quic_demo::{
client_proxy::{
ForwardingParams,
forward_port,
},
prelude::*,
};
use protocol::PeerId; use protocol::PeerId;
#[derive (Debug, StructOpt)] #[derive (Debug, StructOpt)]
@ -60,76 +65,36 @@ impl P2Client {
.. ..
} = protocol::p2_connect_to_p3 (&self.endpoint, &conf.relay_addr, &conf.client_id).await?; } = protocol::p2_connect_to_p3 (&self.endpoint, &conf.relay_addr, &conf.client_id).await?;
let listener = TcpListener::bind (("127.0.0.1", conf.client_tcp_port)).await?; let client_tcp_port = conf.client_tcp_port;
debug! ("Accepting local TCP connections from P1 at {}", conf.client_tcp_port); debug! ("Accepting local TCP connections from P1 at {}", client_tcp_port);
// End of per-port stuff // End of per-port stuff
// Beginning of per-connection stuff // Beginning of per-connection stuff
let (_shutdown_flag_tx, shutdown_flag_rx) = watch::channel (true);
let task_tcp_server = { let task_tcp_server = {
let connection = connection.clone (); let connection = connection.clone ();
tokio::spawn (async move {
Self::run_tcp_server (listener, connection, &*conf).await
})
};
task_tcp_server.await??;
Ok (())
}
/// Runs a TCP listen to forward a single TCP port
async fn run_tcp_server (
listener: TcpListener,
connection: quinn::Connection,
conf: &Config,
) -> anyhow::Result <()> {
let running = true;
while running {
let (tcp_stream, _) = listener.accept ().await?;
let connection = connection.clone ();
let server_id = conf.server_id.clone (); let server_id = conf.server_id.clone ();
let server_tcp_port = conf.server_tcp_port; let server_tcp_port = conf.server_tcp_port;
tokio::spawn (async move { tokio::spawn (async move {
Self::run_tcp_stream (&connection, tcp_stream, &server_id, server_tcp_port).await?; forward_port (
connection,
ForwardingParams {
client_tcp_port,
server_id,
server_tcp_port,
},
shutdown_flag_rx,
).await?;
Ok::<_, anyhow::Error> (()) Ok::<_, anyhow::Error> (())
}); })
} };
Ok (()) task_tcp_server.await??;
}
/// Forwards a single TCP stream
async fn run_tcp_stream (
connection: &quinn::Connection,
tcp_stream: TcpStream,
server_id: &str,
server_tcp_port: u16,
) -> anyhow::Result <()> {
let (local_recv, local_send) = tcp_stream.into_split ();
debug! ("Starting PTTH connection");
let (relay_send, relay_recv) = protocol::p2_connect_to_p5 (connection, server_id, server_tcp_port).await?;
trace! ("Relaying bytes...");
let ptth_conn = quic_demo::connection::NewConnection {
local_send,
local_recv,
relay_send,
relay_recv,
}.build ();
ptth_conn.wait_for_close ().await?;
debug! ("Ended PTTH connection");
Ok (()) Ok (())
} }

View File

@ -44,7 +44,10 @@ pub struct ForwardingParams {
pub server_tcp_port: u16, pub server_tcp_port: u16,
} }
async fn forward_port ( /// Starts a TCP listener that can forward any number of TCP streams to
/// the same client:server port combination
pub async fn forward_port (
connection_p2_p3: quinn::Connection, connection_p2_p3: quinn::Connection,
params: ForwardingParams, params: ForwardingParams,
shutdown_flag_rx: tokio::sync::watch::Receiver <bool>, shutdown_flag_rx: tokio::sync::watch::Receiver <bool>,