swaits
90a306853d
UART boards (CP2102 bridge) keep the serial port alive across ESP32 resets, so the mux never detected the board rebooting. Now the writer sends a keepalive Ping every 5s of idle and declares the dongle lost after a 2s timeout. On reconnect, saved radio config and RX state are re-sent to the dongle so connected clients continue working without re-establishing state.
712 lines
25 KiB
Rust
712 lines
25 KiB
Rust
//! Core mux daemon: dongle I/O, client management, reconnect loop.
|
|
//!
|
|
//! The daemon owns the USB serial connection and exposes Unix domain socket
|
|
//! and optional TCP listeners. It routes commands from clients to the dongle
|
|
//! and broadcasts/routes responses back.
|
|
|
|
use std::collections::HashMap;
|
|
use std::path::Path;
|
|
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
|
|
use donglora_client::protocol::{
|
|
CMD_TAG_SET_CONFIG, CMD_TAG_START_RX, CMD_TAG_STOP_RX, RADIO_CONFIG_SIZE, RESP_TAG_OK, RESP_TAG_RX_PACKET,
|
|
};
|
|
use donglora_client::{FrameReader, encode_frame};
|
|
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
|
use tokio::net::{TcpListener, UnixListener};
|
|
use tokio::sync::{Mutex, mpsc, oneshot};
|
|
|
|
/// A queued command: (client_id, raw_cmd_bytes).
|
|
type QueuedCmd = (u64, Vec<u8>);
|
|
|
|
/// Shared command receiver, wrapped for use across reconnect cycles.
|
|
type SharedCmdRx = Arc<Mutex<mpsc::Receiver<QueuedCmd>>>;
|
|
use tokio_util::sync::CancellationToken;
|
|
use tracing::{debug, error, info, warn};
|
|
|
|
use crate::intercept::{self, MuxState};
|
|
use crate::session::ClientSession;
|
|
|
|
/// Command timeout waiting for dongle response.
|
|
const COMMAND_TIMEOUT: Duration = Duration::from_secs(10);
|
|
|
|
/// Delay between dongle reconnect attempts.
|
|
const RECONNECT_DELAY: Duration = Duration::from_secs(2);
|
|
|
|
/// Max consecutive dongle read errors before declaring disconnect.
|
|
const MAX_CONSECUTIVE_ERRORS: u32 = 3;
|
|
|
|
/// Sentinel client ID for synthetic commands (e.g. auto StopRx).
|
|
const SYNTHETIC_CLIENT_ID: u64 = u64::MAX;
|
|
|
|
/// Keepalive ping interval: if no commands flow for this long, ping the dongle.
|
|
const KEEPALIVE_INTERVAL: Duration = Duration::from_secs(5);
|
|
|
|
/// How long to wait for a keepalive Pong before declaring the dongle lost.
|
|
const KEEPALIVE_TIMEOUT: Duration = Duration::from_secs(2);
|
|
|
|
/// Shared mux state, protected by async mutex.
|
|
type SharedState = Arc<Mutex<MuxInner>>;
|
|
|
|
struct MuxInner {
|
|
sessions: HashMap<u64, ClientSession>,
|
|
mux_state: MuxState,
|
|
}
|
|
|
|
/// The multiplexer daemon.
|
|
pub struct MuxDaemon {
|
|
serial_port: String,
|
|
socket_path: String,
|
|
tcp_addr: Option<(String, u16)>,
|
|
shutdown: CancellationToken,
|
|
}
|
|
|
|
impl MuxDaemon {
|
|
pub fn new(
|
|
serial_port: String,
|
|
socket_path: String,
|
|
tcp_addr: Option<(String, u16)>,
|
|
shutdown: CancellationToken,
|
|
) -> Self {
|
|
Self { serial_port, socket_path, tcp_addr, shutdown }
|
|
}
|
|
|
|
/// Run the mux daemon. Returns when shutdown is signalled.
|
|
pub async fn run(&self) -> anyhow::Result<()> {
|
|
let state: SharedState =
|
|
Arc::new(Mutex::new(MuxInner { sessions: HashMap::new(), mux_state: MuxState::new() }));
|
|
|
|
// Command channel: client handlers → dongle writer
|
|
let (cmd_tx, cmd_rx) = mpsc::channel::<(u64, Vec<u8>)>(64);
|
|
let cmd_rx = Arc::new(Mutex::new(cmd_rx));
|
|
|
|
// Pending response slot: dongle writer installs, dongle reader resolves
|
|
let pending_response: Arc<Mutex<Option<oneshot::Sender<Vec<u8>>>>> = Arc::new(Mutex::new(None));
|
|
|
|
// Prepare socket directory
|
|
let sock_path = Path::new(&self.socket_path);
|
|
if let Some(parent) = sock_path.parent() {
|
|
tokio::fs::create_dir_all(parent).await.ok();
|
|
}
|
|
// Acquire an exclusive lock to prevent multiple mux instances.
|
|
// The lock is held for the process lifetime and auto-released on exit (even SIGKILL).
|
|
let lock_path = format!("{}.lock", self.socket_path);
|
|
let lock_file = std::fs::OpenOptions::new()
|
|
.create(true)
|
|
.write(true)
|
|
.truncate(false)
|
|
.open(&lock_path)
|
|
.map_err(|e| anyhow::anyhow!("failed to open lock file {lock_path}: {e}"))?;
|
|
let mut lock = fd_lock::RwLock::new(lock_file);
|
|
let _lock_guard = match lock.try_write() {
|
|
Ok(guard) => guard,
|
|
Err(_) => anyhow::bail!("another donglora-mux is already running (lock held on {lock_path})"),
|
|
};
|
|
|
|
if sock_path.exists() {
|
|
// Lock acquired but socket file exists — stale from a crashed instance.
|
|
info!("removing stale socket {}", self.socket_path);
|
|
tokio::fs::remove_file(sock_path).await.ok();
|
|
}
|
|
|
|
// Start Unix socket server
|
|
let unix_listener = UnixListener::bind(&self.socket_path)
|
|
.map_err(|e| anyhow::anyhow!("failed to bind Unix socket {}: {e}", self.socket_path))?;
|
|
info!("Unix socket listening on {}", self.socket_path);
|
|
|
|
// Start TCP server (optional)
|
|
let tcp_listener = if let Some((ref host, port)) = self.tcp_addr {
|
|
let addr = format!("{host}:{port}");
|
|
let listener = TcpListener::bind(&addr).await.map_err(|e| {
|
|
if e.kind() == std::io::ErrorKind::AddrInUse {
|
|
anyhow::anyhow!("another donglora-mux is already listening on TCP {addr}")
|
|
} else {
|
|
anyhow::anyhow!("failed to bind TCP {addr}: {e}")
|
|
}
|
|
})?;
|
|
info!("TCP listening on {addr}");
|
|
Some(listener)
|
|
} else {
|
|
None
|
|
};
|
|
|
|
// Spawn client accept tasks
|
|
tokio::spawn(accept_unix_clients(unix_listener, state.clone(), cmd_tx.clone(), self.shutdown.clone()));
|
|
if let Some(listener) = tcp_listener {
|
|
tokio::spawn(accept_tcp_clients(listener, state.clone(), cmd_tx.clone(), self.shutdown.clone()));
|
|
}
|
|
|
|
// Connect/reconnect loop
|
|
self.reconnect_loop(state, cmd_tx, cmd_rx, pending_response).await;
|
|
|
|
// Cleanup
|
|
info!("shutting down...");
|
|
if Path::new(&self.socket_path).exists() {
|
|
tokio::fs::remove_file(&self.socket_path).await.ok();
|
|
}
|
|
info!("stopped.");
|
|
Ok(())
|
|
}
|
|
|
|
async fn reconnect_loop(
|
|
&self,
|
|
state: SharedState,
|
|
cmd_tx: mpsc::Sender<(u64, Vec<u8>)>,
|
|
cmd_rx: SharedCmdRx,
|
|
pending_response: Arc<Mutex<Option<oneshot::Sender<Vec<u8>>>>>,
|
|
) {
|
|
let mut saved_config: Option<[u8; RADIO_CONFIG_SIZE]> = None;
|
|
let mut had_rx_interest = false;
|
|
|
|
loop {
|
|
if self.shutdown.is_cancelled() {
|
|
break;
|
|
}
|
|
|
|
info!("opening dongle on {}", self.serial_port);
|
|
let (serial, _serial_lock_guard) = loop {
|
|
if self.shutdown.is_cancelled() {
|
|
return;
|
|
}
|
|
match open_and_ping(&self.serial_port).await {
|
|
Ok(pair) => break pair,
|
|
Err(e) => {
|
|
debug!("could not open dongle: {e}");
|
|
tokio::select! {
|
|
() = self.shutdown.cancelled() => return,
|
|
() = tokio::time::sleep(RECONNECT_DELAY) => continue,
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
info!("dongle connected");
|
|
let (serial_read, serial_write) = tokio::io::split(serial);
|
|
|
|
let dongle_lost = CancellationToken::new();
|
|
|
|
let reader_handle = tokio::spawn(dongle_reader(
|
|
serial_read,
|
|
state.clone(),
|
|
pending_response.clone(),
|
|
dongle_lost.clone(),
|
|
self.shutdown.clone(),
|
|
));
|
|
|
|
let writer_handle = tokio::spawn(dongle_writer(
|
|
serial_write,
|
|
cmd_rx.clone(),
|
|
state.clone(),
|
|
pending_response.clone(),
|
|
dongle_lost.clone(),
|
|
self.shutdown.clone(),
|
|
));
|
|
|
|
// Re-establish radio state from previous session
|
|
if let Some(config) = saved_config.take() {
|
|
info!("restoring radio config after reconnect");
|
|
let mut cmd = vec![CMD_TAG_SET_CONFIG];
|
|
cmd.extend_from_slice(&config);
|
|
let _ = cmd_tx.send((SYNTHETIC_CLIENT_ID, cmd)).await;
|
|
if had_rx_interest {
|
|
let _ = cmd_tx.send((SYNTHETIC_CLIENT_ID, vec![CMD_TAG_START_RX])).await;
|
|
}
|
|
}
|
|
|
|
// Wait for dongle loss or shutdown
|
|
tokio::select! {
|
|
() = self.shutdown.cancelled() => {},
|
|
() = dongle_lost.cancelled() => {},
|
|
}
|
|
|
|
reader_handle.abort();
|
|
writer_handle.abort();
|
|
let _ = reader_handle.await;
|
|
let _ = writer_handle.await;
|
|
|
|
// Drain pending commands
|
|
{
|
|
let mut rx = cmd_rx.lock().await;
|
|
while rx.try_recv().is_ok() {}
|
|
}
|
|
|
|
// Save radio state for restoration after reconnect
|
|
{
|
|
let mut inner = state.lock().await;
|
|
saved_config = inner.mux_state.locked_config.take();
|
|
had_rx_interest = intercept::rx_interest_count(&inner.sessions) > 0;
|
|
}
|
|
|
|
if !self.shutdown.is_cancelled() {
|
|
info!("reconnecting in 2 seconds...");
|
|
tokio::select! {
|
|
() = self.shutdown.cancelled() => return,
|
|
() = tokio::time::sleep(RECONNECT_DELAY) => {},
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// ── Dongle I/O ─────────────────────────────────────────────────────
|
|
|
|
/// Lock file path for exclusive serial port access.
|
|
fn serial_lock_path(port: &str) -> String {
|
|
format!("/tmp/donglora-serial-{}.lock", port.replace('/', "_"))
|
|
}
|
|
|
|
/// Open the serial port with an exclusive lock and verify the dongle responds to ping.
|
|
///
|
|
/// Returns the serial stream and the lock. The caller must keep `_serial_lock` alive
|
|
/// for the duration of the dongle session. Dropping it releases the exclusive flock,
|
|
/// allowing other processes to claim the port.
|
|
async fn open_and_ping(
|
|
port: &str,
|
|
) -> anyhow::Result<(tokio_serial::SerialStream, fd_lock::RwLock<std::fs::File>)> {
|
|
// Acquire exclusive flock before opening the serial port. This prevents bridges
|
|
// or other processes from directly opening the same device while the mux owns it.
|
|
let lock_path = serial_lock_path(port);
|
|
let lock_file = std::fs::OpenOptions::new()
|
|
.create(true)
|
|
.write(true)
|
|
.truncate(false)
|
|
.open(&lock_path)
|
|
.map_err(|e| anyhow::anyhow!("failed to open serial lock file {lock_path}: {e}"))?;
|
|
let mut lock = fd_lock::RwLock::new(lock_file);
|
|
// Acquire exclusive flock (non-blocking). Fails if another process holds it.
|
|
let guard = lock
|
|
.try_write()
|
|
.map_err(|_| anyhow::anyhow!("serial port {port} is locked by another process"))?;
|
|
// Forget the guard so its Drop doesn't call flock(UNLOCK). The underlying flock
|
|
// stays held as long as the file descriptor (owned by `lock`) remains open.
|
|
// When `lock` is dropped at the end of the dongle session, the fd closes and
|
|
// the OS releases the flock automatically.
|
|
std::mem::forget(guard);
|
|
|
|
let builder = tokio_serial::new(port, 115_200).timeout(Duration::from_secs(2));
|
|
let mut serial =
|
|
tokio_serial::SerialStream::open(&builder).map_err(|e| anyhow::anyhow!("failed to open {port}: {e}"))?;
|
|
|
|
let ping_frame = encode_frame(&[0]);
|
|
AsyncWriteExt::write_all(&mut serial, &ping_frame)
|
|
.await
|
|
.map_err(|e| anyhow::anyhow!("failed to write ping: {e}"))?;
|
|
|
|
let mut reader = FrameReader::new();
|
|
let mut buf = [0u8; 64];
|
|
let deadline = tokio::time::Instant::now() + Duration::from_secs(2);
|
|
|
|
loop {
|
|
let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
|
|
if remaining.is_zero() {
|
|
anyhow::bail!("ping timeout");
|
|
}
|
|
|
|
let n = tokio::select! {
|
|
result = serial.read(&mut buf) => {
|
|
result.map_err(|e| anyhow::anyhow!("serial read error: {e}"))?
|
|
}
|
|
() = tokio::time::sleep(remaining) => {
|
|
anyhow::bail!("ping timeout");
|
|
}
|
|
};
|
|
|
|
if n == 0 {
|
|
anyhow::bail!("serial port closed during ping");
|
|
}
|
|
|
|
for frame in reader.feed(&buf[..n]) {
|
|
if frame.first() == Some(&0) {
|
|
info!("dongle responded to Ping");
|
|
return Ok((serial, lock));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn dongle_reader(
|
|
mut serial: tokio::io::ReadHalf<tokio_serial::SerialStream>,
|
|
state: SharedState,
|
|
pending_response: Arc<Mutex<Option<oneshot::Sender<Vec<u8>>>>>,
|
|
dongle_lost: CancellationToken,
|
|
shutdown: CancellationToken,
|
|
) {
|
|
let mut reader = FrameReader::new();
|
|
let mut buf = [0u8; 512];
|
|
let mut consecutive_errors: u32 = 0;
|
|
|
|
loop {
|
|
if shutdown.is_cancelled() || dongle_lost.is_cancelled() {
|
|
return;
|
|
}
|
|
|
|
let read_result = tokio::select! {
|
|
result = serial.read(&mut buf) => result,
|
|
() = shutdown.cancelled() => return,
|
|
() = dongle_lost.cancelled() => return,
|
|
};
|
|
|
|
match read_result {
|
|
Ok(0) => {
|
|
error!("serial port closed");
|
|
resolve_pending_with_error(&pending_response).await;
|
|
dongle_lost.cancel();
|
|
return;
|
|
}
|
|
Ok(n) => {
|
|
consecutive_errors = 0;
|
|
for frame in reader.feed(&buf[..n]) {
|
|
dispatch_frame(&frame, &state, &pending_response).await;
|
|
}
|
|
}
|
|
Err(e) => {
|
|
consecutive_errors += 1;
|
|
if consecutive_errors >= MAX_CONSECUTIVE_ERRORS {
|
|
error!("dongle read error (giving up after {consecutive_errors}): {e}");
|
|
resolve_pending_with_error(&pending_response).await;
|
|
dongle_lost.cancel();
|
|
return;
|
|
}
|
|
warn!("dongle read glitch ({consecutive_errors}/{MAX_CONSECUTIVE_ERRORS}): {e}");
|
|
tokio::time::sleep(Duration::from_millis(100)).await;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn dispatch_frame(
|
|
frame: &[u8],
|
|
state: &SharedState,
|
|
pending_response: &Arc<Mutex<Option<oneshot::Sender<Vec<u8>>>>>,
|
|
) {
|
|
let Some(&tag) = frame.first() else {
|
|
return;
|
|
};
|
|
|
|
if tag == RESP_TAG_RX_PACKET {
|
|
let cobs_frame = encode_frame(frame);
|
|
let inner = state.lock().await;
|
|
for session in inner.sessions.values() {
|
|
session.enqueue(cobs_frame.clone());
|
|
}
|
|
} else {
|
|
let mut slot = pending_response.lock().await;
|
|
if let Some(sender) = slot.take() {
|
|
let _ = sender.send(frame.to_vec());
|
|
} else {
|
|
warn!("unsolicited non-RxPacket response (tag {tag}) — dropped");
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn resolve_pending_with_error(pending_response: &Arc<Mutex<Option<oneshot::Sender<Vec<u8>>>>>) {
|
|
let mut slot = pending_response.lock().await;
|
|
if let Some(sender) = slot.take() {
|
|
let _ = sender.send(vec![5, 1]); // Error(RadioBusy)
|
|
}
|
|
}
|
|
|
|
/// Send a Ping and wait for Pong. Returns `true` if the dongle responded.
|
|
async fn keepalive_ping(
|
|
serial: &mut tokio::io::WriteHalf<tokio_serial::SerialStream>,
|
|
pending_response: &Arc<Mutex<Option<oneshot::Sender<Vec<u8>>>>>,
|
|
) -> bool {
|
|
let (resp_tx, resp_rx) = oneshot::channel();
|
|
{
|
|
let mut slot = pending_response.lock().await;
|
|
*slot = Some(resp_tx);
|
|
}
|
|
|
|
let ping_frame = encode_frame(&[0]);
|
|
if AsyncWriteExt::write_all(serial, &ping_frame).await.is_err() {
|
|
return false;
|
|
}
|
|
|
|
match tokio::time::timeout(KEEPALIVE_TIMEOUT, resp_rx).await {
|
|
Ok(Ok(_)) => true,
|
|
_ => {
|
|
let mut slot = pending_response.lock().await;
|
|
*slot = None;
|
|
false
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn dongle_writer(
|
|
mut serial: tokio::io::WriteHalf<tokio_serial::SerialStream>,
|
|
cmd_rx: SharedCmdRx,
|
|
state: SharedState,
|
|
pending_response: Arc<Mutex<Option<oneshot::Sender<Vec<u8>>>>>,
|
|
dongle_lost: CancellationToken,
|
|
shutdown: CancellationToken,
|
|
) {
|
|
let keepalive = tokio::time::sleep(KEEPALIVE_INTERVAL);
|
|
tokio::pin!(keepalive);
|
|
|
|
loop {
|
|
// Pull next command, or fire keepalive if idle too long
|
|
let cmd = {
|
|
let mut rx = cmd_rx.lock().await;
|
|
tokio::select! {
|
|
cmd = rx.recv() => cmd,
|
|
() = &mut keepalive => None,
|
|
() = shutdown.cancelled() => return,
|
|
() = dongle_lost.cancelled() => return,
|
|
}
|
|
};
|
|
|
|
keepalive.as_mut().reset(tokio::time::Instant::now() + KEEPALIVE_INTERVAL);
|
|
|
|
let (client_id, raw_cmd) = match cmd {
|
|
Some(c) => c,
|
|
None => {
|
|
debug!("keepalive ping");
|
|
if !keepalive_ping(&mut serial, &pending_response).await {
|
|
error!("keepalive ping failed - dongle lost");
|
|
dongle_lost.cancel();
|
|
return;
|
|
}
|
|
continue;
|
|
}
|
|
};
|
|
|
|
// Check for interception
|
|
let intercepted = {
|
|
let mut inner = state.lock().await;
|
|
let locked = inner.mux_state.locked_config;
|
|
intercept::maybe_intercept(&raw_cmd, client_id, &mut inner.sessions, &locked)
|
|
};
|
|
if let Some(synthetic) = intercepted {
|
|
let inner = state.lock().await;
|
|
if let Some(session) = inner.sessions.get(&client_id) {
|
|
session.enqueue(encode_frame(&synthetic));
|
|
}
|
|
continue;
|
|
}
|
|
|
|
// Install oneshot for the response
|
|
let (resp_tx, resp_rx) = oneshot::channel();
|
|
{
|
|
let mut slot = pending_response.lock().await;
|
|
*slot = Some(resp_tx);
|
|
}
|
|
|
|
// Send command to dongle
|
|
let cobs_frame = encode_frame(&raw_cmd);
|
|
if let Err(e) = AsyncWriteExt::write_all(&mut serial, &cobs_frame).await {
|
|
error!("dongle write error: {e}");
|
|
resolve_pending_with_error(&pending_response).await;
|
|
dongle_lost.cancel();
|
|
return;
|
|
}
|
|
|
|
// Wait for response
|
|
let response = match tokio::time::timeout(COMMAND_TIMEOUT, resp_rx).await {
|
|
Ok(Ok(resp)) => resp,
|
|
Ok(Err(_)) => {
|
|
warn!("response channel dropped");
|
|
vec![5, 1]
|
|
}
|
|
Err(_) => {
|
|
error!("dongle response timeout");
|
|
let mut slot = pending_response.lock().await;
|
|
*slot = None;
|
|
vec![5, 1]
|
|
}
|
|
};
|
|
|
|
// Track state on successful forward
|
|
let cmd_tag = raw_cmd.first().copied().unwrap_or(0xFF);
|
|
let resp_tag = response.first().copied().unwrap_or(0xFF);
|
|
{
|
|
let mut inner = state.lock().await;
|
|
if cmd_tag == CMD_TAG_SET_CONFIG && resp_tag == RESP_TAG_OK {
|
|
if let Some(config_bytes) = raw_cmd.get(1..1 + RADIO_CONFIG_SIZE) {
|
|
let mut locked = [0u8; RADIO_CONFIG_SIZE];
|
|
locked.copy_from_slice(config_bytes);
|
|
inner.mux_state.locked_config = Some(locked);
|
|
info!("radio config locked");
|
|
}
|
|
} else if cmd_tag == CMD_TAG_START_RX && resp_tag == RESP_TAG_OK {
|
|
if let Some(session) = inner.sessions.get_mut(&client_id) {
|
|
session.rx_interested = true;
|
|
}
|
|
} else if cmd_tag == CMD_TAG_STOP_RX
|
|
&& resp_tag == RESP_TAG_OK
|
|
&& let Some(session) = inner.sessions.get_mut(&client_id)
|
|
{
|
|
session.rx_interested = false;
|
|
}
|
|
}
|
|
|
|
// Route response to commanding client
|
|
if client_id != SYNTHETIC_CLIENT_ID {
|
|
let inner = state.lock().await;
|
|
if let Some(session) = inner.sessions.get(&client_id) {
|
|
session.enqueue(encode_frame(&response));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// ── Client management ──────────────────────────────────────────────
|
|
|
|
async fn accept_unix_clients(
|
|
listener: UnixListener,
|
|
state: SharedState,
|
|
cmd_tx: mpsc::Sender<(u64, Vec<u8>)>,
|
|
shutdown: CancellationToken,
|
|
) {
|
|
loop {
|
|
let accept_result = tokio::select! {
|
|
result = listener.accept() => result,
|
|
() = shutdown.cancelled() => return,
|
|
};
|
|
|
|
match accept_result {
|
|
Ok((stream, _addr)) => {
|
|
let (read_half, write_half) = stream.into_split();
|
|
spawn_client(read_half, write_half, state.clone(), cmd_tx.clone(), shutdown.clone());
|
|
}
|
|
Err(e) => warn!("Unix accept error: {e}"),
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn accept_tcp_clients(
|
|
listener: TcpListener,
|
|
state: SharedState,
|
|
cmd_tx: mpsc::Sender<(u64, Vec<u8>)>,
|
|
shutdown: CancellationToken,
|
|
) {
|
|
loop {
|
|
let accept_result = tokio::select! {
|
|
result = listener.accept() => result,
|
|
() = shutdown.cancelled() => return,
|
|
};
|
|
|
|
match accept_result {
|
|
Ok((stream, addr)) => {
|
|
debug!("TCP connection from {addr}");
|
|
// Low-latency: disable Nagle (COBS frames are small)
|
|
let _ = stream.set_nodelay(true);
|
|
let (read_half, write_half) = stream.into_split();
|
|
spawn_client(read_half, write_half, state.clone(), cmd_tx.clone(), shutdown.clone());
|
|
}
|
|
Err(e) => warn!("TCP accept error: {e}"),
|
|
}
|
|
}
|
|
}
|
|
|
|
fn spawn_client<R, W>(
|
|
read_half: R,
|
|
write_half: W,
|
|
state: SharedState,
|
|
cmd_tx: mpsc::Sender<(u64, Vec<u8>)>,
|
|
shutdown: CancellationToken,
|
|
) where
|
|
R: AsyncReadExt + Unpin + Send + 'static,
|
|
W: AsyncWriteExt + Unpin + Send + 'static,
|
|
{
|
|
let (session, send_rx) = ClientSession::new();
|
|
let client_id = session.id;
|
|
let label = session.label();
|
|
|
|
let state_clone = state.clone();
|
|
let cmd_tx_clone = cmd_tx.clone();
|
|
|
|
tokio::spawn(async move {
|
|
{
|
|
let mut inner = state_clone.lock().await;
|
|
inner.sessions.insert(client_id, session);
|
|
info!("{label} connected ({} total)", inner.sessions.len());
|
|
}
|
|
|
|
let writer_shutdown = shutdown.clone();
|
|
let writer_handle = tokio::spawn(client_writer(write_half, send_rx, writer_shutdown));
|
|
|
|
client_reader(read_half, client_id, cmd_tx_clone.clone(), shutdown).await;
|
|
|
|
writer_handle.abort();
|
|
let _ = writer_handle.await;
|
|
|
|
remove_client(client_id, &state_clone, &label, &cmd_tx_clone).await;
|
|
});
|
|
}
|
|
|
|
async fn client_reader<R: AsyncReadExt + Unpin>(
|
|
mut read_half: R,
|
|
client_id: u64,
|
|
cmd_tx: mpsc::Sender<(u64, Vec<u8>)>,
|
|
shutdown: CancellationToken,
|
|
) {
|
|
let mut reader = FrameReader::new();
|
|
let mut buf = [0u8; 4096];
|
|
|
|
loop {
|
|
let read_result = tokio::select! {
|
|
result = read_half.read(&mut buf) => result,
|
|
() = shutdown.cancelled() => return,
|
|
};
|
|
|
|
match read_result {
|
|
Ok(0) => return,
|
|
Ok(n) => {
|
|
for frame in reader.feed(&buf[..n]) {
|
|
if cmd_tx.send((client_id, frame)).await.is_err() {
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
Err(_) => return,
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn client_writer<W: AsyncWriteExt + Unpin>(
|
|
mut write_half: W,
|
|
mut send_rx: mpsc::Receiver<Vec<u8>>,
|
|
shutdown: CancellationToken,
|
|
) {
|
|
loop {
|
|
let frame = tokio::select! {
|
|
f = send_rx.recv() => {
|
|
match f {
|
|
Some(f) => f,
|
|
None => return,
|
|
}
|
|
}
|
|
() = shutdown.cancelled() => return,
|
|
};
|
|
|
|
if AsyncWriteExt::write_all(&mut write_half, &frame).await.is_err() {
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn remove_client(client_id: u64, state: &SharedState, label: &str, cmd_tx: &mpsc::Sender<(u64, Vec<u8>)>) {
|
|
let mut inner = state.lock().await;
|
|
|
|
let was_interested = inner.sessions.get(&client_id).is_some_and(|s| s.rx_interested);
|
|
|
|
inner.sessions.remove(&client_id);
|
|
info!("{label} disconnected ({} remain)", inner.sessions.len());
|
|
|
|
// If last RX-interested client gone, send synthetic StopRx
|
|
if was_interested && intercept::rx_interest_count(&inner.sessions) == 0 {
|
|
info!("last RX-interested client gone — sending StopRx");
|
|
// Drop the lock before sending to avoid deadlock
|
|
drop(inner);
|
|
let _ = cmd_tx.send((SYNTHETIC_CLIENT_ID, vec![CMD_TAG_STOP_RX])).await;
|
|
return;
|
|
}
|
|
|
|
// Reset locked config when all clients are gone
|
|
if inner.sessions.is_empty() && inner.mux_state.locked_config.is_some() {
|
|
info!("all clients gone — config lock released");
|
|
inner.mux_state.locked_config = None;
|
|
}
|
|
}
|