/// blit-proxy library — all proxy logic, usable in-process or as a binary.
///
/// Call [`proxy_socket_path`] to find the socket, then [`run`] to start the
/// proxy on the current thread (blocking, runs its own tokio runtime).
use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicI64, AtomicUsize, Ordering};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::sync::{Mutex, RwLock};
static VERBOSE: AtomicBool = AtomicBool::new(false);
macro_rules! log {
($($arg:tt)*) => {
if VERBOSE.load(Ordering::Relaxed) {
eprintln!($($arg)*);
}
};
}
type BoxRead = Box<dyn AsyncRead + Unpin + Send>;
type BoxWrite = Box<dyn AsyncWrite + Unpin + Send>;
// ---------------------------------------------------------------------------
// Proxy socket path (single stable path for the whole process)
// ---------------------------------------------------------------------------
pub fn proxy_socket_path() -> String {
if let Ok(p) = std::env::var("BLIT_PROXY_SOCK") {
return p;
}
#[cfg(unix)]
{
let dir = std::env::var("XDG_RUNTIME_DIR").unwrap_or_else(|_| "/tmp".into());
format!("{dir}/blit-proxy.sock")
}
#[cfg(windows)]
{
r"\\.\pipe\blit-proxy".into()
}
}
// ---------------------------------------------------------------------------
// Upstream connection
// ---------------------------------------------------------------------------
struct UpstreamConn {
reader: BoxRead,
writer: BoxWrite,
}
// ---------------------------------------------------------------------------
// Per-target pool
// ---------------------------------------------------------------------------
struct Pool {
idle: Mutex<VecDeque<UpstreamConn>>,
/// Number of currently active (proxied) downstream clients.
active: AtomicUsize,
/// Unix seconds of last client connect or disconnect.
/// `i64::MAX` while any client is active.
last_activity: AtomicI64,
pool_size: usize,
upstream_uri: String,
}
impl Pool {
fn new(upstream_uri: String, pool_size: usize) -> Arc<Self> {
Arc::new(Self {
idle: Mutex::new(VecDeque::new()),
active: AtomicUsize::new(0),
last_activity: AtomicI64::new(now_secs()),
pool_size,
upstream_uri,
})
}
async fn acquire(&self) -> Result<UpstreamConn, String> {
{
let mut idle = self.idle.lock().await;
if let Some(conn) = idle.pop_front() {
return Ok(conn);
}
}
self.connect_one().await
}
async fn connect_one(&self) -> Result<UpstreamConn, String> {
connect_upstream(&self.upstream_uri).await
}
fn client_connected(&self) {
self.active.fetch_add(1, Ordering::Relaxed);
self.last_activity.store(i64::MAX, Ordering::Relaxed);
}
fn client_disconnected(&self) {
let prev = self.active.fetch_sub(1, Ordering::Relaxed);
if prev == 1 {
self.last_activity.store(now_secs(), Ordering::Relaxed);
}
}
/// Background task: keep idle slots full.
async fn refill_loop(self: Arc<Self>) {
loop {
let need = {
let idle = self.idle.lock().await;
self.pool_size.saturating_sub(idle.len())
};
for _ in 0..need {
match self.connect_one().await {
Ok(conn) => {
self.idle.lock().await.push_back(conn);
}
Err(e) => {
log!(
"blit-proxy: [{uri}] upstream connect failed: {e}",
uri = self.upstream_uri
);
tokio::time::sleep(std::time::Duration::from_secs(2)).await;
break;
}
}
}
tokio::time::sleep(std::time::Duration::from_millis(200)).await;
}
}
}
fn now_secs() -> i64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs() as i64
}
// ---------------------------------------------------------------------------
// Pool registry: one pool per distinct upstream URI
// ---------------------------------------------------------------------------
struct Registry {
pools: RwLock<HashMap<String, Arc<Pool>>>,
pool_size: usize,
}
impl Registry {
fn new(pool_size: usize) -> Arc<Self> {
Arc::new(Self {
pools: RwLock::new(HashMap::new()),
pool_size,
})
}
/// Get or create the pool for `uri`, spawning its refill task if new.
async fn get_or_create(self: &Arc<Self>, uri: &str) -> Arc<Pool> {
// Fast path: pool already exists.
{
let pools = self.pools.read().await;
if let Some(p) = pools.get(uri) {
return p.clone();
}
}
// Slow path: create and seed.
let mut pools = self.pools.write().await;
// Re-check under write lock.
if let Some(p) = pools.get(uri) {
return p.clone();
}
let pool = Pool::new(uri.to_string(), self.pool_size);
pools.insert(uri.to_string(), pool.clone());
drop(pools);
// Seed eagerly in a background task.
let pool_seed = pool.clone();
tokio::spawn(async move {
for _ in 0..pool_seed.pool_size {
match pool_seed.connect_one().await {
Ok(conn) => pool_seed.idle.lock().await.push_back(conn),
Err(e) => {
log!(
"blit-proxy: [{uri}] initial connect: {e}",
uri = pool_seed.upstream_uri
);
break;
}
}
}
pool_seed.refill_loop().await;
});
pool
}
/// Returns the most-recent `last_activity` across all pools.
/// Returns `i64::MAX` if any pool has an active client.
async fn latest_activity(&self) -> i64 {
let pools = self.pools.read().await;
if pools.is_empty() {
return now_secs();
}
pools
.values()
.map(|p| p.last_activity.load(Ordering::Relaxed))
.max()
.unwrap_or_else(now_secs)
}
}
// ---------------------------------------------------------------------------
// Upstream transport implementations
// ---------------------------------------------------------------------------
async fn connect_upstream(uri: &str) -> Result<UpstreamConn, String> {
if let Some(rest) = uri.strip_prefix("share:") {
return connect_share(rest).await;
}
#[cfg(unix)]
if let Some(rest) = uri.strip_prefix("ssh:") {
return connect_ssh(rest).await;
}
// Extract query parameters from URIs that support them.
let (base_uri, passphrase, cert_hash) = extract_uri_params(uri);
if let Some(path) = base_uri.strip_prefix("socket:") {
return connect_socket(path).await;
}
if let Some(addr) = base_uri.strip_prefix("tcp:") {
return connect_tcp(addr).await;
}
if base_uri.starts_with("ws://") || base_uri.starts_with("wss://") {
return connect_ws(&base_uri, passphrase.as_deref()).await;
}
if let Some(rest) = base_uri.strip_prefix("wt://") {
let cert_bytes = cert_hash.as_deref().and_then(parse_hex);
return connect_wt(rest, passphrase.as_deref(), &cert_bytes).await;
}
Err(format!(
"unknown upstream URI scheme in '{uri}' \
(expected socket:, tcp:, ws://, wss://, wt://, share:, or ssh:)"
))
}
/// Split `share:` URI rest into (passphrase, hub_url).
///
/// Accepted forms:
/// `myphrase` — use default hub
/// `myphrase?hub=wss://custom.hub` — use specific hub
fn parse_share_uri(rest: &str) -> (String, String) {
let (passphrase_raw, hub) = if let Some(q_pos) = rest.find('?') {
let phrase = &rest[..q_pos];
let query = &rest[q_pos + 1..];
let hub = query
.split('&')
.find_map(|kv| kv.strip_prefix("hub=").map(percent_decode))
.unwrap_or_else(|| blit_webrtc_forwarder::DEFAULT_HUB_URL.to_string());
(phrase.to_string(), hub)
} else {
(
rest.to_string(),
blit_webrtc_forwarder::DEFAULT_HUB_URL.to_string(),
)
};
let passphrase = percent_decode(&passphrase_raw);
let hub = blit_webrtc_forwarder::normalize_hub(&hub);
(passphrase, hub)
}
async fn connect_share(rest: &str) -> Result<UpstreamConn, String> {
let (passphrase, hub) = parse_share_uri(rest);
let stream = blit_webrtc_forwarder::client::connect(&passphrase, &hub)
.await
.map_err(|e| format!("share:{rest}: {e}"))?;
let (r, w) = tokio::io::split(stream);
Ok(UpstreamConn {
reader: Box::new(r),
writer: Box::new(w),
})
}
// ---------------------------------------------------------------------------
// SSH stdio bridge
// ---------------------------------------------------------------------------
/// Shell snippet that locates the blit-server Unix socket on the remote host.
#[cfg(unix)]
const SSH_SOCK_SEARCH: &str = r#"if [ -n "$BLIT_SOCK" ]; then S="$BLIT_SOCK"; elif [ -n "$TMPDIR" ] && [ -S "$TMPDIR/blit.sock" ]; then S="$TMPDIR/blit.sock"; elif [ -S "/tmp/blit-$(id -un).sock" ]; then S="/tmp/blit-$(id -un).sock"; elif [ -S "/run/blit/$(id -un).sock" ]; then S="/run/blit/$(id -un).sock"; elif [ -n "$XDG_RUNTIME_DIR" ] && [ -S "$XDG_RUNTIME_DIR/blit.sock" ]; then S="$XDG_RUNTIME_DIR/blit.sock"; else S=/tmp/blit.sock; fi"#;
/// Shell snippet that starts blit-server on the remote if the socket is absent/stale.
#[cfg(unix)]
const SSH_AUTOSTART: &str = "if [ -S \"$S\" ]; then if command -v nc >/dev/null 2>&1; then nc -z -U \"$S\" 2>/dev/null || rm -f \"$S\"; elif command -v socat >/dev/null 2>&1; then socat /dev/null \"UNIX-CONNECT:$S\" 2>/dev/null || rm -f \"$S\"; fi; fi; if ! [ -S \"$S\" ]; then if command -v blit >/dev/null 2>&1; then blit server & i=0; while ! [ -S \"$S\" ] && [ $i -lt 50 ]; do sleep 0.1; i=$((i+1)); done; elif command -v blit-server >/dev/null 2>&1; then blit-server & i=0; while ! [ -S \"$S\" ] && [ $i -lt 50 ]; do sleep 0.1; i=$((i+1)); done; fi; fi";
/// Escape a string for use inside double quotes in a POSIX shell.
/// Handles `\`, `$`, `` ` ``, and `"`. Used instead of single-quote
/// escaping because the result is embedded inside an outer `sh -c '…'`
/// wrapper where nested single quotes would break.
#[cfg(unix)]
fn dq_escape(s: &str) -> String {
let mut out = String::with_capacity(s.len());
for ch in s.chars() {
match ch {
'\\' | '$' | '`' | '"' => {
out.push('\\');
out.push(ch);
}
_ => out.push(ch),
}
}
out
}
/// Parse `[user@]host[:socket_path]` from an `ssh:` URI remainder.
/// Returns `(ssh_target, remote_socket)`.
#[cfg(unix)]
fn parse_ssh_target(rest: &str) -> (&str, Option<&str>) {
let colon_search_start = rest.find('@').map(|a| a + 1).unwrap_or(0);
if let Some(rel) = rest[colon_search_start..].find(':') {
let pos = colon_search_start + rel;
let path = &rest[pos + 1..];
if path.is_empty() {
(rest, None)
} else {
(&rest[..pos], Some(path))
}
} else {
(rest, None)
}
}
/// Connect to a remote blit-server via an SSH stdio bridge.
///
/// Spawns `ssh <host> -- sh -c '…find-socket…; nc -U "$S"'` and wraps
/// the child's stdin/stdout as an `UpstreamConn`.
#[cfg(unix)]
async fn connect_ssh(rest: &str) -> Result<UpstreamConn, String> {
if rest.is_empty() {
return Err("ssh: destination requires a host".into());
}
let (host, remote_socket) = parse_ssh_target(rest);
if host.starts_with('-') {
return Err(format!(
"invalid ssh host '{host}': must not start with '-'"
));
}
let resolve = match remote_socket {
Some(path) => format!("S=\"{}\"", dq_escape(path)),
None => SSH_SOCK_SEARCH.to_string(),
};
let bridge = format!(
"sh -c '{resolve}; {SSH_AUTOSTART}; if command -v nc >/dev/null 2>&1; then exec nc -U \"$S\"; else exec socat - \"UNIX-CONNECT:$S\"; fi'"
);
let mut child = tokio::process::Command::new("ssh")
.arg("-T")
.arg("-o")
.arg("ControlMaster=auto")
.arg("-o")
.arg("ControlPath=/tmp/blit-ssh-%r@%h:%p")
.arg("-o")
.arg("ControlPersist=300")
.arg(host)
.arg("--")
.arg(&bridge)
.stdin(std::process::Stdio::piped())
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::inherit())
.spawn()
.map_err(|e| format!("ssh: {e}"))?;
let stdin = child
.stdin
.take()
.ok_or_else(|| "ssh: could not get stdin".to_string())?;
let stdout = child
.stdout
.take()
.ok_or_else(|| "ssh: could not get stdout".to_string())?;
// Keep the child alive by moving it into a background task that waits on it.
tokio::spawn(async move {
let _ = child.wait().await;
});
Ok(UpstreamConn {
reader: Box::new(stdout),
writer: Box::new(stdin),
})
}
/// Split a URI into (base, passphrase, certHash) by parsing query params.
/// Only applies to ws://, wss://, wt:// — socket: and tcp: are returned as-is.
fn extract_uri_params(uri: &str) -> (String, Option<String>, Option<String>) {
if !uri.starts_with("ws://") && !uri.starts_with("wss://") && !uri.starts_with("wt://") {
return (uri.to_string(), None, None);
}
// Find '?'.
let (base, query) = match uri.find('?') {
Some(pos) => (&uri[..pos], Some(&uri[pos + 1..])),
None => (uri, None),
};
let mut passphrase = None;
let mut cert_hash = None;
if let Some(q) = query {
for param in q.split('&') {
if let Some(v) = param.strip_prefix("passphrase=") {
passphrase = Some(percent_decode(v));
} else if let Some(v) = param.strip_prefix("certHash=") {
cert_hash = Some(v.to_string());
}
}
}
(base.to_string(), passphrase, cert_hash)
}
fn percent_decode(s: &str) -> String {
// Minimal %XX decoder sufficient for passphrases.
let mut out = String::with_capacity(s.len());
let mut chars = s.chars().peekable();
while let Some(c) = chars.next() {
if c == '%' {
let h1 = chars.next().unwrap_or('0');
let h2 = chars.next().unwrap_or('0');
if let Ok(b) = u8::from_str_radix(&format!("{h1}{h2}"), 16) {
out.push(b as char);
continue;
}
}
out.push(c);
}
out
}
async fn connect_socket(path: &str) -> Result<UpstreamConn, String> {
#[cfg(unix)]
{
let stream = tokio::net::UnixStream::connect(path)
.await
.map_err(|e| format!("socket:{path}: {e}"))?;
let (r, w) = tokio::io::split(stream);
Ok(UpstreamConn {
reader: Box::new(r),
writer: Box::new(w),
})
}
#[cfg(not(unix))]
{
use tokio::net::windows::named_pipe::ClientOptions;
let pipe = ClientOptions::new()
.open(path)
.map_err(|e| format!("socket:{path}: {e}"))?;
let (r, w) = tokio::io::split(pipe);
Ok(UpstreamConn {
reader: Box::new(r),
writer: Box::new(w),
})
}
}
async fn connect_tcp(addr: &str) -> Result<UpstreamConn, String> {
let stream = tokio::net::TcpStream::connect(addr)
.await
.map_err(|e| format!("tcp:{addr}: {e}"))?;
let _ = stream.set_nodelay(true);
let (r, w) = tokio::io::split(stream);
Ok(UpstreamConn {
reader: Box::new(r),
writer: Box::new(w),
})
}
async fn connect_ws(uri: &str, passphrase: Option<&str>) -> Result<UpstreamConn, String> {
use futures_util::{SinkExt, StreamExt};
use tokio_tungstenite::tungstenite::Message;
let (mut ws, _) = tokio_tungstenite::connect_async(uri)
.await
.map_err(|e| format!("{uri}: {e}"))?;
let pass = passphrase.unwrap_or("");
ws.send(Message::Text(pass.into()))
.await
.map_err(|e| format!("{uri}: auth send: {e}"))?;
match ws.next().await {
Some(Ok(Message::Text(t))) if t.trim() == "ok" => {}
Some(Ok(Message::Text(t))) => {
return Err(format!("{uri}: auth rejected: {}", t.trim()));
}
other => {
return Err(format!("{uri}: unexpected auth response: {other:?}"));
}
}
let (ws_write, ws_read) = ws.split();
Ok(UpstreamConn {
reader: Box::new(WsFrameReader {
inner: ws_read,
buf: bytes::Bytes::new(),
}),
writer: Box::new(WsFrameWriter { inner: ws_write }),
})
}
async fn connect_wt(
rest: &str,
passphrase: Option<&str>,
cert_hash: &Option<Vec<u8>>,
) -> Result<UpstreamConn, String> {
use web_transport_quinn as wt;
// Build the URL for the WT session (must use https: scheme).
let (host, port) = parse_wt_host_port(rest)?;
let url: url::Url = format!("https://{host}:{port}/")
.parse()
.map_err(|e| format!("wt: url: {e}"))?;
// Build the client with appropriate certificate verification.
let client: wt::Client = if let Some(hash) = cert_hash {
wt::ClientBuilder::new()
.with_server_certificate_hashes(vec![hash.clone()])
.map_err(|e| format!("wt: client build: {e}"))?
} else {
wt::ClientBuilder::new()
.with_system_roots()
.map_err(|e| format!("wt: client build: {e}"))?
};
let session = client
.connect(url)
.await
.map_err(|e| format!("wt: connect {host}:{port}: {e}"))?;
let (mut send, mut recv) = session
.open_bi()
.await
.map_err(|e| format!("wt: open_bi: {e}"))?;
// Auth: 2-byte-LE passphrase length + passphrase bytes, then read 1-byte response.
let pass = passphrase.unwrap_or("").as_bytes();
let mut auth_buf = Vec::with_capacity(2 + pass.len());
auth_buf.extend_from_slice(&(pass.len() as u16).to_le_bytes());
auth_buf.extend_from_slice(pass);
send.write_all(&auth_buf)
.await
.map_err(|e| format!("wt: auth send: {e}"))?;
let mut resp = [0u8; 1];
recv.read_exact(&mut resp)
.await
.map_err(|e| format!("wt: auth recv: {e}"))?;
if resp[0] != 1 {
return Err(format!(
"wt: auth rejected (response byte {:#04x})",
resp[0]
));
}
Ok(UpstreamConn {
reader: Box::new(recv),
writer: Box::new(send),
})
}
fn parse_wt_host_port(rest: &str) -> Result<(String, u16), String> {
let without_path = rest.split('/').next().unwrap_or(rest);
if let Some(colon) = without_path.rfind(':') {
let host = without_path[..colon].to_string();
let port: u16 = without_path[colon + 1..]
.parse()
.map_err(|_| format!("wt: invalid port in '{rest}'"))?;
Ok((host, port))
} else {
Ok((without_path.to_string(), 443))
}
}
fn parse_hex(s: &str) -> Option<Vec<u8>> {
if !s.len().is_multiple_of(2) {
return None;
}
(0..s.len())
.step_by(2)
.map(|i| u8::from_str_radix(&s[i..i + 2], 16).ok())
.collect()
}
// ---------------------------------------------------------------------------
// WS ↔ raw blit-frame adapters
// ---------------------------------------------------------------------------
use futures_util::{SinkExt, StreamExt};
use tokio_tungstenite::tungstenite::Message;
type WsSink = futures_util::stream::SplitSink<
tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>,
Message,
>;
type WsStream = futures_util::stream::SplitStream<
tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>,
>;
struct WsFrameReader {
inner: WsStream,
buf: bytes::Bytes,
}
impl AsyncRead for WsFrameReader {
fn poll_read(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
out: &mut tokio::io::ReadBuf<'_>,
) -> std::task::Poll<std::io::Result<()>> {
loop {
if !self.buf.is_empty() {
let n = out.remaining().min(self.buf.len());
out.put_slice(&self.buf[..n]);
self.buf = self.buf.slice(n..);
return std::task::Poll::Ready(Ok(()));
}
match self.inner.poll_next_unpin(cx) {
std::task::Poll::Pending => return std::task::Poll::Pending,
std::task::Poll::Ready(None) => return std::task::Poll::Ready(Ok(())),
std::task::Poll::Ready(Some(Err(e))) => {
return std::task::Poll::Ready(Err(std::io::Error::other(e)));
}
std::task::Poll::Ready(Some(Ok(msg))) => {
let data = match msg {
Message::Binary(d) => d,
Message::Close(_) => return std::task::Poll::Ready(Ok(())),
_ => continue,
};
let len = data.len() as u32;
let mut framed = Vec::with_capacity(4 + data.len());
framed.extend_from_slice(&len.to_le_bytes());
framed.extend_from_slice(&data);
self.buf = bytes::Bytes::from(framed);
}
}
}
}
}
struct WsFrameWriter {
inner: WsSink,
}
impl AsyncWrite for WsFrameWriter {
fn poll_write(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
buf: &[u8],
) -> std::task::Poll<std::io::Result<usize>> {
if buf.len() < 4 {
return std::task::Poll::Ready(Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"ws frame writer: buffer too small for length prefix",
)));
}
let len = u32::from_le_bytes(buf[..4].try_into().unwrap()) as usize;
if buf.len() < 4 + len {
return std::task::Poll::Ready(Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!(
"ws frame writer: buffer ({}) too small for frame payload ({})",
buf.len(),
4 + len
),
)));
}
let payload = buf[4..4 + len].to_vec();
match self.inner.poll_ready_unpin(cx) {
std::task::Poll::Pending => return std::task::Poll::Pending,
std::task::Poll::Ready(Err(e)) => {
return std::task::Poll::Ready(Err(std::io::Error::other(e)));
}
std::task::Poll::Ready(Ok(())) => {}
}
match self.inner.start_send_unpin(Message::Binary(payload.into())) {
Err(e) => std::task::Poll::Ready(Err(std::io::Error::other(e))),
Ok(()) => std::task::Poll::Ready(Ok(4 + len)),
}
}
fn poll_flush(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<std::io::Result<()>> {
self.inner
.poll_flush_unpin(cx)
.map_err(std::io::Error::other)
}
fn poll_shutdown(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<std::io::Result<()>> {
self.inner
.poll_close_unpin(cx)
.map_err(std::io::Error::other)
}
}
// ---------------------------------------------------------------------------
// Downstream listener
// ---------------------------------------------------------------------------
/// Read one line from the downstream socket (up to 4 KiB).
#[cfg(unix)]
async fn read_line(stream: &mut tokio::net::UnixStream) -> Option<String> {
let mut buf = Vec::with_capacity(128);
let mut byte = [0u8; 1];
loop {
match stream.read(&mut byte).await {
Ok(0) | Err(_) => return None,
Ok(_) => {
if byte[0] == b'\n' {
break;
}
buf.push(byte[0]);
if buf.len() > 4096 {
return None;
}
}
}
}
Some(
String::from_utf8_lossy(&buf)
.trim_end_matches('\r')
.to_string(),
)
}
#[cfg(unix)]
async fn run_listener(registry: Arc<Registry>, sock_path: &str) {
let _ = std::fs::remove_file(sock_path);
let listener = tokio::net::UnixListener::bind(sock_path).unwrap_or_else(|e| {
eprintln!("blit-proxy: cannot bind to {sock_path}: {e}");
std::process::exit(1);
});
log!("blit-proxy: listening on {sock_path}");
loop {
match listener.accept().await {
Ok((stream, _)) => {
let registry = registry.clone();
tokio::spawn(async move {
handle_downstream(registry, stream).await;
});
}
Err(e) => log!("blit-proxy: accept error: {e}"),
}
}
}
#[cfg(unix)]
async fn handle_downstream(registry: Arc<Registry>, mut downstream: tokio::net::UnixStream) {
// Handshake: read "target <uri>\n" or "shutdown\n"
let line = match read_line(&mut downstream).await {
Some(l) => l,
None => return,
};
if line == "shutdown" {
let _ = downstream.write_all(b"ok\n").await;
log!("blit-proxy: shutdown requested, exiting");
std::process::exit(0);
}
let uri = match line.strip_prefix("target ") {
Some(u) if !u.is_empty() => u.to_string(),
_ => {
let _ = downstream.write_all(b"error invalid handshake\n").await;
return;
}
};
let pool = registry.get_or_create(&uri).await;
let upstream = match pool.acquire().await {
Ok(u) => u,
Err(e) => {
log!("blit-proxy: [{uri}] upstream unavailable: {e}");
let msg = format!("error {e}\n");
let _ = downstream.write_all(msg.as_bytes()).await;
return;
}
};
if downstream.write_all(b"ok\n").await.is_err() {
return;
}
pool.client_connected();
let (mut ds_read, mut ds_write) = tokio::io::split(downstream);
let (mut us_read, mut us_write) = (upstream.reader, upstream.writer);
let mut d2u = tokio::spawn(async move { tokio::io::copy(&mut ds_read, &mut us_write).await });
let mut u2d = tokio::spawn(async move { tokio::io::copy(&mut us_read, &mut ds_write).await });
tokio::select! { _ = &mut d2u => { u2d.abort(); }, _ = &mut u2d => { d2u.abort(); } }
pool.client_disconnected();
}
// ---------------------------------------------------------------------------
// Public entry point
// ---------------------------------------------------------------------------
/// Run the proxy on the current thread (blocks until the process exits).
///
/// Reads `BLIT_PROXY_SOCK`, `BLIT_PROXY_POOL`, and `BLIT_PROXY_IDLE` from the
/// environment. When called from within the `blit` binary instead of the
/// standalone `blit-proxy` binary, `verbose` should be `false`.
pub fn run(verbose: bool) {
if verbose {
VERBOSE.store(true, Ordering::Relaxed);
}
let sock_path = proxy_socket_path();
let pool_size: usize = std::env::var("BLIT_PROXY_POOL")
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(4);
let idle_secs: Option<u64> = std::env::var("BLIT_PROXY_IDLE")
.ok()
.and_then(|v| v.parse().ok());
tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()
.expect("blit-proxy: tokio runtime")
.block_on(async move {
rustls::crypto::ring::default_provider()
.install_default()
.ok(); // may already be installed by the CLI's runtime
let registry = Registry::new(pool_size);
// Idle-timeout watcher.
if let Some(idle) = idle_secs {
let reg = registry.clone();
let sock = sock_path.clone();
tokio::spawn(async move {
loop {
tokio::time::sleep(std::time::Duration::from_secs(5)).await;
let last = reg.latest_activity().await;
if last == i64::MAX {
continue;
}
let elapsed = now_secs().saturating_sub(last) as u64;
if elapsed >= idle {
log!("blit-proxy: idle for {elapsed}s (limit {idle}s), exiting");
let _ = std::fs::remove_file(&sock);
std::process::exit(0);
}
}
});
}
#[cfg(unix)]
{
let sock_cleanup = sock_path.clone();
tokio::spawn(async move {
let _ = tokio::signal::ctrl_c().await;
let _ = std::fs::remove_file(&sock_cleanup);
std::process::exit(0);
});
run_listener(registry, &sock_path).await;
}
#[cfg(not(unix))]
{
eprintln!("blit-proxy: Windows not yet supported");
std::process::exit(1);
}
});
}
// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn proxy_socket_path_default() {
// Just verify it produces a non-empty string without panicking.
let p = proxy_socket_path();
assert!(!p.is_empty());
assert!(p.contains("blit-proxy"));
}
#[test]
fn parse_hex_valid() {
assert_eq!(parse_hex("deadbeef"), Some(vec![0xde, 0xad, 0xbe, 0xef]));
assert_eq!(parse_hex(""), Some(vec![]));
}
#[test]
fn parse_hex_odd_length() {
assert_eq!(parse_hex("abc"), None);
}
#[test]
fn parse_hex_invalid_char() {
assert_eq!(parse_hex("zz"), None);
}
#[test]
fn extract_uri_params_no_query() {
let (base, pass, cert) = extract_uri_params("wss://host:3264/");
assert_eq!(base, "wss://host:3264/");
assert_eq!(pass, None);
assert_eq!(cert, None);
}
#[test]
fn extract_uri_params_passphrase() {
let (base, pass, cert) = extract_uri_params("wss://host:3264/?passphrase=secret");
assert_eq!(base, "wss://host:3264/");
assert_eq!(pass, Some("secret".into()));
assert_eq!(cert, None);
}
#[test]
fn extract_uri_params_both() {
let (base, pass, cert) =
extract_uri_params("wt://host:4433/?passphrase=abc&certHash=deadbeef");
assert_eq!(base, "wt://host:4433/");
assert_eq!(pass, Some("abc".into()));
assert_eq!(cert, Some("deadbeef".into()));
}
#[test]
fn extract_uri_params_socket_unchanged() {
let (base, pass, cert) = extract_uri_params("socket:/tmp/blit.sock");
assert_eq!(base, "socket:/tmp/blit.sock");
assert_eq!(pass, None);
assert_eq!(cert, None);
}
#[test]
fn parse_wt_host_port_with_port() {
assert_eq!(parse_wt_host_port("host:4433"), Ok(("host".into(), 4433)));
}
#[test]
fn parse_wt_host_port_default() {
assert_eq!(parse_wt_host_port("host"), Ok(("host".into(), 443)));
}
#[test]
fn percent_decode_basic() {
assert_eq!(percent_decode("hello%20world"), "hello world");
assert_eq!(percent_decode("plain"), "plain");
}
#[test]
fn parse_share_uri_no_hub() {
let (pass, hub) = parse_share_uri("myphrase");
assert_eq!(pass, "myphrase");
assert_eq!(
hub,
blit_webrtc_forwarder::normalize_hub(blit_webrtc_forwarder::DEFAULT_HUB_URL)
);
}
#[test]
fn parse_share_uri_with_hub() {
let (pass, hub) = parse_share_uri("myphrase?hub=wss://custom.example.com");
assert_eq!(pass, "myphrase");
assert_eq!(hub, "wss://custom.example.com");
}
#[test]
fn parse_share_uri_hub_normalized() {
let (pass, hub) = parse_share_uri("myphrase?hub=custom.example.com");
assert_eq!(pass, "myphrase");
assert_eq!(hub, "wss://custom.example.com");
}
#[test]
fn parse_share_uri_percent_encoded_passphrase() {
let (pass, hub) = parse_share_uri("my%3Fphrase");
assert_eq!(pass, "my?phrase");
assert_eq!(
hub,
blit_webrtc_forwarder::normalize_hub(blit_webrtc_forwarder::DEFAULT_HUB_URL)
);
}
}