feat: scaffold relay client auth workspace
This commit is contained in:
L
16
relay/Cargo.toml
Normal file
16
relay/Cargo.toml
Normal file
@@ -0,0 +1,16 @@
|
||||
[package]
|
||||
name = "relay"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
anyhow.workspace = true
|
||||
tokio.workspace = true
|
||||
tracing.workspace = true
|
||||
tracing-subscriber.workspace = true
|
||||
uuid.workspace = true
|
||||
fastrand.workspace = true
|
||||
redis.workspace = true
|
||||
serde_json.workspace = true
|
||||
chrono.workspace = true
|
||||
common = { path = "../common" }
|
||||
699
relay/src/main.rs
Normal file
699
relay/src/main.rs
Normal file
@@ -0,0 +1,699 @@
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
net::SocketAddr,
|
||||
sync::Arc,
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use common::{
|
||||
codec::{read_frame, write_frame},
|
||||
minecraft::read_handshake_hostname_and_bytes,
|
||||
protocol::{
|
||||
ClientFrame, Heartbeat, IncomingTcp, RegisterAccepted, RegisterRequest, ServerFrame,
|
||||
StreamClosed, StreamData,
|
||||
},
|
||||
};
|
||||
use redis::AsyncCommands;
|
||||
use tokio::{
|
||||
io::{AsyncReadExt, AsyncWriteExt},
|
||||
net::{TcpListener, TcpStream},
|
||||
sync::{Notify, RwLock, mpsc},
|
||||
time::{MissedTickBehavior, interval, timeout},
|
||||
};
|
||||
use tracing::{debug, info, warn};
|
||||
use uuid::Uuid;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct RelayConfig {
|
||||
instance_id: String,
|
||||
region: String,
|
||||
control_bind: String,
|
||||
player_bind: String,
|
||||
domain: String,
|
||||
heartbeat_timeout: Duration,
|
||||
registry_ttl_secs: u64,
|
||||
}
|
||||
|
||||
impl RelayConfig {
|
||||
fn from_env() -> Self {
|
||||
Self {
|
||||
instance_id: std::env::var("RELAY_INSTANCE_ID")
|
||||
.unwrap_or_else(|_| format!("relay-{}", Uuid::new_v4())),
|
||||
region: std::env::var("RELAY_REGION").unwrap_or_else(|_| "eu".to_string()),
|
||||
control_bind: std::env::var("RELAY_CONTROL_BIND")
|
||||
.unwrap_or_else(|_| "0.0.0.0:7000".to_string()),
|
||||
player_bind: std::env::var("RELAY_PLAYER_BIND")
|
||||
.unwrap_or_else(|_| "0.0.0.0:25565".to_string()),
|
||||
domain: std::env::var("RELAY_BASE_DOMAIN").unwrap_or_else(|_| "dvv.one".to_string()),
|
||||
heartbeat_timeout: Duration::from_secs(
|
||||
std::env::var("RELAY_HEARTBEAT_TIMEOUT_SECS")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(30),
|
||||
),
|
||||
registry_ttl_secs: std::env::var("RELAY_REGISTRY_TTL_SECS")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(20),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct SessionHandle {
|
||||
session_id: String,
|
||||
tx: mpsc::Sender<ServerFrame>,
|
||||
stream_sinks: Arc<RwLock<HashMap<String, mpsc::Sender<Vec<u8>>>>>,
|
||||
last_heartbeat: Instant,
|
||||
}
|
||||
|
||||
struct RelayState {
|
||||
by_fqdn: HashMap<String, SessionHandle>,
|
||||
by_session: HashMap<String, String>,
|
||||
}
|
||||
|
||||
impl RelayState {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
by_fqdn: HashMap::new(),
|
||||
by_session: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn session_count(&self) -> usize {
|
||||
self.by_session.len()
|
||||
}
|
||||
}
|
||||
|
||||
type SharedState = Arc<RwLock<RelayState>>;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct RedisRegistry {
|
||||
conn: Option<redis::aio::ConnectionManager>,
|
||||
instance_id: String,
|
||||
region: String,
|
||||
control_addr: String,
|
||||
player_addr: String,
|
||||
ttl_secs: u64,
|
||||
}
|
||||
|
||||
impl RedisRegistry {
|
||||
async fn from_env(cfg: &RelayConfig) -> Self {
|
||||
let conn = match std::env::var("REDIS_URL") {
|
||||
Ok(url) => match redis::Client::open(url.clone()) {
|
||||
Ok(client) => match redis::aio::ConnectionManager::new(client).await {
|
||||
Ok(cm) => {
|
||||
info!("connected to redis");
|
||||
Some(cm)
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(error = %e, "redis connection manager failed; continuing without redis");
|
||||
None
|
||||
}
|
||||
},
|
||||
Err(e) => {
|
||||
warn!(error = %e, "invalid REDIS_URL; continuing without redis");
|
||||
None
|
||||
}
|
||||
},
|
||||
Err(_) => None,
|
||||
};
|
||||
|
||||
Self {
|
||||
conn,
|
||||
instance_id: cfg.instance_id.clone(),
|
||||
region: cfg.region.clone(),
|
||||
control_addr: cfg.control_bind.clone(),
|
||||
player_addr: cfg.player_bind.clone(),
|
||||
ttl_secs: cfg.registry_ttl_secs,
|
||||
}
|
||||
}
|
||||
|
||||
async fn register_instance(&self) {
|
||||
let Some(mut conn) = self.conn.clone() else {
|
||||
return;
|
||||
};
|
||||
let key = format!("relay:instance:{}", self.instance_id);
|
||||
let payload = serde_json::json!({
|
||||
"instance_id": self.instance_id,
|
||||
"region": self.region,
|
||||
"status": "active",
|
||||
"control_addr": self.control_addr,
|
||||
"player_addr": self.player_addr,
|
||||
"started_at": chrono::Utc::now().timestamp(),
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let res: redis::RedisResult<()> = async {
|
||||
let _: usize = conn.sadd("relay:instances", &self.instance_id).await?;
|
||||
let _: () = conn.set_ex(&key, payload, self.ttl_secs).await?;
|
||||
let hb_key = format!("relay:heartbeat:{}", self.instance_id);
|
||||
let _: () = conn.set_ex(hb_key, "1", self.ttl_secs).await?;
|
||||
Ok(())
|
||||
}
|
||||
.await;
|
||||
if let Err(e) = res {
|
||||
warn!(error = %e, "failed to register instance in redis");
|
||||
}
|
||||
}
|
||||
|
||||
async fn heartbeat_instance(&self, tunnel_count: usize) {
|
||||
let Some(mut conn) = self.conn.clone() else {
|
||||
return;
|
||||
};
|
||||
let hb_key = format!("relay:heartbeat:{}", self.instance_id);
|
||||
let load_key = format!("relay:load:{}", self.region);
|
||||
let score = tunnel_count as f64;
|
||||
let key = format!("relay:instance:{}", self.instance_id);
|
||||
let payload = serde_json::json!({
|
||||
"instance_id": self.instance_id,
|
||||
"region": self.region,
|
||||
"status": "active",
|
||||
"control_addr": self.control_addr,
|
||||
"player_addr": self.player_addr,
|
||||
"tunnel_count": tunnel_count,
|
||||
"updated_at": chrono::Utc::now().timestamp(),
|
||||
})
|
||||
.to_string();
|
||||
let res: redis::RedisResult<()> = async {
|
||||
let _: () = conn.set_ex(hb_key, "1", self.ttl_secs).await?;
|
||||
let _: () = conn.set_ex(key, payload, self.ttl_secs).await?;
|
||||
let _: () = conn.zadd(load_key, &self.instance_id, score).await?;
|
||||
Ok(())
|
||||
}
|
||||
.await;
|
||||
if let Err(e) = res {
|
||||
warn!(error = %e, "redis instance heartbeat failed");
|
||||
}
|
||||
}
|
||||
|
||||
async fn set_draining(&self) {
|
||||
let Some(mut conn) = self.conn.clone() else {
|
||||
return;
|
||||
};
|
||||
let key = format!("relay:instance:{}", self.instance_id);
|
||||
let payload = serde_json::json!({
|
||||
"instance_id": self.instance_id,
|
||||
"region": self.region,
|
||||
"status": "draining",
|
||||
"control_addr": self.control_addr,
|
||||
"player_addr": self.player_addr,
|
||||
"updated_at": chrono::Utc::now().timestamp(),
|
||||
})
|
||||
.to_string();
|
||||
let _: redis::RedisResult<()> = async {
|
||||
let _: () = conn.set_ex(key, payload, self.ttl_secs).await?;
|
||||
let load_key = format!("relay:load:{}", self.region);
|
||||
let _: () = conn.zadd(load_key, &self.instance_id, 1e12f64).await?;
|
||||
Ok(())
|
||||
}
|
||||
.await;
|
||||
}
|
||||
|
||||
async fn register_tunnel(&self, fqdn: &str, session_id: &str, user_id: &str) {
|
||||
let Some(mut conn) = self.conn.clone() else {
|
||||
return;
|
||||
};
|
||||
let key = format!("tunnel:sub:{fqdn}");
|
||||
let session_key = format!("tunnel:session:{session_id}");
|
||||
let payload = serde_json::json!({
|
||||
"instance_id": self.instance_id,
|
||||
"session_id": session_id,
|
||||
"user_id": user_id,
|
||||
"region": self.region,
|
||||
"fqdn": fqdn,
|
||||
})
|
||||
.to_string();
|
||||
let _: redis::RedisResult<()> = async {
|
||||
let _: () = conn.set_ex(key, &payload, self.ttl_secs).await?;
|
||||
let _: () = conn.set_ex(session_key, payload, self.ttl_secs).await?;
|
||||
Ok(())
|
||||
}
|
||||
.await;
|
||||
}
|
||||
|
||||
async fn refresh_tunnel_session(&self, fqdn: &str, session_id: &str, user_id: &str) {
|
||||
self.register_tunnel(fqdn, session_id, user_id).await;
|
||||
}
|
||||
|
||||
async fn remove_tunnel(&self, fqdn: &str, session_id: &str) {
|
||||
let Some(mut conn) = self.conn.clone() else {
|
||||
return;
|
||||
};
|
||||
let _: redis::RedisResult<()> = async {
|
||||
let _: usize = conn.del(format!("tunnel:sub:{fqdn}")).await?;
|
||||
let _: usize = conn.del(format!("tunnel:session:{session_id}")).await?;
|
||||
Ok(())
|
||||
}
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| "relay=info".into()),
|
||||
)
|
||||
.init();
|
||||
|
||||
let cfg = RelayConfig::from_env();
|
||||
let registry = RedisRegistry::from_env(&cfg).await;
|
||||
registry.register_instance().await;
|
||||
|
||||
let control_listener = TcpListener::bind(&cfg.control_bind)
|
||||
.await
|
||||
.with_context(|| format!("bind control {}", cfg.control_bind))?;
|
||||
let player_listener = TcpListener::bind(&cfg.player_bind)
|
||||
.await
|
||||
.with_context(|| format!("bind player {}", cfg.player_bind))?;
|
||||
|
||||
info!(instance_id = %cfg.instance_id, region = %cfg.region, "relay started");
|
||||
|
||||
let shutdown = Arc::new(Notify::new());
|
||||
let state: SharedState = Arc::new(RwLock::new(RelayState::new()));
|
||||
|
||||
let heartbeat_task = tokio::spawn(run_registry_heartbeat(
|
||||
state.clone(),
|
||||
registry.clone(),
|
||||
shutdown.clone(),
|
||||
));
|
||||
let control_task = tokio::spawn(run_control_accept_loop(
|
||||
control_listener,
|
||||
cfg.clone(),
|
||||
state.clone(),
|
||||
registry.clone(),
|
||||
shutdown.clone(),
|
||||
));
|
||||
let player_task = tokio::spawn(run_player_accept_loop(
|
||||
player_listener,
|
||||
state.clone(),
|
||||
shutdown.clone(),
|
||||
));
|
||||
tokio::pin!(heartbeat_task);
|
||||
tokio::pin!(control_task);
|
||||
tokio::pin!(player_task);
|
||||
|
||||
tokio::select! {
|
||||
_ = tokio::signal::ctrl_c() => info!("shutdown signal received"),
|
||||
res = &mut control_task => warn!("control accept loop ended: {:?}", res),
|
||||
res = &mut player_task => warn!("player accept loop ended: {:?}", res),
|
||||
res = &mut heartbeat_task => warn!("registry heartbeat task ended: {:?}", res),
|
||||
}
|
||||
|
||||
registry.set_draining().await;
|
||||
shutdown.notify_waiters();
|
||||
info!("draining relay");
|
||||
tokio::time::sleep(Duration::from_secs(1)).await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn run_registry_heartbeat(state: SharedState, registry: RedisRegistry, shutdown: Arc<Notify>) {
|
||||
let mut ticker = interval(Duration::from_secs(5));
|
||||
ticker.set_missed_tick_behavior(MissedTickBehavior::Skip);
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = shutdown.notified() => break,
|
||||
_ = ticker.tick() => {
|
||||
let count = state.read().await.session_count();
|
||||
registry.heartbeat_instance(count).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_control_accept_loop(
|
||||
listener: TcpListener,
|
||||
cfg: RelayConfig,
|
||||
state: SharedState,
|
||||
registry: RedisRegistry,
|
||||
shutdown: Arc<Notify>,
|
||||
) -> Result<()> {
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = shutdown.notified() => break,
|
||||
res = listener.accept() => {
|
||||
let (stream, addr) = match res {
|
||||
Ok(v) => v,
|
||||
Err(e) => { warn!(error = %e, "control accept failed"); continue; }
|
||||
};
|
||||
let cfg = cfg.clone();
|
||||
let state = state.clone();
|
||||
let registry = registry.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = handle_control_conn(stream, addr, cfg, state, registry).await {
|
||||
warn!(peer = %addr, error = %e, "control connection ended with error");
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn run_player_accept_loop(
|
||||
listener: TcpListener,
|
||||
state: SharedState,
|
||||
shutdown: Arc<Notify>,
|
||||
) -> Result<()> {
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = shutdown.notified() => break,
|
||||
res = listener.accept() => {
|
||||
let (stream, addr) = match res {
|
||||
Ok(v) => v,
|
||||
Err(e) => { warn!(error = %e, "player accept failed"); continue; }
|
||||
};
|
||||
let state = state.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = handle_player_conn(stream, addr, state).await {
|
||||
debug!(peer = %addr, error = %e, "player connection closed");
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_control_conn(
|
||||
stream: TcpStream,
|
||||
addr: SocketAddr,
|
||||
cfg: RelayConfig,
|
||||
state: SharedState,
|
||||
registry: RedisRegistry,
|
||||
) -> Result<()> {
|
||||
let (mut reader, mut writer) = stream.into_split();
|
||||
let first: ClientFrame = read_frame(&mut reader).await.context("read initial frame")?;
|
||||
let register = match first {
|
||||
ClientFrame::Register(req) => req,
|
||||
_ => {
|
||||
write_frame(
|
||||
&mut writer,
|
||||
&ServerFrame::RegisterRejected { reason: "expected Register frame".to_string() },
|
||||
).await.ok();
|
||||
anyhow::bail!("expected register frame");
|
||||
}
|
||||
};
|
||||
|
||||
if !token_looks_valid(®ister.token) {
|
||||
write_frame(
|
||||
&mut writer,
|
||||
&ServerFrame::RegisterRejected { reason: "invalid token".to_string() },
|
||||
).await.ok();
|
||||
anyhow::bail!("invalid token");
|
||||
}
|
||||
|
||||
let (tx, mut rx) = mpsc::channel::<ServerFrame>(512);
|
||||
let session_id = Uuid::new_v4().to_string();
|
||||
let fqdn = assign_fqdn(&cfg, ®ister);
|
||||
let user_id = fake_user_id_from_token(®ister.token);
|
||||
let stream_sinks = Arc::new(RwLock::new(HashMap::<String, mpsc::Sender<Vec<u8>>>::new()));
|
||||
|
||||
{
|
||||
let mut guard = state.write().await;
|
||||
let handle = SessionHandle {
|
||||
session_id: session_id.clone(),
|
||||
tx: tx.clone(),
|
||||
stream_sinks: stream_sinks.clone(),
|
||||
last_heartbeat: Instant::now(),
|
||||
};
|
||||
guard.by_session.insert(session_id.clone(), fqdn.clone());
|
||||
guard.by_fqdn.insert(fqdn.clone(), handle);
|
||||
}
|
||||
registry.register_tunnel(&fqdn, &session_id, &user_id).await;
|
||||
|
||||
write_frame(
|
||||
&mut writer,
|
||||
&ServerFrame::RegisterAccepted(RegisterAccepted {
|
||||
session_id: session_id.clone(),
|
||||
fqdn: fqdn.clone(),
|
||||
heartbeat_interval_secs: 5,
|
||||
owner_instance_id: cfg.instance_id.clone(),
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
info!(peer = %addr, user_id = %user_id, fqdn = %fqdn, session_id = %session_id, "client registered");
|
||||
|
||||
let write_task = tokio::spawn(async move {
|
||||
while let Some(frame) = rx.recv().await {
|
||||
write_frame(&mut writer, &frame).await?;
|
||||
}
|
||||
Ok::<(), anyhow::Error>(())
|
||||
});
|
||||
|
||||
let read_result = control_read_loop(
|
||||
&mut reader,
|
||||
&state,
|
||||
®istry,
|
||||
&session_id,
|
||||
&fqdn,
|
||||
&user_id,
|
||||
cfg.heartbeat_timeout,
|
||||
)
|
||||
.await;
|
||||
|
||||
if let Err(e) = &read_result {
|
||||
warn!(session_id = %session_id, error = %e, "control read loop error");
|
||||
}
|
||||
|
||||
{
|
||||
let mut guard = state.write().await;
|
||||
if let Some(fqdn) = guard.by_session.remove(&session_id) {
|
||||
guard.by_fqdn.remove(&fqdn);
|
||||
}
|
||||
}
|
||||
registry.remove_tunnel(&fqdn, &session_id).await;
|
||||
write_task.abort();
|
||||
info!(session_id = %session_id, "client session removed");
|
||||
read_result
|
||||
}
|
||||
|
||||
async fn control_read_loop(
|
||||
reader: &mut tokio::net::tcp::OwnedReadHalf,
|
||||
state: &SharedState,
|
||||
registry: &RedisRegistry,
|
||||
session_id: &str,
|
||||
fqdn: &str,
|
||||
user_id: &str,
|
||||
heartbeat_timeout: Duration,
|
||||
) -> Result<()> {
|
||||
loop {
|
||||
let frame: ClientFrame = timeout(heartbeat_timeout, read_frame(reader))
|
||||
.await
|
||||
.context("heartbeat timeout")??;
|
||||
match frame {
|
||||
ClientFrame::Heartbeat(Heartbeat { session_id: hb_id, .. }) => {
|
||||
if hb_id != session_id {
|
||||
anyhow::bail!("heartbeat session mismatch");
|
||||
}
|
||||
let mut guard = state.write().await;
|
||||
if let Some(route_fqdn) = guard.by_session.get(session_id).cloned()
|
||||
&& let Some(handle) = guard.by_fqdn.get_mut(&route_fqdn)
|
||||
{
|
||||
handle.last_heartbeat = Instant::now();
|
||||
}
|
||||
drop(guard);
|
||||
registry.refresh_tunnel_session(fqdn, session_id, user_id).await;
|
||||
}
|
||||
ClientFrame::StreamData(StreamData { stream_id, data }) => {
|
||||
let sink = lookup_stream_sink(state, session_id, &stream_id).await;
|
||||
if let Some(tx) = sink {
|
||||
if tx.send(data).await.is_err() {
|
||||
remove_stream_sink(state, session_id, &stream_id).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
ClientFrame::StreamClosed(StreamClosed { stream_id, .. }) => {
|
||||
remove_stream_sink(state, session_id, &stream_id).await;
|
||||
}
|
||||
ClientFrame::Pong => {}
|
||||
ClientFrame::Register(_) => anyhow::bail!("unexpected Register frame after registration"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_player_conn(
|
||||
mut stream: TcpStream,
|
||||
addr: SocketAddr,
|
||||
state: SharedState,
|
||||
) -> Result<()> {
|
||||
let (hostname, initial_data) = read_handshake_hostname_and_bytes(&mut stream)
|
||||
.await
|
||||
.context("parse minecraft handshake")?;
|
||||
|
||||
let session = {
|
||||
let guard = state.read().await;
|
||||
guard.by_fqdn.get(&hostname).cloned()
|
||||
};
|
||||
let Some(session) = session else {
|
||||
debug!(peer = %addr, hostname = %hostname, "no tunnel for hostname");
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let stream_id = Uuid::new_v4().to_string();
|
||||
let (player_read, player_write) = stream.into_split();
|
||||
let (to_player_tx, to_player_rx) = mpsc::channel::<Vec<u8>>(128);
|
||||
session
|
||||
.stream_sinks
|
||||
.write()
|
||||
.await
|
||||
.insert(stream_id.clone(), to_player_tx);
|
||||
|
||||
session
|
||||
.tx
|
||||
.send(ServerFrame::IncomingTcp(IncomingTcp {
|
||||
stream_id: stream_id.clone(),
|
||||
session_id: session.session_id.clone(),
|
||||
peer_addr: addr.to_string(),
|
||||
hostname: hostname.clone(),
|
||||
initial_data,
|
||||
}))
|
||||
.await
|
||||
.context("send IncomingTcp to client")?;
|
||||
|
||||
let tx_control = session.tx.clone();
|
||||
let stream_id_clone = stream_id.clone();
|
||||
let session_id_clone = session.session_id.clone();
|
||||
let sinks = session.stream_sinks.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = run_player_writer(player_write, to_player_rx).await {
|
||||
debug!(stream_id = %stream_id_clone, error = %e, "player writer ended");
|
||||
}
|
||||
let _ = tx_control
|
||||
.send(ServerFrame::StreamClosed(StreamClosed {
|
||||
stream_id: stream_id_clone.clone(),
|
||||
reason: Some("player_writer_closed".into()),
|
||||
}))
|
||||
.await;
|
||||
let _ = remove_stream_sink_by_store(sinks, &stream_id_clone).await;
|
||||
let _ = session_id_clone;
|
||||
});
|
||||
|
||||
let tx_control = session.tx.clone();
|
||||
let stream_id_clone = stream_id.clone();
|
||||
let sinks = session.stream_sinks.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = run_player_reader(player_read, tx_control.clone(), stream_id_clone.clone()).await {
|
||||
debug!(stream_id = %stream_id_clone, error = %e, "player reader ended");
|
||||
}
|
||||
let _ = tx_control
|
||||
.send(ServerFrame::StreamClosed(StreamClosed {
|
||||
stream_id: stream_id_clone.clone(),
|
||||
reason: Some("player_reader_closed".into()),
|
||||
}))
|
||||
.await;
|
||||
let _ = remove_stream_sink_by_store(sinks, &stream_id_clone).await;
|
||||
});
|
||||
|
||||
info!(peer = %addr, hostname = %hostname, session_id = %session.session_id, stream_id = %stream_id, "player proxied via client stream");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn run_player_reader(
|
||||
mut reader: tokio::net::tcp::OwnedReadHalf,
|
||||
tx_control: mpsc::Sender<ServerFrame>,
|
||||
stream_id: String,
|
||||
) -> Result<()> {
|
||||
let mut buf = vec![0u8; 16 * 1024];
|
||||
loop {
|
||||
let n = reader.read(&mut buf).await?;
|
||||
if n == 0 {
|
||||
break;
|
||||
}
|
||||
tx_control
|
||||
.send(ServerFrame::StreamData(StreamData {
|
||||
stream_id: stream_id.clone(),
|
||||
data: buf[..n].to_vec(),
|
||||
}))
|
||||
.await
|
||||
.context("send stream data to client")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn run_player_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(())
|
||||
}
|
||||
|
||||
async fn lookup_stream_sink(
|
||||
state: &SharedState,
|
||||
session_id: &str,
|
||||
stream_id: &str,
|
||||
) -> Option<mpsc::Sender<Vec<u8>>> {
|
||||
let store = {
|
||||
let guard = state.read().await;
|
||||
let fqdn = guard.by_session.get(session_id)?.clone();
|
||||
guard.by_fqdn.get(&fqdn)?.stream_sinks.clone()
|
||||
};
|
||||
store.read().await.get(stream_id).cloned()
|
||||
}
|
||||
|
||||
async fn remove_stream_sink(state: &SharedState, session_id: &str, stream_id: &str) {
|
||||
let store = {
|
||||
let guard = state.read().await;
|
||||
let Some(fqdn) = guard.by_session.get(session_id).cloned() else {
|
||||
return;
|
||||
};
|
||||
let Some(handle) = guard.by_fqdn.get(&fqdn) else {
|
||||
return;
|
||||
};
|
||||
handle.stream_sinks.clone()
|
||||
};
|
||||
let _ = remove_stream_sink_by_store(store, stream_id).await;
|
||||
}
|
||||
|
||||
async fn remove_stream_sink_by_store(
|
||||
store: Arc<RwLock<HashMap<String, mpsc::Sender<Vec<u8>>>>>,
|
||||
stream_id: &str,
|
||||
) -> Option<mpsc::Sender<Vec<u8>>> {
|
||||
store.write().await.remove(stream_id)
|
||||
}
|
||||
|
||||
fn token_looks_valid(token: &str) -> bool {
|
||||
!token.trim().is_empty()
|
||||
}
|
||||
|
||||
fn fake_user_id_from_token(token: &str) -> String {
|
||||
let suffix: String = token.chars().rev().take(6).collect();
|
||||
format!("user-{}", suffix.chars().rev().collect::<String>())
|
||||
}
|
||||
|
||||
fn assign_fqdn(cfg: &RelayConfig, req: &RegisterRequest) -> String {
|
||||
let label = req
|
||||
.requested_subdomain
|
||||
.as_ref()
|
||||
.filter(|s| !s.trim().is_empty())
|
||||
.cloned()
|
||||
.unwrap_or_else(random_label);
|
||||
format!("{}.{}.{}", sanitize_label(&label), cfg.region, cfg.domain)
|
||||
}
|
||||
|
||||
fn random_label() -> String {
|
||||
const ADJ: &[&str] = &["sleepy", "swift", "brave", "quiet", "mossy"];
|
||||
const NOUN: &[&str] = &["creeper", "ghast", "axolotl", "wolf", "beacon"];
|
||||
format!(
|
||||
"{}-{}-{}",
|
||||
ADJ[fastrand::usize(..ADJ.len())],
|
||||
NOUN[fastrand::usize(..NOUN.len())],
|
||||
fastrand::u16(..9999)
|
||||
)
|
||||
}
|
||||
|
||||
fn sanitize_label(input: &str) -> String {
|
||||
input
|
||||
.chars()
|
||||
.filter(|c| c.is_ascii_alphanumeric() || *c == '-')
|
||||
.collect::<String>()
|
||||
.trim_matches('-')
|
||||
.to_ascii_lowercase()
|
||||
}
|
||||
Reference in New Issue
Block a user