feat: scaffold relay client auth workspace
This commit is contained in:
L
254
client/src/main.rs
Normal file
254
client/src/main.rs
Normal file
@@ -0,0 +1,254 @@
|
||||
use std::{collections::HashMap, sync::Arc, time::Duration};
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use common::{
|
||||
codec::{read_frame, write_frame},
|
||||
protocol::{
|
||||
ClientFrame, Heartbeat, RegisterRequest, ServerFrame, StreamClosed, StreamData,
|
||||
},
|
||||
};
|
||||
use tokio::{
|
||||
io::{AsyncReadExt, AsyncWriteExt},
|
||||
net::TcpStream,
|
||||
sync::{RwLock, mpsc},
|
||||
time::{MissedTickBehavior, sleep},
|
||||
};
|
||||
use tracing::{error, info, warn};
|
||||
|
||||
#[derive(Clone)]
|
||||
struct ClientConfig {
|
||||
relay_addr: String,
|
||||
token: String,
|
||||
region: String,
|
||||
local_addr: String,
|
||||
requested_subdomain: Option<String>,
|
||||
}
|
||||
|
||||
impl ClientConfig {
|
||||
fn from_env() -> Self {
|
||||
Self {
|
||||
relay_addr: std::env::var("DVV_RELAY_ADDR").unwrap_or_else(|_| "127.0.0.1:7000".into()),
|
||||
token: std::env::var("DVV_TOKEN").unwrap_or_else(|_| "dev-token-local".into()),
|
||||
region: std::env::var("DVV_REGION").unwrap_or_else(|_| "eu".into()),
|
||||
local_addr: std::env::var("DVV_LOCAL_ADDR").unwrap_or_else(|_| "127.0.0.1:25565".into()),
|
||||
requested_subdomain: std::env::var("DVV_SUBDOMAIN").ok(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type StreamSinkMap = Arc<RwLock<HashMap<String, mpsc::Sender<Vec<u8>>>>>;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| "client=info".into()),
|
||||
)
|
||||
.init();
|
||||
|
||||
let cfg = ClientConfig::from_env();
|
||||
let mut backoff = Duration::from_millis(300);
|
||||
|
||||
loop {
|
||||
match run_session(cfg.clone()).await {
|
||||
Ok(()) => warn!("session ended; reconnecting"),
|
||||
Err(e) => error!(error = %e, "session failed"),
|
||||
}
|
||||
sleep(backoff).await;
|
||||
backoff = (backoff * 2).min(Duration::from_secs(10));
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_session(cfg: ClientConfig) -> Result<()> {
|
||||
let stream = TcpStream::connect(&cfg.relay_addr)
|
||||
.await
|
||||
.with_context(|| format!("connect relay {}", cfg.relay_addr))?;
|
||||
let (mut reader, mut writer) = stream.into_split();
|
||||
|
||||
write_frame(
|
||||
&mut writer,
|
||||
&ClientFrame::Register(RegisterRequest {
|
||||
token: cfg.token.clone(),
|
||||
region: cfg.region.clone(),
|
||||
requested_subdomain: cfg.requested_subdomain.clone(),
|
||||
local_addr: cfg.local_addr.clone(),
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
let accepted = match read_frame::<_, ServerFrame>(&mut reader)
|
||||
.await
|
||||
.context("read register response")?
|
||||
{
|
||||
ServerFrame::RegisterAccepted(ok) => ok,
|
||||
ServerFrame::RegisterRejected { reason } => anyhow::bail!("register rejected: {reason}"),
|
||||
other => anyhow::bail!("unexpected frame at register: {other:?}"),
|
||||
};
|
||||
|
||||
info!(
|
||||
fqdn = %accepted.fqdn,
|
||||
session_id = %accepted.session_id,
|
||||
owner = %accepted.owner_instance_id,
|
||||
"registered tunnel"
|
||||
);
|
||||
|
||||
let sinks: StreamSinkMap = Arc::new(RwLock::new(HashMap::new()));
|
||||
let (out_tx, mut out_rx) = mpsc::channel::<ClientFrame>(1024);
|
||||
|
||||
let writer_task = tokio::spawn(async move {
|
||||
while let Some(frame) = out_rx.recv().await {
|
||||
write_frame(&mut writer, &frame).await?;
|
||||
}
|
||||
Ok::<(), anyhow::Error>(())
|
||||
});
|
||||
|
||||
let hb_tx = out_tx.clone();
|
||||
let hb_sinks = sinks.clone();
|
||||
let hb_session_id = accepted.session_id.clone();
|
||||
let hb_interval = Duration::from_secs(accepted.heartbeat_interval_secs.max(1));
|
||||
let hb_task = tokio::spawn(async move {
|
||||
let mut ticker = tokio::time::interval(hb_interval);
|
||||
ticker.set_missed_tick_behavior(MissedTickBehavior::Delay);
|
||||
loop {
|
||||
ticker.tick().await;
|
||||
let active_streams = hb_sinks.read().await.len() as u32;
|
||||
let frame = ClientFrame::Heartbeat(Heartbeat {
|
||||
session_id: hb_session_id.clone(),
|
||||
active_streams,
|
||||
bytes_in: 0,
|
||||
bytes_out: 0,
|
||||
});
|
||||
if hb_tx.send(frame).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
loop {
|
||||
let frame: ServerFrame = read_frame(&mut reader).await?;
|
||||
match frame {
|
||||
ServerFrame::Ping => {
|
||||
let _ = out_tx.send(ClientFrame::Pong).await;
|
||||
}
|
||||
ServerFrame::IncomingTcp(incoming) => {
|
||||
let cfg_clone = cfg.clone();
|
||||
let out_tx_clone = out_tx.clone();
|
||||
let sinks_clone = sinks.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = handle_incoming_stream(cfg_clone, incoming, out_tx_clone, sinks_clone).await {
|
||||
warn!(error = %e, "incoming stream handling failed");
|
||||
}
|
||||
});
|
||||
}
|
||||
ServerFrame::StreamData(StreamData { stream_id, data }) => {
|
||||
if let Some(tx) = sinks.read().await.get(&stream_id).cloned() {
|
||||
if tx.send(data).await.is_err() {
|
||||
sinks.write().await.remove(&stream_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
ServerFrame::StreamClosed(StreamClosed { stream_id, .. }) => {
|
||||
sinks.write().await.remove(&stream_id);
|
||||
}
|
||||
ServerFrame::DrainNotice { retry_after_ms, reason } => {
|
||||
warn!(retry_after_ms, reason = %reason, "relay draining");
|
||||
break;
|
||||
}
|
||||
ServerFrame::Error { message } => warn!(message = %message, "server error frame"),
|
||||
ServerFrame::RegisterAccepted(_) | ServerFrame::RegisterRejected { .. } => {
|
||||
warn!("ignoring unexpected register response after session start")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
hb_task.abort();
|
||||
writer_task.abort();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_incoming_stream(
|
||||
cfg: ClientConfig,
|
||||
incoming: common::protocol::IncomingTcp,
|
||||
out_tx: mpsc::Sender<ClientFrame>,
|
||||
sinks: StreamSinkMap,
|
||||
) -> Result<()> {
|
||||
let mut local = TcpStream::connect(&cfg.local_addr)
|
||||
.await
|
||||
.with_context(|| format!("connect local mc {}", cfg.local_addr))?;
|
||||
|
||||
if !incoming.initial_data.is_empty() {
|
||||
local.write_all(&incoming.initial_data).await?;
|
||||
}
|
||||
|
||||
let (local_read, local_write) = local.into_split();
|
||||
let (to_local_tx, to_local_rx) = mpsc::channel::<Vec<u8>>(128);
|
||||
sinks.write().await.insert(incoming.stream_id.clone(), to_local_tx);
|
||||
|
||||
let stream_id = incoming.stream_id.clone();
|
||||
let sinks_clone = sinks.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = run_local_writer(local_write, to_local_rx).await {
|
||||
warn!(stream_id = %stream_id, error = %e, "local writer ended");
|
||||
}
|
||||
sinks_clone.write().await.remove(&stream_id);
|
||||
});
|
||||
|
||||
let stream_id = incoming.stream_id.clone();
|
||||
let out_tx_clone = out_tx.clone();
|
||||
let sinks_clone = sinks.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = run_local_reader(local_read, out_tx_clone.clone(), stream_id.clone()).await {
|
||||
warn!(stream_id = %stream_id, error = %e, "local reader ended");
|
||||
}
|
||||
let _ = out_tx_clone
|
||||
.send(ClientFrame::StreamClosed(StreamClosed {
|
||||
stream_id: stream_id.clone(),
|
||||
reason: Some("local_reader_closed".into()),
|
||||
}))
|
||||
.await;
|
||||
sinks_clone.write().await.remove(&stream_id);
|
||||
});
|
||||
|
||||
info!(
|
||||
stream_id = %incoming.stream_id,
|
||||
peer = %incoming.peer_addr,
|
||||
hostname = %incoming.hostname,
|
||||
local = %cfg.local_addr,
|
||||
"connected relay stream to local minecraft"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn run_local_reader(
|
||||
mut reader: tokio::net::tcp::OwnedReadHalf,
|
||||
out_tx: mpsc::Sender<ClientFrame>,
|
||||
stream_id: String,
|
||||
) -> Result<()> {
|
||||
let mut buf = vec![0u8; 16 * 1024];
|
||||
loop {
|
||||
let n = reader.read(&mut buf).await?;
|
||||
if n == 0 {
|
||||
break;
|
||||
}
|
||||
out_tx
|
||||
.send(ClientFrame::StreamData(StreamData {
|
||||
stream_id: stream_id.clone(),
|
||||
data: buf[..n].to_vec(),
|
||||
}))
|
||||
.await
|
||||
.context("send local data to relay")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn run_local_writer(
|
||||
mut writer: tokio::net::tcp::OwnedWriteHalf,
|
||||
mut rx: mpsc::Receiver<Vec<u8>>,
|
||||
) -> Result<()> {
|
||||
while let Some(chunk) = rx.recv().await {
|
||||
writer.write_all(&chunk).await?;
|
||||
}
|
||||
writer.shutdown().await.ok();
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user