2021-07-17 23:58:11 +00:00
|
|
|
use structopt::StructOpt;
|
|
|
|
|
2021-07-17 00:58:35 +00:00
|
|
|
use quic_demo::prelude::*;
|
2021-07-18 17:57:54 +00:00
|
|
|
use protocol::Command;
|
2021-07-17 00:43:20 +00:00
|
|
|
|
2021-07-17 23:58:11 +00:00
|
|
|
#[derive (Debug, StructOpt)]
|
|
|
|
struct Opt {
|
|
|
|
#[structopt (long)]
|
|
|
|
listen_addr: Option <String>,
|
|
|
|
}
|
|
|
|
|
2021-07-17 00:43:20 +00:00
|
|
|
#[tokio::main]
|
|
|
|
async fn main () -> anyhow::Result <()> {
|
|
|
|
tracing_subscriber::fmt::init ();
|
|
|
|
|
2021-07-17 23:58:11 +00:00
|
|
|
let opt = Opt::from_args ();
|
|
|
|
|
|
|
|
let listen_addr = opt.listen_addr.unwrap_or_else (|| String::from ("0.0.0.0:30380")).parse ()?;
|
|
|
|
let (mut incoming, server_cert) = make_server_endpoint (listen_addr)?;
|
2021-07-18 00:33:01 +00:00
|
|
|
println! ("Base64 cert: {}", base64::encode (&server_cert));
|
2021-07-17 00:43:20 +00:00
|
|
|
tokio::fs::write ("quic_server.crt", &server_cert).await?;
|
|
|
|
|
2021-07-17 06:43:14 +00:00
|
|
|
let relay_state = RelayState::default ();
|
2021-07-17 21:46:17 +00:00
|
|
|
let relay_state = Arc::new (relay_state);
|
2021-07-17 06:43:14 +00:00
|
|
|
|
|
|
|
while let Some (conn) = incoming.next ().await {
|
|
|
|
let relay_state = Arc::clone (&relay_state);
|
|
|
|
|
2021-07-17 06:50:54 +00:00
|
|
|
// Each new peer QUIC connection gets its own task
|
2021-07-17 06:43:14 +00:00
|
|
|
tokio::spawn (async move {
|
2021-07-17 22:03:55 +00:00
|
|
|
let active = relay_state.stats.quic.connect ();
|
|
|
|
debug! ("QUIC connections: {}", active);
|
|
|
|
|
|
|
|
match handle_quic_connection (Arc::clone (&relay_state), conn).await {
|
|
|
|
Ok (_) => (),
|
|
|
|
Err (e) => warn! ("handle_quic_connection {:?}", e),
|
|
|
|
}
|
|
|
|
|
|
|
|
let active = relay_state.stats.quic.disconnect ();
|
|
|
|
debug! ("QUIC connections: {}", active);
|
2021-07-17 06:43:14 +00:00
|
|
|
});
|
|
|
|
}
|
|
|
|
|
|
|
|
Ok (())
|
|
|
|
}
|
|
|
|
|
|
|
|
#[derive (Default)]
|
|
|
|
struct RelayState {
|
2021-07-17 21:46:17 +00:00
|
|
|
p4_server_proxies: Mutex <HashMap <u8, P4State>>,
|
2021-07-17 22:03:55 +00:00
|
|
|
stats: Stats,
|
|
|
|
}
|
|
|
|
|
|
|
|
#[derive (Default)]
|
|
|
|
struct Stats {
|
|
|
|
quic: ConnectEvents,
|
|
|
|
}
|
|
|
|
|
|
|
|
#[derive (Default)]
|
|
|
|
struct ConnectEvents {
|
|
|
|
connects: AtomicU64,
|
|
|
|
disconnects: AtomicU64,
|
|
|
|
}
|
|
|
|
|
|
|
|
impl ConnectEvents {
|
|
|
|
fn connect (&self) -> u64 {
|
|
|
|
let connects = self.connects.fetch_add (1, Ordering::Relaxed) + 1;
|
|
|
|
let disconnects = self.disconnects.load (Ordering::Relaxed);
|
|
|
|
connects - disconnects
|
|
|
|
}
|
|
|
|
|
|
|
|
fn disconnect (&self) -> u64 {
|
|
|
|
let disconnects = self.disconnects.fetch_add (1, Ordering::Relaxed) + 1;
|
|
|
|
let connects = self.connects.load (Ordering::Relaxed);
|
|
|
|
connects - disconnects
|
|
|
|
}
|
|
|
|
|
|
|
|
fn active (&self) -> u64 {
|
|
|
|
let connects = self.connects.load (Ordering::Relaxed);
|
|
|
|
let disconnects = self.disconnects.load (Ordering::Relaxed);
|
|
|
|
connects - disconnects
|
|
|
|
}
|
2021-07-17 07:11:34 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
struct P4State {
|
|
|
|
req_channel: mpsc::Sender <RequestP2ToP4>,
|
2021-07-17 22:03:55 +00:00
|
|
|
|
2021-07-17 06:43:14 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
impl RelayState {
|
2021-07-17 00:43:20 +00:00
|
|
|
|
2021-07-17 06:26:06 +00:00
|
|
|
}
|
|
|
|
|
2021-07-17 07:43:21 +00:00
|
|
|
struct RequestP2ToP4 {
|
|
|
|
client_send: quinn::SendStream,
|
|
|
|
client_recv: quinn::RecvStream,
|
|
|
|
client_id: u8,
|
2021-07-17 07:11:34 +00:00
|
|
|
}
|
|
|
|
|
2021-07-17 06:26:06 +00:00
|
|
|
struct PtthNewConnection {
|
|
|
|
client_send: quinn::SendStream,
|
|
|
|
client_recv: quinn::RecvStream,
|
|
|
|
server_send: quinn::SendStream,
|
|
|
|
server_recv: quinn::RecvStream,
|
|
|
|
}
|
|
|
|
|
|
|
|
struct PtthConnection {
|
|
|
|
uplink_task: JoinHandle <anyhow::Result <()>>,
|
|
|
|
downlink_task: JoinHandle <anyhow::Result <()>>,
|
|
|
|
}
|
|
|
|
|
|
|
|
impl PtthNewConnection {
|
|
|
|
fn build (self) -> PtthConnection {
|
|
|
|
let Self {
|
|
|
|
mut client_send,
|
|
|
|
mut client_recv,
|
|
|
|
mut server_send,
|
|
|
|
mut server_recv,
|
|
|
|
} = self;
|
2021-07-17 00:43:20 +00:00
|
|
|
|
2021-07-17 06:26:06 +00:00
|
|
|
let uplink_task = tokio::spawn (async move {
|
|
|
|
// Uplink - Client to end server
|
|
|
|
|
|
|
|
let mut buf = vec! [0u8; 65_536];
|
|
|
|
while let Some (bytes_read) = client_recv.read (&mut buf).await? {
|
2021-07-17 07:43:21 +00:00
|
|
|
if bytes_read == 0 {
|
|
|
|
break;
|
|
|
|
}
|
2021-07-17 06:26:06 +00:00
|
|
|
let buf_slice = &buf [0..bytes_read];
|
2021-07-17 21:42:51 +00:00
|
|
|
trace! ("Uplink relaying {} bytes", bytes_read);
|
2021-07-17 06:26:06 +00:00
|
|
|
server_send.write_all (buf_slice).await?;
|
|
|
|
}
|
|
|
|
|
2021-07-17 21:42:51 +00:00
|
|
|
trace! ("Uplink closed");
|
2021-07-17 06:26:06 +00:00
|
|
|
|
|
|
|
Ok::<_, anyhow::Error> (())
|
|
|
|
});
|
2021-07-17 00:43:20 +00:00
|
|
|
|
2021-07-17 06:26:06 +00:00
|
|
|
let downlink_task = tokio::spawn (async move {
|
|
|
|
// Downlink - End server to client
|
|
|
|
|
|
|
|
let mut buf = vec! [0u8; 65_536];
|
|
|
|
while let Some (bytes_read) = server_recv.read (&mut buf).await? {
|
|
|
|
let buf_slice = &buf [0..bytes_read];
|
2021-07-17 21:42:51 +00:00
|
|
|
trace! ("Downlink relaying {} bytes", bytes_read);
|
2021-07-17 06:26:06 +00:00
|
|
|
client_send.write_all (buf_slice).await?;
|
|
|
|
}
|
|
|
|
|
2021-07-17 21:42:51 +00:00
|
|
|
trace! ("Downlink closed");
|
2021-07-17 06:26:06 +00:00
|
|
|
|
|
|
|
Ok::<_, anyhow::Error> (())
|
|
|
|
});
|
2021-07-17 00:43:20 +00:00
|
|
|
|
2021-07-17 06:26:06 +00:00
|
|
|
PtthConnection {
|
|
|
|
uplink_task,
|
|
|
|
downlink_task,
|
|
|
|
}
|
|
|
|
}
|
2021-07-17 00:43:20 +00:00
|
|
|
}
|
2021-07-17 07:51:47 +00:00
|
|
|
|
|
|
|
async fn handle_quic_connection (
|
2021-07-17 21:46:17 +00:00
|
|
|
relay_state: Arc <RelayState>,
|
2021-07-17 07:51:47 +00:00
|
|
|
conn: quinn::Connecting,
|
2021-07-17 07:57:55 +00:00
|
|
|
) -> anyhow::Result <()>
|
|
|
|
{
|
|
|
|
let mut conn = conn.await?;
|
2021-07-17 07:51:47 +00:00
|
|
|
|
|
|
|
// Everyone who connects must identify themselves with the first
|
|
|
|
// bi stream
|
|
|
|
// TODO: Timeout
|
|
|
|
|
2021-07-17 07:57:55 +00:00
|
|
|
let (mut send, mut recv) = conn.bi_streams.next ().await.ok_or_else (|| anyhow::anyhow! ("QUIC client didn't identify itself"))??;
|
2021-07-17 07:51:47 +00:00
|
|
|
|
2021-07-17 23:14:50 +00:00
|
|
|
let mut req_buf = [0, 0, 0, 0];
|
2021-07-17 07:51:47 +00:00
|
|
|
recv.read_exact (&mut req_buf).await?;
|
|
|
|
|
|
|
|
let peer_type = req_buf [0];
|
|
|
|
let peer_id = req_buf [1];
|
|
|
|
|
|
|
|
match peer_type {
|
|
|
|
4 => debug! ("Server-side proxy (P4) connected, ID {}", peer_id),
|
|
|
|
2 => debug! ("Client-side proxy (P2) connected, ID {}", peer_id),
|
|
|
|
_ => bail! ("Unknown QUIC client type"),
|
|
|
|
}
|
|
|
|
|
2021-07-17 23:14:50 +00:00
|
|
|
let resp_buf = [
|
|
|
|
Command::OKAY.0,
|
|
|
|
peer_type,
|
|
|
|
0,
|
|
|
|
0,
|
|
|
|
];
|
2021-07-17 07:51:47 +00:00
|
|
|
send.write_all (&resp_buf).await?;
|
|
|
|
|
|
|
|
match peer_type {
|
|
|
|
2 => {
|
2021-07-17 07:57:55 +00:00
|
|
|
handle_p2_connection (relay_state, peer_id, conn).await?;
|
2021-07-17 07:51:47 +00:00
|
|
|
},
|
|
|
|
4 => {
|
2021-07-17 07:57:55 +00:00
|
|
|
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 (
|
2021-07-17 21:46:17 +00:00
|
|
|
relay_state: Arc <RelayState>,
|
2021-07-17 07:57:55 +00:00
|
|
|
client_id: u8,
|
|
|
|
conn: quinn::NewConnection,
|
|
|
|
) -> anyhow::Result <()>
|
|
|
|
{
|
|
|
|
let quinn::NewConnection {
|
|
|
|
mut bi_streams,
|
|
|
|
..
|
|
|
|
} = conn;
|
|
|
|
|
|
|
|
while let Some (bi_stream) = bi_streams.next ().await {
|
2021-07-17 08:02:53 +00:00
|
|
|
let (client_send, mut client_recv) = bi_stream?;
|
2021-07-17 07:57:55 +00:00
|
|
|
let relay_state = Arc::clone (&relay_state);
|
|
|
|
|
|
|
|
tokio::spawn (async move {
|
2021-07-17 21:42:51 +00:00
|
|
|
debug! ("Request started for P2");
|
|
|
|
|
2021-07-17 23:14:50 +00:00
|
|
|
let mut req_buf = [0, 0, 0, 0];
|
2021-07-17 07:57:55 +00:00
|
|
|
client_recv.read_exact (&mut req_buf).await?;
|
2021-07-17 07:51:47 +00:00
|
|
|
|
2021-07-17 07:57:55 +00:00
|
|
|
let cmd_type = req_buf [0];
|
2021-07-17 23:14:50 +00:00
|
|
|
match Command (cmd_type) {
|
|
|
|
Command::CONNECT_P2_TO_P4 => {
|
2021-07-17 07:57:55 +00:00
|
|
|
let server_id = req_buf [1];
|
2021-07-17 08:02:53 +00:00
|
|
|
handle_request_p2_to_p4 (relay_state, client_id, server_id, client_send, client_recv).await?;
|
2021-07-17 07:57:55 +00:00
|
|
|
},
|
|
|
|
_ => bail! ("Unknown command type from P2"),
|
2021-07-17 07:51:47 +00:00
|
|
|
}
|
|
|
|
|
2021-07-17 07:57:55 +00:00
|
|
|
debug! ("Request ended for P2");
|
|
|
|
|
|
|
|
Ok::<_, anyhow::Error> (())
|
|
|
|
});
|
2021-07-17 07:51:47 +00:00
|
|
|
}
|
|
|
|
|
2021-07-17 07:57:55 +00:00
|
|
|
debug! ("P2 {} disconnected", client_id);
|
|
|
|
Ok (())
|
|
|
|
}
|
|
|
|
|
2021-07-17 08:02:53 +00:00
|
|
|
async fn handle_request_p2_to_p4 (
|
2021-07-17 21:46:17 +00:00
|
|
|
relay_state: Arc <RelayState>,
|
2021-07-17 08:02:53 +00:00
|
|
|
client_id: u8,
|
|
|
|
server_id: u8,
|
|
|
|
mut client_send: quinn::SendStream,
|
|
|
|
client_recv: quinn::RecvStream,
|
|
|
|
) -> anyhow::Result <()>
|
|
|
|
{
|
2021-07-17 21:42:51 +00:00
|
|
|
trace! ("P2 {} wants to connect to P4 {}", client_id, server_id);
|
2021-07-17 08:02:53 +00:00
|
|
|
|
|
|
|
// TODO: Auth checks
|
|
|
|
|
2021-07-17 23:14:50 +00:00
|
|
|
let resp_buf = [
|
|
|
|
Command::OKAY.0,
|
|
|
|
Command::CONNECT_P2_TO_P4.0,
|
|
|
|
0,
|
|
|
|
0,
|
|
|
|
];
|
2021-07-17 08:02:53 +00:00
|
|
|
client_send.write_all (&resp_buf).await?;
|
|
|
|
|
|
|
|
{
|
2021-07-17 21:46:17 +00:00
|
|
|
let p4_server_proxies = relay_state.p4_server_proxies.lock ().await;
|
|
|
|
match p4_server_proxies.get (&server_id) {
|
2021-07-17 08:02:53 +00:00
|
|
|
Some (p4_state) => {
|
|
|
|
p4_state.req_channel.send (RequestP2ToP4 {
|
|
|
|
client_send,
|
|
|
|
client_recv,
|
|
|
|
client_id,
|
|
|
|
}).await.map_err (|_| anyhow::anyhow! ("Can't send request to P4 server"))?;
|
|
|
|
},
|
|
|
|
None => warn! ("That server isn't connected"),
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
Ok (())
|
|
|
|
}
|
|
|
|
|
2021-07-17 07:57:55 +00:00
|
|
|
async fn handle_p4_connection (
|
2021-07-17 21:46:17 +00:00
|
|
|
relay_state: Arc <RelayState>,
|
2021-07-17 07:57:55 +00:00
|
|
|
server_id: u8,
|
|
|
|
conn: quinn::NewConnection,
|
|
|
|
) -> anyhow::Result <()>
|
|
|
|
{
|
|
|
|
let quinn::NewConnection {
|
|
|
|
connection,
|
|
|
|
..
|
|
|
|
} = conn;
|
|
|
|
let (tx, mut rx) = mpsc::channel (2);
|
|
|
|
|
|
|
|
let p4_state = P4State {
|
|
|
|
req_channel: tx,
|
|
|
|
};
|
|
|
|
|
|
|
|
{
|
2021-07-17 21:46:17 +00:00
|
|
|
let mut p4_server_proxies = relay_state.p4_server_proxies.lock ().await;
|
|
|
|
p4_server_proxies.insert (server_id, p4_state);
|
2021-07-17 07:57:55 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
while let Some (req) = rx.recv ().await {
|
|
|
|
let connection = connection.clone ();
|
|
|
|
|
|
|
|
tokio::spawn (async move {
|
|
|
|
let RequestP2ToP4 {
|
|
|
|
client_send,
|
|
|
|
client_recv,
|
|
|
|
client_id,
|
|
|
|
} = req;
|
|
|
|
|
|
|
|
debug! ("P4 {} got a request from P2 {}", server_id, req.client_id);
|
|
|
|
|
|
|
|
let (mut server_send, mut server_recv) = connection.open_bi ().await?;
|
|
|
|
|
2021-07-17 23:14:50 +00:00
|
|
|
let req_buf = [
|
|
|
|
Command::CONNECT_P2_TO_P4_STEP_2.0,
|
|
|
|
client_id,
|
|
|
|
0,
|
|
|
|
0,
|
|
|
|
];
|
2021-07-17 07:57:55 +00:00
|
|
|
server_send.write_all (&req_buf).await?;
|
|
|
|
|
2021-07-17 23:14:50 +00:00
|
|
|
let mut resp_buf = [0, 0, 0, 0];
|
2021-07-17 07:57:55 +00:00
|
|
|
server_recv.read_exact (&mut resp_buf).await?;
|
2021-07-17 23:14:50 +00:00
|
|
|
assert_eq! (resp_buf, [
|
|
|
|
Command::OKAY.0,
|
|
|
|
Command::CONNECT_P2_TO_P4_STEP_2.0,
|
|
|
|
0,
|
|
|
|
0,
|
|
|
|
]);
|
2021-07-17 07:57:55 +00:00
|
|
|
|
2021-07-17 21:42:51 +00:00
|
|
|
trace! ("Relaying bytes...");
|
2021-07-17 07:57:55 +00:00
|
|
|
|
|
|
|
let ptth_conn = PtthNewConnection {
|
|
|
|
client_send,
|
|
|
|
client_recv,
|
|
|
|
server_send,
|
|
|
|
server_recv,
|
|
|
|
}.build ();
|
|
|
|
|
|
|
|
ptth_conn.uplink_task.await??;
|
|
|
|
ptth_conn.downlink_task.await??;
|
|
|
|
|
|
|
|
debug! ("Request ended for P4");
|
|
|
|
Ok::<_, anyhow::Error> (())
|
|
|
|
});
|
|
|
|
}
|
|
|
|
|
|
|
|
debug! ("P4 {} disconnected", server_id);
|
|
|
|
Ok (())
|
2021-07-17 07:51:47 +00:00
|
|
|
}
|