use std::ffi::OsString;
use std::io::{self, ErrorKind, Read, Write};
use std::net::{Shutdown, TcpListener, TcpStream};
use std::path::PathBuf;
use std::process::{Child, Command, Stdio};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::thread;
use std::time::{Duration, Instant};
use intuituive::App;
use tungstenite::WebSocket;
use tungstenite::protocol::{Role, WebSocketConfig};
use crate::{Renderer, session};
pub const SESSION_ENV: &str = "RS_RICH_WEB_SESSION";
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum Isolation {
#[default]
Processes,
Threads,
}
const CONNECT_TIMEOUT: Duration = Duration::from_secs(15);
const HELLO_TIMEOUT: Duration = Duration::from_secs(2);
const MAX_WAITING: usize = 64;
const EXIT_WAIT: Duration = Duration::from_secs(1);
const TURN: Duration = Duration::from_millis(50);
#[derive(Clone, Debug, Default)]
pub(crate) struct Spawn {
pub command: Option<(PathBuf, Vec<OsString>)>,
pub memory: Option<u64>,
pub cpu: Option<Duration>,
}
static BINDS: AtomicUsize = AtomicUsize::new(0);
pub(crate) fn next_bind() -> usize {
BINDS.fetch_add(1, Ordering::SeqCst)
}
#[derive(Debug, PartialEq, Eq)]
pub(crate) struct Spec {
port: u16,
parent: u32,
secret: String,
reply: String,
bind: usize,
columns: u16,
rows: u16,
renderer: Renderer,
memory: Option<u64>,
cpu: Option<u64>,
compress: Option<u64>,
}
impl Spec {
fn encode(&self) -> String {
let limit = |n: Option<u64>| n.map_or("-".to_string(), |n| n.to_string());
format!(
"3 {} {} {} {} {} {} {} {} {} {} {}",
self.port,
self.parent,
self.secret,
self.reply,
self.bind,
self.columns,
self.rows,
self.renderer.name(),
limit(self.memory),
limit(self.cpu),
limit(self.compress)
)
}
fn decode(text: &str) -> Option<Spec> {
let mut parts = text.split(' ');
if parts.next()? != "3" {
return None;
}
let mut next = || parts.next();
let limit = |text: &str| match text {
"-" => Some(None),
n => n.parse().ok().map(Some),
};
let spec = Spec {
port: next()?.parse().ok()?,
parent: next()?.parse().ok()?,
secret: next()?.to_string(),
reply: next()?.to_string(),
bind: next()?.parse().ok()?,
columns: next()?.parse().ok()?,
rows: next()?.parse().ok()?,
renderer: Renderer::parse(next()?)?,
memory: limit(next()?)?,
cpu: limit(next()?)?,
compress: limit(next()?)?,
};
next().is_none().then_some(spec)
}
}
pub(crate) fn session_for(bind: usize) -> Option<Spec> {
let spec = own_spec()?;
(spec.bind == bind).then_some(spec)
}
pub(crate) fn is_session_process() -> bool {
own_spec().is_some()
}
fn own_spec() -> Option<Spec> {
let spec = Spec::decode(&std::env::var(SESSION_ENV).ok()?)?;
#[cfg(unix)]
{
use rustix::process::{Pid, getpgrp, getpid, getppid};
if Pid::as_raw(getppid()) as u32 != spec.parent {
if getpgrp() == getpid() {
eprintln!("rich-web: the session's server has gone");
std::process::exit(1);
}
return None;
}
}
Some(spec)
}
pub(crate) fn run_child(spec: Spec, app: impl FnOnce() -> App) -> i32 {
#[cfg(windows)]
#[allow(unsafe_code)]
unsafe {
std::env::remove_var(SESSION_ENV)
};
#[cfg(unix)]
watch_server(spec.parent);
limit(spec.memory, spec.cpu);
let stream = match TcpStream::connect(("127.0.0.1", spec.port)) {
Ok(stream) => stream,
Err(error) => {
eprintln!("rich-web: the session could not reach its server: {error}");
return 1;
}
};
let _ = stream.set_nodelay(true);
let mut hello = stream.try_clone().expect("a socket clones");
if hello
.write_all(format!("{}\n", spec.secret).as_bytes())
.is_err()
{
return 1;
}
let _ = stream.set_read_timeout(Some(CONNECT_TIMEOUT));
let mut answer = vec![0u8; spec.reply.len() + 1];
if hello.read_exact(&mut answer).is_err()
|| !crate::http::same_token(
std::str::from_utf8(&answer[..spec.reply.len()]).unwrap_or(""),
&spec.reply,
)
|| answer[spec.reply.len()] != b'\n'
{
eprintln!("rich-web: the session's server did not answer");
return 1;
}
let _ = stream.set_read_timeout(None);
let config = WebSocketConfig::default()
.max_message_size(Some(crate::MAX_MESSAGE))
.max_frame_size(Some(crate::MAX_MESSAGE));
let mut socket = WebSocket::from_raw_socket(stream, Role::Server, Some(config));
let stopping = AtomicBool::new(false);
let code = match session::run(
&mut socket,
app(),
spec.columns,
spec.rows,
&stopping,
spec.renderer,
session::Compress(spec.compress.map(|n| n as usize)),
) {
Ok(()) => 0,
Err(_) => 1,
};
#[cfg(unix)]
thread::sleep(ENDING);
code
}
#[cfg(unix)]
const ENDING: Duration = Duration::from_millis(250);
#[cfg(unix)]
const WATCH: Duration = Duration::from_millis(50);
#[cfg(unix)]
fn watch_server(server: u32) {
use rustix::process::{Pid, Signal, getpid, getppid, kill_process_group};
thread::spawn(move || {
loop {
let parent = Pid::as_raw(getppid()) as u32;
if parent != server {
let _ = kill_process_group(getpid(), Signal::KILL);
std::process::exit(1);
}
thread::sleep(WATCH);
}
});
}
#[cfg(unix)]
fn limits(
memory: Option<u64>,
cpu: Option<u64>,
) -> impl Iterator<Item = (rustix::process::Resource, rustix::process::Rlimit)> {
use rustix::process::{Resource, Rlimit};
let memory = memory.filter(|_| !cfg!(target_vendor = "apple"));
[(Resource::As, memory), (Resource::Cpu, cpu)]
.into_iter()
.filter_map(|(resource, value)| {
let value = value?;
let limit = Rlimit {
current: Some(value),
maximum: Some(value),
};
Some((resource, limit))
})
}
#[cfg(unix)]
fn limit(memory: Option<u64>, cpu: Option<u64>) {
for (resource, limit) in limits(memory, cpu) {
if let Err(error) = rustix::process::setrlimit(resource, limit) {
eprintln!("rich-web: the session's limit could not be set: {error}");
}
}
}
#[cfg(target_os = "linux")]
fn limit_child(pid: u32, memory: Option<u64>, cpu: Option<u64>) -> io::Result<()> {
let pid = i32::try_from(pid)
.ok()
.and_then(rustix::process::Pid::from_raw)
.ok_or_else(|| io::Error::other("no such process"))?;
for (resource, limit) in limits(memory, cpu) {
rustix::process::prlimit(Some(pid), resource, limit)?;
}
Ok(())
}
#[cfg(not(unix))]
fn limit(_memory: Option<u64>, _cpu: Option<u64>) {}
struct Running<'a>(&'a AtomicUsize);
impl Drop for Running<'_> {
fn drop(&mut self) {
self.0.fetch_sub(1, Ordering::SeqCst);
}
}
pub(crate) struct Session<'a> {
pub spawn: &'a Spawn,
pub bind: usize,
pub columns: u16,
pub rows: u16,
pub renderer: Renderer,
pub compress: Option<usize>,
pub stopping: &'a AtomicBool,
pub running: &'a AtomicUsize,
}
pub(crate) fn run_parent(browser: TcpStream, session: Session<'_>) {
session.running.fetch_add(1, Ordering::SeqCst);
let _counted = Running(session.running);
if session.stopping.load(Ordering::SeqCst) {
return goodbye(browser, "The server stopped.");
}
let started = match start(&session) {
Ok(started) => started,
Err(error) => {
return goodbye(
browser,
&format!("The session's process could not start: {error}."),
);
}
};
let (mut child, tree, listener, secrets) = started;
let memory = session.spawn.memory;
let child_stream = match connected(&listener, &secrets, &mut child, memory, session.stopping) {
Ok(stream) => stream,
Err(why) => {
end(&mut child, tree);
return goodbye(browser, &why);
}
};
drop(listener);
let relayed = relay(&browser, &child_stream, session.stopping);
let _ = child_stream.shutdown(Shutdown::Both);
let status = end(&mut child, tree);
let reason = match relayed.ended {
_ if !relayed.can_close => None,
_ if session.stopping.load(Ordering::SeqCst) => Some("The server stopped.".to_string()),
Ended::Page => None,
Ended::Child => Some(match status {
Some(status) if status.success() => "The session ended.".to_string(),
Some(status) => format!("The session's process ended unexpectedly ({status})."),
None => "The session's process ended unexpectedly.".to_string(),
}),
};
match reason {
Some(reason) => goodbye(browser, &reason),
None => {
let _ = browser.shutdown(Shutdown::Both);
}
}
}
fn start(session: &Session<'_>) -> io::Result<(Child, Tree, TcpListener, Secrets)> {
let listener = TcpListener::bind(("127.0.0.1", 0))?;
let secrets = Secrets {
secret: crate::random_token()?,
reply: crate::random_token()?,
};
let spec = Spec {
port: listener.local_addr()?.port(),
parent: std::process::id(),
secret: secrets.secret.clone(),
reply: secrets.reply.clone(),
bind: session.bind,
columns: session.columns,
rows: session.rows,
renderer: session.renderer,
memory: session.spawn.memory,
cpu: session.spawn.cpu.map(|cpu| cpu.as_secs().max(1)),
compress: session.compress.map(|n| n as u64),
};
let (program, args) = match &session.spawn.command {
Some((program, args)) => (program.clone(), args.clone()),
None => (
std::env::current_exe()?,
std::env::args_os().skip(1).collect(),
),
};
let mut command = Command::new(program);
command
.args(args)
.env(SESSION_ENV, spec.encode())
.stdin(Stdio::null());
#[cfg(unix)]
std::os::unix::process::CommandExt::process_group(&mut command, 0);
let mut child = command.spawn()?;
#[cfg(target_os = "linux")]
if let Err(error) = limit_child(child.id(), spec.memory, spec.cpu)
.and_then(|()| within_memory(child.id(), spec.memory))
{
let _ = child.kill();
let _ = child.wait();
return Err(io::Error::other(format!(
"its limits could not be set: {error}"
)));
}
let tree = match Tree::new(&child) {
Ok(tree) => tree,
Err(error) => {
let _ = child.kill();
let _ = child.wait();
return Err(io::Error::other(format!(
"it could not be put in a job object: {error}"
)));
}
};
listener.set_nonblocking(true)?;
Ok((child, tree, listener, secrets))
}
struct Tree {
#[cfg(unix)]
leader: u32,
#[cfg(windows)]
_job: win32job::Job,
}
impl Tree {
#[cfg(unix)]
fn new(child: &Child) -> io::Result<Tree> {
Ok(Tree { leader: child.id() })
}
#[cfg(windows)]
fn new(child: &Child) -> io::Result<Tree> {
use std::os::windows::io::AsRawHandle;
let mut limits = win32job::ExtendedLimitInfo::new();
limits.limit_kill_on_job_close();
let job = win32job::Job::create_with_limit_info(&limits)?;
job.assign_process(child.as_raw_handle() as isize)?;
Ok(Tree { _job: job })
}
#[cfg(not(any(unix, windows)))]
fn new(_child: &Child) -> io::Result<Tree> {
Ok(Tree {})
}
fn kill(self) {
#[cfg(unix)]
kill_group(self.leader);
}
}
struct Secrets {
secret: String,
reply: String,
}
#[cfg(target_os = "linux")]
fn within_memory(pid: u32, memory: Option<u64>) -> io::Result<()> {
let Some(memory) = memory else {
return Ok(());
};
let status = std::fs::read_to_string(format!("/proc/{pid}/status"))?;
let peak = status
.lines()
.find_map(|line| line.strip_prefix("VmPeak:"))
.and_then(|kb| kb.trim().trim_end_matches("kB").trim().parse::<u64>().ok())
.unwrap_or(0)
* 1024;
if peak > memory {
return Err(io::Error::other(format!(
"it took {peak} bytes before they were set"
)));
}
Ok(())
}
fn connected(
listener: &TcpListener,
secrets: &Secrets,
child: &mut Child,
memory: Option<u64>,
stopping: &AtomicBool,
) -> Result<TcpStream, String> {
let until = Instant::now() + CONNECT_TIMEOUT;
let mut waiting: Vec<(TcpStream, Vec<u8>, Instant)> = Vec::new();
loop {
if stopping.load(Ordering::SeqCst) {
return Err("The server stopped.".into());
}
if let Ok(Some(status)) = child.try_wait() {
return Err(format!(
"The session's process ended before it started the app ({status}): \
it must reach the same serve() call as the server."
));
}
let now = Instant::now();
if now > until {
return Err("The session's process did not start the app in time.".into());
}
while let Ok((stream, _)) = listener.accept() {
if waiting.len() < MAX_WAITING && stream.set_nonblocking(true).is_ok() {
waiting.push((stream, Vec::new(), now + HELLO_TIMEOUT));
}
}
let mut found = None;
waiting.retain_mut(|(stream, line, by)| {
if found.is_some() {
return true;
}
match hello(stream, line, &secrets.secret) {
Some(true) => {
found = stream.try_clone().ok();
false
}
Some(false) => false,
None => now < *by,
}
});
if let Some(mut stream) = found {
#[cfg(not(target_os = "linux"))]
let _ = memory;
#[cfg(target_os = "linux")]
if let Err(error) = within_memory(child.id(), memory) {
return Err(format!(
"The session's process could not start: its limits could not be set: {error}."
));
}
let answered = stream
.set_nonblocking(false)
.and_then(|()| stream.write_all(format!("{}\n", secrets.reply).as_bytes()));
if answered.is_ok() {
let _ = stream.set_nodelay(true);
return Ok(stream);
}
continue;
}
thread::sleep(Duration::from_millis(10));
}
}
fn hello(stream: &TcpStream, line: &mut Vec<u8>, secret: &str) -> Option<bool> {
let mut byte = [0u8];
loop {
match (&*stream).read(&mut byte) {
Ok(1) if byte[0] == b'\n' => {
return Some(
std::str::from_utf8(line)
.is_ok_and(|line| crate::http::same_token(line, secret)),
);
}
Ok(1) if line.len() < secret.len() => line.push(byte[0]),
Err(e) if e.kind() == ErrorKind::WouldBlock => return None,
Err(e) if e.kind() == ErrorKind::Interrupted => {}
_ => return Some(false),
}
}
}
#[derive(Debug, PartialEq, Eq)]
enum Ended {
Page,
Child,
}
struct Relayed {
ended: Ended,
can_close: bool,
}
fn relay(browser: &TcpStream, child: &TcpStream, stopping: &AtomicBool) -> Relayed {
let failed = Relayed {
ended: Ended::Page,
can_close: true,
};
let done = Arc::new(AtomicBool::new(false));
let (Ok(page_in), Ok(child_out)) = (browser.try_clone(), child.try_clone()) else {
return failed;
};
let _ = page_in.set_read_timeout(Some(TURN));
let _ = child.set_read_timeout(Some(TURN));
let _ = child.set_write_timeout(Some(TURN));
let upstream = {
let done = done.clone();
thread::Builder::new()
.name("rich-web relay".into())
.spawn(move || {
pump(page_in, child_out, &done, &mut |_| {}, true);
done.store(true, Ordering::SeqCst);
})
};
let mut frames = Frames::default();
let (Ok(child_in), Ok(page_out)) = (child.try_clone(), browser.try_clone()) else {
done.store(true, Ordering::SeqCst);
return failed;
};
let child_ended = {
let stop = StopOr(&done, stopping);
pump_until(
child_in,
page_out,
&stop,
&mut |bytes| frames.feed(bytes),
false,
)
};
done.store(true, Ordering::SeqCst);
if let Ok(upstream) = upstream {
let _ = upstream.join();
}
Relayed {
ended: if child_ended {
Ended::Child
} else {
Ended::Page
},
can_close: frames.at_boundary() && !frames.closed,
}
}
struct StopOr<'a>(&'a AtomicBool, &'a AtomicBool);
impl StopOr<'_> {
fn is_set(&self) -> bool {
self.0.load(Ordering::SeqCst) || self.1.load(Ordering::SeqCst)
}
}
fn pump(
from: TcpStream,
to: TcpStream,
done: &AtomicBool,
seen: &mut dyn FnMut(&[u8]),
patient: bool,
) {
let never = AtomicBool::new(false);
pump_until(from, to, &StopOr(done, &never), seen, patient);
}
fn pump_until(
mut from: TcpStream,
mut to: TcpStream,
stop: &StopOr<'_>,
seen: &mut dyn FnMut(&[u8]),
patient: bool,
) -> bool {
let mut buffer = vec![0u8; 64 * 1024];
loop {
if stop.is_set() {
return false;
}
match from.read(&mut buffer) {
Ok(0) => return true,
Ok(n) => {
seen(&buffer[..n]);
if !write_until(&mut to, &buffer[..n], stop, patient) {
return false;
}
}
Err(e) if matches!(e.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut) => {}
Err(e) if e.kind() == ErrorKind::Interrupted => {}
Err(_) => return true,
}
}
}
fn write_until(to: &mut TcpStream, mut bytes: &[u8], stop: &StopOr<'_>, patient: bool) -> bool {
while !bytes.is_empty() {
match to.write(bytes) {
Ok(0) => return false,
Ok(n) => bytes = &bytes[n..],
Err(e) if e.kind() == ErrorKind::Interrupted => {}
Err(e)
if patient && matches!(e.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut) =>
{
if stop.is_set() {
return false;
}
}
Err(_) => return false,
}
}
true
}
#[derive(Default)]
struct Frames {
header: Vec<u8>,
left: u64,
closed: bool,
}
impl Frames {
fn at_boundary(&self) -> bool {
self.header.is_empty() && self.left == 0
}
fn feed(&mut self, mut bytes: &[u8]) {
while !bytes.is_empty() {
if self.left > 0 {
let n = (self.left.min(bytes.len() as u64)) as usize;
self.left -= n as u64;
bytes = &bytes[n..];
continue;
}
self.header.push(bytes[0]);
bytes = &bytes[1..];
if let Some((length, opcode)) = header_done(&self.header) {
if opcode == 0x8 {
self.closed = true;
}
self.header.clear();
self.left = length;
}
}
}
}
fn header_done(header: &[u8]) -> Option<(u64, u8)> {
if header.len() < 2 {
return None;
}
let opcode = header[0] & 0x0f;
let masked = header[1] & 0x80 != 0;
let short = u64::from(header[1] & 0x7f);
let (extended, length) = match short {
126 => (2, None),
127 => (8, None),
n => (0, Some(n)),
};
let mask = if masked { 4 } else { 0 };
if header.len() < 2 + extended + mask {
return None;
}
let length = length.unwrap_or_else(|| {
header[2..2 + extended]
.iter()
.fold(0u64, |n, b| (n << 8) | u64::from(*b))
});
Some((length, opcode))
}
fn end(child: &mut Child, tree: Tree) -> Option<std::process::ExitStatus> {
let until = Instant::now() + EXIT_WAIT;
while Instant::now() < until {
match child.try_wait() {
Ok(Some(status)) => {
tree.kill();
return Some(status);
}
Ok(None) => thread::sleep(Duration::from_millis(10)),
Err(_) => break,
}
}
tree.kill();
let _ = child.kill();
let _ = child.wait();
None
}
#[cfg(unix)]
fn kill_group(leader: u32) {
use rustix::process::{Pid, Signal, kill_process_group};
if let Some(pid) = i32::try_from(leader).ok().and_then(Pid::from_raw) {
let _ = kill_process_group(pid, Signal::KILL);
}
}
fn goodbye(mut browser: TcpStream, reason: &str) {
let mut end = reason.len().min(123);
while !reason.is_char_boundary(end) {
end -= 1;
}
let mut payload = 1000u16.to_be_bytes().to_vec();
payload.extend_from_slice(&reason.as_bytes()[..end]);
let mut frame = vec![0x88, payload.len() as u8];
frame.extend_from_slice(&payload);
let _ = browser.write_all(&frame);
let _ = browser.flush();
let _ = browser.set_read_timeout(Some(Duration::from_millis(200)));
let _ = browser.read(&mut [0u8; 256]);
let _ = browser.shutdown(Shutdown::Both);
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn specs_round_trip() {
let spec = Spec {
port: 4321,
parent: 99,
secret: "abc".into(),
reply: "def".into(),
bind: 2,
columns: 80,
rows: 24,
renderer: Renderer::Dom,
memory: Some(1 << 28),
cpu: None,
compress: Some(1024),
};
assert_eq!(Spec::decode(&spec.encode()), Some(spec));
assert_eq!(Spec::decode("1 1 a 0 80 24 xterm - -"), None);
assert_eq!(Spec::decode("3 1 9 a b 0 80 24 xterm - -"), None);
assert_eq!(Spec::decode("2 1 9 a b 0 80 24 xterm - -"), None);
assert_eq!(Spec::decode("3 1 9 a b 0 80 24 canvas - - -"), None);
}
#[test]
fn frames_are_followed_across_reads() {
let mut frames = Frames::default();
frames.feed(&[0x81, 0x03, b'a']);
assert!(!frames.at_boundary());
frames.feed(b"bc");
assert!(frames.at_boundary());
let mut big = vec![0x82, 126, 0x01, 0x2c];
big.extend(vec![0u8; 300]);
frames.feed(&big[..3]);
assert!(!frames.at_boundary());
frames.feed(&big[3..]);
assert!(frames.at_boundary());
assert!(!frames.closed);
frames.feed(&[0x88, 0x02, 0x03, 0xe8]);
assert!(frames.closed && frames.at_boundary());
}
#[test]
fn a_write_to_a_process_that_stopped_reading_gives_up_when_told() {
let listener = TcpListener::bind(("127.0.0.1", 0)).unwrap();
let mut writer = TcpStream::connect(listener.local_addr().unwrap()).unwrap();
let (_reader, _) = listener.accept().unwrap();
writer
.set_write_timeout(Some(Duration::from_millis(20)))
.unwrap();
let stopping = Arc::new(AtomicBool::new(false));
let never = AtomicBool::new(false);
let stopper = {
let stopping = stopping.clone();
thread::spawn(move || {
thread::sleep(Duration::from_millis(300));
stopping.store(true, Ordering::SeqCst);
})
};
let started = Instant::now();
let big = vec![0u8; 64 << 20];
let stop = StopOr(&never, &stopping);
assert!(!write_until(&mut writer, &big, &stop, true));
assert!(started.elapsed() < Duration::from_secs(5));
stopper.join().unwrap();
}
#[test]
fn a_masked_header_counts_its_key() {
assert_eq!(header_done(&[0x81, 0x85]), None);
assert_eq!(header_done(&[0x81, 0x85, 1, 2, 3]), None);
assert_eq!(header_done(&[0x81, 0x85, 1, 2, 3, 4]), Some((5, 1)));
}
}