use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use anyhow::Context;
use clap::Args as ClapArgs;
use iroh::EndpointId;
use tokio::signal::unix::{signal, SignalKind};
use tokio_util::sync::CancellationToken;
use crate::server::audit::{auth_event, Outcome};
use crate::server::policy::{load_allow_file, Policy};
use crate::server::{run_attached, session, SessionExit};
use crate::transport_iroh::{
bind_endpoint, bind_endpoint_local, bind_endpoint_with_relay, format_endpoint_id,
load_or_create_secret_key, parse_endpoint_id, parse_relay_url, ALPN,
};
use tracing::{error, info, warn};
const ACCEPT_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);
#[derive(ClapArgs, Debug)]
pub struct ServeArgs {
#[arg(long)]
key_file: Option<PathBuf>,
#[arg(long = "allow", value_name = "ENDPOINT_ID")]
allow: Vec<String>,
#[arg(long)]
allow_any: bool,
#[arg(long, value_name = "PATH")]
allow_file: Option<PathBuf>,
#[arg(long)]
read_only: bool,
#[arg(long)]
shell: Option<String>,
#[arg(long, default_value_t = 1000)]
scrollback: usize,
#[arg(long, default_value_t = 86_400)]
session_ttl_secs: u64,
#[arg(long, env = "KOH_SERVER_NETWORK_TMOUT", default_value_t = 0)]
network_timeout_secs: u64,
#[arg(long, value_name = "URL")]
relay_url: Option<String>,
#[arg(long, conflicts_with = "relay_url")]
local: bool,
#[arg(long, default_value_t = 64, value_parser = clap::value_parser!(u32).range(1..))]
max_connections: u32,
#[arg(long, default_value_t = 64, value_parser = clap::value_parser!(u32).range(1..))]
max_sessions: u32,
}
fn connect_qr(data: &str) -> Option<String> {
use qrcode::render::unicode::Dense1x2;
let code = qrcode::QrCode::new(data).ok()?;
Some(
code.render::<Dense1x2>()
.dark_color(Dense1x2::Light)
.light_color(Dense1x2::Dark)
.quiet_zone(true)
.build(),
)
}
fn default_key_file() -> PathBuf {
crate::transport_iroh::default_key_path("server")
}
pub async fn serve(args: ServeArgs) -> anyhow::Result<()> {
tracing_subscriber::fmt()
.with_env_filter(
tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| "koh=info".into()),
)
.with_writer(std::io::stderr)
.init();
let mut allow: HashMap<EndpointId, Policy> = HashMap::new();
for s in &args.allow {
let id = parse_endpoint_id(s).with_context(|| format!("bad --allow id: {s}"))?;
allow.insert(
id,
Policy {
read_only: args.read_only,
force_command: None,
},
);
}
if let Some(path) = &args.allow_file {
for (id, mut policy) in load_allow_file(path)? {
policy.read_only |= args.read_only;
if allow.insert(id, policy).is_some() {
anyhow::bail!(
"endpoint id authorized by both --allow and --allow-file: {}",
format_endpoint_id(&id)
);
}
}
}
if allow.is_empty() && !args.allow_any {
anyhow::bail!(
"no clients authorized: pass --allow <endpoint-id> (repeatable), --allow-file <path>, or --allow-any for testing"
);
}
let key_file = args.key_file.clone().unwrap_or_else(default_key_file);
let secret = load_or_create_secret_key(&key_file).with_context(|| {
format!(
"loading server key from {} (pass --key-file to use a writable path)",
key_file.display()
)
})?;
let endpoint = if let Some(url) = &args.relay_url {
let relay = parse_relay_url(url)?;
bind_endpoint_with_relay(secret, true, relay)
.await
.context("binding endpoint")?
} else if args.local {
bind_endpoint_local(secret, true)
.await
.context("binding endpoint")?
} else {
bind_endpoint(secret, true)
.await
.context("binding endpoint")?
};
let my_id = endpoint.id();
let id_str = format_endpoint_id(&my_id);
let connect_hint = if let Some(url) = &args.relay_url {
format!("koh connect {id_str} --relay-url {url}")
} else if args.local {
let port = endpoint
.bound_sockets()
.iter()
.find(|s| s.is_ipv4())
.map_or(0, std::net::SocketAddr::port);
format!("koh connect {id_str} --direct <this-host-ip>:{port}")
} else {
format!("koh connect {id_str}")
};
eprintln!("┌─ koh server ready ──────────────────────────────────────");
eprintln!("│ endpoint id : {id_str}");
eprintln!("│ key file : {}", key_file.display());
eprintln!("│ alpn : {}", String::from_utf8_lossy(ALPN));
if args.allow_any {
eprintln!("│ auth : ⚠ ALLOW-ANY (INSECURE) — a shell to ANYONE who can reach this");
eprintln!("│ : endpoint, with no allowlist. Use only on trusted/local nets.");
} else {
eprintln!("│ auth : allowlist ({} client(s))", allow.len());
}
if args.read_only {
eprintln!("│ mode : READ-ONLY (clients can watch, not type)");
}
let restricted = allow.values().filter(|p| p.read_only).count();
let forced = allow.values().filter(|p| p.force_command.is_some()).count();
if (restricted > 0 || forced > 0) && !args.read_only {
eprintln!(
"│ policy : {restricted} restrict, {forced} forced-command (per allow-file)"
);
}
if args.network_timeout_secs > 0 {
eprintln!(
"│ idle-timeout: {}s (server exits — reaping any detached session — after this idle)",
args.network_timeout_secs
);
}
eprintln!("│ connect : {connect_hint}");
eprintln!("└───────────────────────────────────────────────────────────");
if args.allow_any {
warn!(
local = args.local,
"--allow-any is set: ANY peer that can reach this endpoint gets a shell (no allowlist). \
Use only on trusted/local networks; prefer --read-only on a public bind."
);
}
if let Some(qr) = connect_qr(&id_str) {
eprintln!(
"\nScan for the endpoint id (point a phone camera at it). Assumes a dark-background \
terminal;\non a light background it renders inverted — copy the id above instead:\n"
);
eprintln!("{qr}");
} else {
warn!("could not render the connect QR (endpoint id too large to encode)");
}
info!(
transport = "QUIC + TLS 1.3 (iroh)",
kex = "X25519",
post_quantum = false,
"transport crypto posture"
);
let shell = args.shell.clone();
let scrollback = args.scrollback;
let global_read_only = args.read_only;
let allow = std::sync::Arc::new(allow);
let allow_any = args.allow_any;
let store = session::SessionStore::default();
let session_ttl = Duration::from_secs(args.session_ttl_secs);
let reaper_shutdown = tokio_util::sync::CancellationToken::new();
let reaper = tokio::spawn(session::run_reaper(
store.clone(),
session_ttl,
session::REAP_INTERVAL,
reaper_shutdown.clone(),
));
let shutdown = CancellationToken::new();
spawn_signal_drain(shutdown.clone())?;
let conn_limit = Arc::new(tokio::sync::Semaphore::new(args.max_connections as usize));
let pending_cap = (args.max_connections as usize).div_ceil(4).max(4);
let handshake_limit = Arc::new(tokio::sync::Semaphore::new(pending_cap));
let max_sessions = args.max_sessions as usize;
let active = Arc::new(AtomicUsize::new(0));
if args.network_timeout_secs > 0 {
spawn_idle_watchdog(
active.clone(),
Duration::from_secs(args.network_timeout_secs),
shutdown.clone(),
);
}
loop {
let incoming = tokio::select! {
biased;
() = shutdown.cancelled() => break,
inc = endpoint.accept() => match inc {
Some(i) => i,
None => break, },
};
let Ok(permit) = conn_limit.clone().try_acquire_owned() else {
warn!("refusing connection: at max-connections capacity");
incoming.refuse();
continue;
};
let Ok(pending_permit) = handshake_limit.clone().try_acquire_owned() else {
warn!("refusing connection: too many handshakes in flight");
incoming.refuse();
continue;
};
let allow = allow.clone();
let shell = shell.clone();
let store = store.clone();
let active_guard = ConnGuard::new(active.clone());
tokio::spawn(async move {
let _permit = permit;
let _active_guard = active_guard;
let pending_permit = pending_permit;
let conn = match tokio::time::timeout(ACCEPT_HANDSHAKE_TIMEOUT, incoming).await {
Ok(Ok(c)) => c,
Ok(Err(e)) => {
warn!(error = %e, "incoming handshake failed");
return;
}
Err(_) => {
warn!("incoming handshake timed out (stalled QUIC handshake)");
return;
}
};
let peer = conn.remote_id();
if !allow_any && !allow.contains_key(&peer) {
auth_event(Outcome::Rejected, Some(&peer), "not on allowlist");
conn.close(1u32.into(), b"not authorized");
return;
}
match tokio::time::timeout(
Duration::from_secs(3),
crate::transport_iroh::admission::admit(&conn),
)
.await
{
Ok(Ok(())) => {}
Ok(Err(e)) => {
warn!(error = %e, "admission ack failed");
return;
}
Err(_) => {
warn!("admission ack timed out");
return;
}
}
drop(pending_permit);
auth_event(
Outcome::Accepted,
Some(&peer),
"authorized; attaching session",
);
let policy = allow.get(&peer).cloned().unwrap_or(Policy {
read_only: global_read_only,
force_command: None,
});
if policy.read_only || policy.force_command.is_some() {
info!(
peer = %format_endpoint_id(&peer),
read_only = policy.read_only,
forced_command = policy.force_command.is_some(),
"applying per-peer policy"
);
}
let (handle, attach_kind) = match session::attach(
&store,
peer,
shell.as_deref(),
scrollback,
max_sessions,
policy.force_command.as_deref(),
policy.read_only,
)
.await
{
Ok(Some(pair)) => pair,
Ok(None) => {
warn!(peer = %format_endpoint_id(&peer), "refusing session: at max-sessions capacity");
conn.close(1u32.into(), b"server at session capacity");
return;
}
Err(e) => {
error!(error = %e, "failed to start session");
conn.close(1u32.into(), b"session error");
return;
}
};
match attach_kind {
session::AttachKind::Created => {
info!(peer = %format_endpoint_id(&peer), "started a new session");
}
session::AttachKind::Reattached { detached_for } => {
info!(
peer = %format_endpoint_id(&peer),
detached_secs = detached_for.map(|d| d.as_secs()),
"reattaching to this peer's existing session"
);
}
}
let attach_guard = session::AttachGuard::new(store.clone(), peer);
let outcome = run_attached(conn, handle).await;
attach_guard.disarm();
match outcome {
Ok(SessionExit::Detached) => {
session::detach(&store, peer).await;
info!(peer = %format_endpoint_id(&peer), "client detached (session retained)");
}
Ok(SessionExit::ShellExited) => {
session::reap(&store, peer).await;
info!(peer = %format_endpoint_id(&peer), "shell exited; session reaped");
}
Err(e) => {
error!(error = %e, "session loop error");
session::detach(&store, peer).await;
}
}
});
}
info!("draining: stopping reaper and closing endpoint");
shutdown.cancel();
reaper_shutdown.cancel();
let _ = reaper.await;
endpoint.close().await;
Ok(())
}
struct ConnGuard(Arc<AtomicUsize>);
impl ConnGuard {
fn new(active: Arc<AtomicUsize>) -> Self {
active.fetch_add(1, Ordering::SeqCst);
Self(active)
}
}
impl Drop for ConnGuard {
fn drop(&mut self) {
self.0.fetch_sub(1, Ordering::SeqCst);
}
}
fn spawn_signal_drain(shutdown: CancellationToken) -> anyhow::Result<()> {
let mut term = signal(SignalKind::terminate()).context("installing SIGTERM handler")?;
let mut intr = signal(SignalKind::interrupt()).context("installing SIGINT handler")?;
tokio::spawn(async move {
tokio::select! {
_ = term.recv() => {}
_ = intr.recv() => {}
}
info!("received shutdown signal; draining");
shutdown.cancel();
});
Ok(())
}
fn spawn_idle_watchdog(active: Arc<AtomicUsize>, timeout: Duration, shutdown: CancellationToken) {
tokio::spawn(async move {
let tick = Duration::from_secs(1).min(timeout);
let mut idle_since: Option<Instant> = None;
loop {
tokio::select! {
() = shutdown.cancelled() => return,
() = tokio::time::sleep(tick) => {}
}
if active.load(Ordering::SeqCst) == 0 {
let since = *idle_since.get_or_insert_with(Instant::now);
if since.elapsed() >= timeout {
info!(
timeout_secs = timeout.as_secs(),
"network idle timeout; shutting down"
);
shutdown.cancel();
return;
}
} else {
idle_since = None;
}
}
});
}
#[cfg(test)]
mod tests {
use super::connect_qr;
#[test]
fn connect_qr_renders_an_id_and_handles_overlong_input() {
let id = "3f9c".repeat(16);
let qr = connect_qr(&id).expect("an endpoint id must fit in a QR");
assert!(qr.lines().count() > 5, "a QR should be a multi-row block");
assert!(
qr.contains('█') || qr.contains('▀') || qr.contains('▄'),
"the unicode renderer uses half-block glyphs"
);
assert!(
connect_qr(&"a".repeat(10_000)).is_none(),
"overlong input must return None, not panic"
);
}
}