use super::{
assistant_state::{Hub, Pair, Seat},
assistant_wire::{self, Greeting},
config::Config,
session::{lock, session_id},
signin::{Door, Issuer},
};
use crate::serve::storage::Failure;
use serde_json::json;
use std::{
collections::BTreeMap,
io::{self, BufReader, Write},
os::unix::{
fs::{FileTypeExt, MetadataExt, PermissionsExt},
net::{UnixListener, UnixStream},
},
path::{Path, PathBuf},
sync::{Arc, Condvar, Mutex},
thread::JoinHandle,
};
#[derive(Default)]
struct Connections {
sockets: Mutex<BTreeMap<String, UnixStream>>,
released: Condvar,
}
struct Connection {
id: String,
live: Arc<Connections>,
count: Arc<std::sync::atomic::AtomicUsize>,
}
impl Drop for Connection {
fn drop(&mut self) {
self.count.fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
match self.live.sockets.lock() {
Ok(mut sockets) => {
sockets.remove(&self.id);
self.live.released.notify_all();
}
Err(_) => eprintln!("haematite assistant_connections_poisoned"),
}
}
}
pub(super) struct Listener {
stop: UnixStream,
worker: Option<JoinHandle<()>>,
live: Arc<Connections>,
owned: SocketPath,
}
impl Listener {
pub fn start<I: Issuer>(door: &Arc<Door<I>>, config: &Config) -> Result<Option<Self>, Failure> {
let Some(settings) = config.assistant.clone() else {
return Ok(None);
};
let (listener, owned) = bind(&settings.socket)?;
listener
.set_nonblocking(true)
.map_err(|error| socket_error(&error))?;
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.map_err(|error| socket_error(&error))?;
let (stop, stopped) = UnixStream::pair().map_err(|error| socket_error(&error))?;
stopped
.set_nonblocking(true)
.map_err(|error| socket_error(&error))?;
let (listener, mut stopped) = {
let entered = runtime.enter();
let listener = tokio::net::UnixListener::from_std(listener)
.map_err(|error| socket_error(&error))?;
let stopped =
tokio::net::UnixStream::from_std(stopped).map_err(|error| socket_error(&error))?;
drop(entered);
(listener, stopped)
};
let hub = Arc::new(Hub::new(settings, config.site().timeout));
let attached_hub = Arc::clone(&hub);
let live = Arc::new(Connections::default());
let serving = Serving {
live: Arc::clone(&live),
door: Arc::clone(door),
hub,
stores: Arc::new(config.stores.clone()),
maximum: config.max_connections.get(),
};
let worker=std::thread::Builder::new().name("haematite-assistant-listener".into()).spawn(move || {
use tokio::io::AsyncReadExt;
runtime.block_on(async {
loop {
let mut byte=[0];
let accepted=tokio::select! { result=stopped.read(&mut byte)=>{if let Err(error)=result {eprintln!("haematite assistant_listener_stop_failed: {error}");} break;}, result=listener.accept()=>result };
let result=accepted.map_err(|error|socket_error(&error)).and_then(|(stream,_)|serving.accept(stream));
if let Err(error)=result {eprintln!("haematite assistant_accept_failed: {}",error.code);}
}
});
}).map_err(|error|socket_error(&error))?;
let running = Self {
stop,
worker: Some(worker),
live,
owned,
};
door.set_assistant(attached_hub)?;
Ok(Some(running))
}
}
impl Drop for Listener {
fn drop(&mut self) {
if let Err(error) = self.stop.write_all(&[1]) {
eprintln!("haematite assistant_stop_failed: {error}");
}
if let Some(worker) = self.worker.take()
&& worker.join().is_err()
{
eprintln!("haematite assistant_listener_panicked");
}
match self.live.sockets.lock() {
Ok(mut sockets) => {
for socket in sockets.values() {
if let Err(error) = socket.shutdown(std::net::Shutdown::Both) {
eprintln!("haematite assistant_close_failed: {error}");
}
}
while !sockets.is_empty() {
if let Ok(next) = self.live.released.wait(sockets) {
sockets = next;
} else {
eprintln!("haematite assistant_connections_poisoned");
break;
}
}
}
Err(_) => eprintln!("haematite assistant_connections_poisoned"),
}
self.owned.release();
}
}
struct Admitted<'a> {
hub: &'a Hub,
pair: Pair,
generation: String,
counted: Option<Connection>,
}
impl Drop for Admitted<'_> {
fn drop(&mut self) {
drop(self.counted.take());
if let Err(error) = self.hub.depart(&self.pair, &self.generation) {
eprintln!("haematite assistant_seat_release_failed: {}", error.code);
}
}
}
fn connection<I: Issuer>(
door: &Door<I>,
hub: &Hub,
stores: &BTreeMap<String, PathBuf>,
stream: &mut UnixStream,
counted: Connection,
) -> Result<(), Failure> {
stream
.set_read_timeout(Some(hub.timeout))
.map_err(|error| socket_error(&error))?;
let mut reader = BufReader::new(stream.try_clone().map_err(|error| socket_error(&error))?);
let Some(value) = assistant_wire::read(&mut reader, hub.config.frame_bytes.get())? else {
return Ok(());
};
let greeting: Greeting = serde_json::from_value(value).map_err(|_| {
Failure::new(
"assistant_hello_invalid",
"hello needs agent, pass, store, resumable and harness",
)
})?;
let hello = greeting.hello;
if !stores.contains_key(&hello.store) {
return Err(Failure::new(
"store_unknown",
"the door serves no such store",
));
}
if hello.harness.id.is_empty()
|| hello.harness.version.is_empty()
|| !matches!(hello.harness.protocol, 1 | 2)
{
return Err(Failure::new(
"assistant_hello_invalid",
"harness identity and protocol are required",
));
}
let person = door
.issuer()
.assistant_identity(&hello.agent, &hello.pass, &hello.store)?;
drop(hello.pass);
let pair = (person.clone(), hello.store);
let generation = session_id()?;
hub.admit(pair.clone(),Seat {agent:hello.agent,generation:generation.clone(),socket:Arc::new(Mutex::new(stream.try_clone().map_err(|error|socket_error(&error))?)),pending:BTreeMap::new(),presence:json!({"kind":"presence","present":true,"resumable":hello.resumable,"harness":hello.harness})})?;
let admitted = Admitted {
hub,
pair,
generation,
counted: Some(counted),
};
assistant_wire::write(stream, &json!({"welcome":{"person":person}}))
.map_err(|error| socket_error(&error))?;
stream
.set_read_timeout(None)
.map_err(|error| socket_error(&error))?;
while let Some(frame) = assistant_wire::read(&mut reader, hub.config.frame_bytes.get())? {
let object = frame
.as_object()
.ok_or_else(|| Failure::new("assistant_frame_invalid", "expected one frame kind"))?;
if object.len() != 1 {
return Err(Failure::new(
"assistant_frame_invalid",
"expected one frame kind",
));
}
if let Some(item) = frame.get("item") {
assistant_wire::item(item)?;
hub.item(&admitted.pair, item)?;
} else if let Some(ack) = frame.get("ack") {
hub.acknowledge(
&admitted.pair,
ack["id"].as_str().ok_or_else(|| {
Failure::new("assistant_ack_invalid", "ack needs its request id")
})?,
Ok(()),
)?;
} else if let Some(refused) = frame.get("refused") {
let (Some(id), Some(code), Some(message)) = (
refused["id"].as_str(),
refused["code"].as_str(),
refused["message"].as_str(),
) else {
return Err(Failure::new(
"assistant_ack_invalid",
"refusal needs id, code and message",
));
};
hub.acknowledge(&admitted.pair, id, Err(Failure::new(code, message)))?;
} else {
return Err(Failure::new(
"assistant_frame_invalid",
"expected ack, refused or item",
));
}
}
Ok(())
}
fn socket_error(error: &io::Error) -> Failure {
Failure::new("assistant_socket_failed", error.to_string())
}
struct Serving<I> {
door: Arc<Door<I>>,
hub: Arc<Hub>,
stores: Arc<BTreeMap<String, PathBuf>>,
live: Arc<Connections>,
maximum: usize,
}
impl<I: Issuer> Serving<I> {
fn accept(&self, stream: tokio::net::UnixStream) -> Result<(), Failure> {
let mut stream = stream.into_std().map_err(|error| socket_error(&error))?;
stream
.set_nonblocking(false)
.map_err(|error| socket_error(&error))?;
stream
.set_write_timeout(Some(self.hub.timeout))
.map_err(|error| socket_error(&error))?;
let id = session_id()?;
let disconnect = stream.try_clone().map_err(|error| socket_error(&error))?;
let mut sockets = lock(&self.live.sockets)?;
let count = self
.door
.connections
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
if count >= self.maximum {
self.door
.connections
.fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
assistant_wire::refuse(
&mut stream,
&super::assistant_state::bound("max_connections", self.maximum, count),
)
.map_err(|error| socket_error(&error))?;
return Ok(());
}
sockets.insert(id.clone(), disconnect);
drop(sockets);
let counted = Connection {
id,
live: Arc::clone(&self.live),
count: Arc::clone(&self.door.connections),
};
let (door, hub, stores) = (
Arc::clone(&self.door),
Arc::clone(&self.hub),
Arc::clone(&self.stores),
);
std::thread::Builder::new()
.name("haematite-assistant-seat".into())
.spawn(move || {
if let Err(failure) = connection(&door, &hub, &stores, &mut stream, counted)
&& let Err(error) = assistant_wire::refuse(&mut stream, &failure)
&& !matches!(
error.kind(),
io::ErrorKind::BrokenPipe
| io::ErrorKind::ConnectionReset
| io::ErrorKind::NotConnected
)
{
eprintln!("haematite assistant_refusal_write_failed: {error}");
}
})
.map_err(|error| socket_error(&error))?;
Ok(())
}
}
struct SocketPath {
path: Option<PathBuf>,
device: u64,
inode: u64,
}
impl SocketPath {
fn release(&mut self) {
let Some(path) = self.path.take() else {
return;
};
match std::fs::symlink_metadata(&path) {
Ok(metadata) if metadata.dev() == self.device && metadata.ino() == self.inode => {
if let Err(error) = std::fs::remove_file(&path) {
eprintln!("haematite assistant_socket_remove_failed: {error}");
}
}
Ok(_) => eprintln!("haematite assistant_socket_changed_at_close"),
Err(error) if error.kind() == io::ErrorKind::NotFound => {}
Err(error) => eprintln!("haematite assistant_socket_stat_failed: {error}"),
}
}
}
impl Drop for SocketPath {
fn drop(&mut self) {
self.release();
}
}
fn bind(path: &Path) -> Result<(UnixListener, SocketPath), Failure> {
if !path.is_absolute() {
return Err(Failure::new(
"assistant_socket_invalid",
"assistant.socket must be absolute",
));
}
match std::fs::symlink_metadata(path) {
Ok(metadata) => {
if !metadata.file_type().is_socket()
|| metadata.uid() != rustix::process::geteuid().as_raw()
{
return Err(Failure::new(
"assistant_socket_exists",
"the existing path is not an owned socket",
));
}
match UnixStream::connect(path) {
Ok(_) => {
return Err(Failure::new(
"assistant_socket_listening",
"an existing listener answers",
));
}
Err(error) if error.kind() == io::ErrorKind::ConnectionRefused => {}
Err(_) => {
return Err(Failure::new(
"assistant_socket_unchecked",
"the existing listener could not be checked",
));
}
}
let checked = std::fs::symlink_metadata(path).map_err(|error| socket_error(&error))?;
if checked.dev() != metadata.dev() || checked.ino() != metadata.ino() {
return Err(Failure::new(
"assistant_socket_changed",
"the socket changed while checking its listener",
));
}
std::fs::remove_file(path).map_err(|error| socket_error(&error))?;
}
Err(error) if error.kind() == io::ErrorKind::NotFound => {}
Err(error) => return Err(socket_error(&error)),
}
let listener = UnixListener::bind(path).map_err(|error| socket_error(&error))?;
let metadata = std::fs::symlink_metadata(path).map_err(|error| socket_error(&error))?;
let owned = SocketPath {
path: Some(path.to_path_buf()),
device: metadata.dev(),
inode: metadata.ino(),
};
std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600))
.map_err(|error| socket_error(&error))?;
Ok((listener, owned))
}