use super::adapter::{Subscriber, ViewerOutbox};
use crate::ecs::{Inbound, ManagerAction, ManagerOutcome, ViewerRequest};
use crate::ids::ViewerId;
use crate::proto::attach::{
ClientMessage, FRAME_TIMEOUT, MAX_CLIENT_FRAME, MAX_INPUT_CHUNK, MAX_SERVER_FRAME,
ServerMessage, read_frame, write_frame,
};
use crate::proto::control::{
self, CONTROL_PREFACE, ErrorCode, MAX_FRAME_BYTES, MAX_SUBSCRIBER_QUEUE, Reply, Request,
};
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
use tokio::net::{UnixListener, UnixStream};
use tokio::sync::{Notify, mpsc, oneshot};
use tokio::task::JoinSet;
#[derive(Clone)]
pub struct Owner {
pub instance: String,
pub inbound: mpsc::Sender<Inbound>,
pub tokens: Arc<AtomicU64>,
pub control_replies: mpsc::Sender<(u64, oneshot::Sender<Reply>)>,
pub manager_replies: mpsc::Sender<(u64, oneshot::Sender<ManagerOutcome>)>,
pub viewer_outboxes: mpsc::Sender<(ViewerId, ViewerOutbox)>,
pub viewer_ids: Arc<AtomicU64>,
}
impl Owner {
fn token(&self) -> u64 {
self.tokens.fetch_add(1, Ordering::Relaxed)
}
}
fn authenticate(stream: UnixStream) -> std::io::Result<UnixStream> {
let stream = stream.into_std()?;
crate::proto::socket::authorize_peer(&stream)?;
UnixStream::from_std(stream)
}
async fn accept_loop<F>(
listener: UnixListener,
stop: &Notify,
admit: impl Fn(&JoinSet<()>) -> bool,
mut serve: impl FnMut(UnixStream) -> F,
) -> JoinSet<()>
where
F: Future<Output = ()> + Send + 'static,
{
let mut tasks = JoinSet::new();
loop {
tokio::select! {
() = stop.notified() => break,
Some(_) = tasks.join_next(), if !tasks.is_empty() => {}
accepted = listener.accept() => {
let Ok((stream, _)) = accepted else { break };
let Ok(stream) = authenticate(stream) else { continue };
if !admit(&tasks) {
continue;
}
tasks.spawn(serve(stream));
}
}
}
tasks
}
async fn drain(mut tasks: JoinSet<()>) {
tasks.abort_all();
while tasks.join_next().await.is_some() {}
}
pub async fn serve_attachments(
listener: UnixListener,
workspace: String,
owner: Owner,
stop: Arc<Notify>,
) {
let active = Arc::new(AtomicUsize::new(0));
let mut tasks = accept_loop(
listener,
&stop,
|_| active.load(Ordering::Acquire) < crate::proto::attach::MAX_VIEWERS_PER_WORKSPACE,
|stream| {
active.fetch_add(1, Ordering::AcqRel);
let active = Arc::clone(&active);
let owner = owner.clone();
let workspace = workspace.clone();
async move {
if let Err(error) = serve_viewer(stream, workspace, owner).await {
tracing::debug!(%error, "viewer connection ended");
}
active.fetch_sub(1, Ordering::AcqRel);
}
},
)
.await;
let grace = tokio::time::sleep(Duration::from_secs(2));
tokio::pin!(grace);
while !tasks.is_empty() {
tokio::select! {
_ = tasks.join_next() => {}
() = &mut grace => break,
}
}
drain(tasks).await;
}
async fn serve_viewer(
mut stream: UnixStream,
workspace: String,
owner: Owner,
) -> anyhow::Result<()> {
let hello: ClientMessage =
tokio::time::timeout(FRAME_TIMEOUT, read_frame(&mut stream, MAX_CLIENT_FRAME)).await??;
let (rows, cols) = match hello {
ClientMessage::Hello { rows, columns } => (rows, columns),
_ => {
write_frame(
&mut stream,
&ServerMessage::Error {
message: "the first attachment frame must be a hello".into(),
},
MAX_SERVER_FRAME,
)
.await?;
anyhow::bail!("attachment did not start with a hello");
}
};
write_frame(&mut stream, &ServerMessage::Hello {}, MAX_SERVER_FRAME).await?;
let viewer = ViewerId(owner.viewer_ids.fetch_add(1, Ordering::Relaxed));
let outbox = ViewerOutbox::default();
owner.viewer_outboxes.send((viewer, outbox.clone())).await?;
owner
.inbound
.send(Inbound::ViewerAttached {
viewer,
workspace,
rows,
cols,
})
.await?;
let (mut reader, mut writer) = stream.into_split();
let writer_outbox = outbox.clone();
let mut writer_task = tokio::spawn(async move {
while let Some(message) = writer_outbox.next().await {
let exit = matches!(
message,
ServerMessage::Exited { .. } | ServerMessage::Error { .. }
);
if write_frame(&mut writer, &message, MAX_SERVER_FRAME)
.await
.is_err()
{
break;
}
if exit {
break;
}
}
let _ = writer.shutdown().await;
});
let inbound = owner.inbound.clone();
let result = async {
loop {
let message: ClientMessage = read_frame(&mut reader, MAX_CLIENT_FRAME).await?;
let request = match message {
ClientMessage::Input { bytes } => {
anyhow::ensure!(bytes.len() <= MAX_INPUT_CHUNK, "oversized input chunk");
ViewerRequest::Input(bytes)
}
ClientMessage::Mouse { event, generation } => {
anyhow::ensure!(event.column > 0 && event.row > 0, "invalid mouse report");
ViewerRequest::Mouse { event, generation }
}
ClientMessage::Control { request } => ViewerRequest::Control(request),
ClientMessage::View {
request,
pane,
offset,
} => ViewerRequest::View {
request,
pane,
offset,
},
ClientMessage::Resize { rows, columns } => ViewerRequest::Resize {
rows,
cols: columns,
},
ClientMessage::Detach => {
inbound
.send(Inbound::ViewerRequest {
viewer,
request: ViewerRequest::Detach,
})
.await?;
return Ok::<(), anyhow::Error>(());
}
ClientMessage::Hello { .. } => anyhow::bail!("duplicate hello"),
};
inbound
.send(Inbound::ViewerRequest { viewer, request })
.await?;
}
};
let outcome = tokio::select! {
result = result => result,
_ = &mut writer_task => Ok(()),
};
let detached = outcome.is_ok();
if detached {
let _ = tokio::time::timeout(FRAME_TIMEOUT, &mut writer_task).await;
}
writer_task.abort();
let _ = owner.inbound.send(Inbound::ViewerGone { viewer }).await;
outbox.close();
outcome
}
pub async fn serve_control(
listener: UnixListener,
workspace: String,
owner: Owner,
subscribers: Arc<Mutex<Vec<Subscriber>>>,
stop: Arc<Notify>,
) {
let tasks = accept_loop(
listener,
&stop,
|tasks| tasks.len() < control::MAX_CONTROL_CONNECTIONS,
|stream| {
let owner = owner.clone();
let workspace = workspace.clone();
let subscribers = Arc::clone(&subscribers);
async move {
if let Err(error) =
serve_control_connection(stream, workspace, owner, subscribers).await
{
tracing::debug!(%error, "control connection ended");
}
}
},
)
.await;
drain(tasks).await;
}
pub async fn negotiate(stream: &mut UnixStream) -> anyhow::Result<()> {
tokio::time::timeout(Duration::from_secs(2), async {
let mut preface = [0_u8; CONTROL_PREFACE.len()];
stream.read_exact(&mut preface).await?;
stream.write_all(CONTROL_PREFACE).await?;
anyhow::ensure!(&preface == CONTROL_PREFACE, "not a fux control preface");
Ok(())
})
.await
.map_err(|_| anyhow::anyhow!("control negotiation timed out"))?
}
async fn read_line(
reader: &mut BufReader<tokio::net::unix::OwnedReadHalf>,
) -> anyhow::Result<Option<Vec<u8>>> {
let mut line = Vec::new();
let limit = MAX_FRAME_BYTES + 1;
let count = tokio::time::timeout(Duration::from_secs(30), async {
(&mut *reader)
.take(limit as u64)
.read_until(b'\n', &mut line)
.await
})
.await
.map_err(|_| anyhow::anyhow!("control frame stalled"))??;
if count == 0 {
return Ok(None);
}
if line.last() == Some(&b'\n') {
line.pop();
} else if line.len() > MAX_FRAME_BYTES {
anyhow::bail!("control frame exceeds limit");
}
Ok(Some(line))
}
async fn serve_control_connection(
mut stream: UnixStream,
workspace: String,
owner: Owner,
subscribers: Arc<Mutex<Vec<Subscriber>>>,
) -> anyhow::Result<()> {
negotiate(&mut stream).await?;
let (reader, mut writer) = stream.into_split();
let mut reader = BufReader::new(reader);
while let Some(line) = read_line(&mut reader).await? {
let request = match control::decode_request_frame(&line) {
Ok(request) => request,
Err(error) => {
write_line(&mut writer, &control::error_reply(&error)).await?;
continue;
}
};
if request
.instance()
.is_some_and(|instance| instance != owner.instance)
{
write_line(
&mut writer,
&Reply::failed(
request.id(),
control::ErrorCode::Conflict,
"server instance changed; rediscover before retrying",
),
)
.await?;
continue;
}
if let Request::Subscribe {
id, events, after, ..
} = request
{
let (sender, mut receiver) = mpsc::channel(MAX_SUBSCRIBER_QUEUE);
{
let mut active = subscribers
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
active.retain(|subscriber| !subscriber.sender.is_closed());
active.push(Subscriber {
filters: events.clone(),
sender,
bytes: Arc::new(AtomicUsize::new(0)),
});
}
let mut boundary = after;
let mut replay = Vec::new();
if let Some(after) = after {
let reply = dispatch_control(
&owner,
&workspace,
Request::Events {
id,
instance: Some(owner.instance.clone()),
after,
},
)
.await?;
match reply {
Reply::Completed {
result:
control::CommandResult::Events {
cursor,
events: entries,
},
..
} => {
boundary = Some(cursor);
replay = entries;
}
other => {
write_line(&mut writer, &bounded(other)).await?;
return Ok(());
}
}
}
write_line(&mut writer, &Reply::Accepted { id }).await?;
for entry in replay {
if events.is_empty() || events.contains(&entry.event.kind()) {
write_event(&mut writer, entry, id).await?;
}
}
let mut probe = [0_u8; 1];
loop {
tokio::select! {
event = receiver.recv() => {
let Some(event) = event else { break };
if boundary.is_some_and(|cursor| cursor.stream == event.entry.cursor.stream
&& event.entry.cursor.sequence <= cursor.sequence) {
continue;
}
write_event(&mut writer, (*event.entry).clone(), id).await?;
}
read = reader.read(&mut probe) => {
let _ = read;
break;
}
}
}
return Ok(());
}
let reply = dispatch_control(&owner, &workspace, request).await?;
write_line(&mut writer, &bounded(reply)).await?;
}
Ok(())
}
async fn dispatch_control(
owner: &Owner,
workspace: &str,
request: Request,
) -> anyhow::Result<Reply> {
let request_id = request.id();
let answer_window = match &request {
Request::Wait { timeout_ms, .. } => {
Duration::from_millis(*timeout_ms).saturating_add(Duration::from_secs(5))
}
_ => Duration::from_secs(30),
};
let token = owner.token();
let (sender, receiver) = oneshot::channel();
owner.control_replies.send((token, sender)).await?;
owner
.inbound
.send(Inbound::ControlRequest {
workspace: workspace.to_owned(),
request,
token,
})
.await?;
let reply = match tokio::time::timeout(answer_window, receiver).await {
Ok(Ok(reply)) => reply,
_ => Reply::failed(
request_id,
ErrorCode::Internal,
"control request was not answered",
),
};
Ok(reply)
}
async fn write_event(
writer: &mut tokio::net::unix::OwnedWriteHalf,
mut event: control::SequencedEvent,
id: u64,
) -> anyhow::Result<()> {
event.event = event.event.with_id(id);
let mut bytes = serde_json::to_vec(&event)?;
anyhow::ensure!(bytes.len() <= MAX_FRAME_BYTES, "event exceeds frame limit");
bytes.push(b'\n');
tokio::time::timeout(FRAME_TIMEOUT, writer.write_all(&bytes)).await??;
Ok(())
}
fn bounded(reply: Reply) -> Reply {
if serde_json::to_vec(&reply).is_ok_and(|bytes| bytes.len() <= MAX_FRAME_BYTES) {
reply
} else {
Reply::failed(
reply.id(),
ErrorCode::FrameTooLarge,
"control response exceeds the 1 MiB frame limit",
)
}
}
async fn write_line<T: serde::Serialize>(
writer: &mut tokio::net::unix::OwnedWriteHalf,
value: &T,
) -> anyhow::Result<()> {
let mut bytes = serde_json::to_vec(value)?;
bytes.push(b'\n');
tokio::time::timeout(FRAME_TIMEOUT, writer.write_all(&bytes)).await??;
Ok(())
}
pub async fn serve_manager(listener: UnixListener, owner: Owner, stop: Arc<Notify>) {
let tasks = accept_loop(
listener,
&stop,
|tasks| tasks.len() < 64,
|stream| {
let owner = owner.clone();
async move {
if let Err(error) = serve_manager_connection(stream, owner).await {
tracing::debug!(%error, "manager connection ended");
}
}
},
)
.await;
drain(tasks).await;
}
pub type DescriptorHook = Arc<dyn Fn(&str, u64) -> Option<crate::daemon::Descriptor> + Send + Sync>;
async fn serve_manager_connection(mut stream: UnixStream, owner: Owner) -> anyhow::Result<()> {
negotiate(&mut stream).await?;
let (reader, mut writer) = stream.into_split();
let mut reader = BufReader::new(reader);
let Some(line) = read_line(&mut reader).await? else {
return Ok(());
};
let request: crate::daemon::ManagerRequest = match serde_json::from_slice(&line) {
Ok(request) => request,
Err(error) => {
let reply = crate::daemon::ManagerReply::Failed {
message: format!("invalid manager request: {error}"),
};
return write_line(&mut writer, &reply).await;
}
};
let action = match request {
crate::daemon::ManagerRequest::Final { instance, pane } => {
ManagerAction::Final { instance, pane }
}
crate::daemon::ManagerRequest::Create { name } => ManagerAction::Create { name },
crate::daemon::ManagerRequest::List => ManagerAction::List,
crate::daemon::ManagerRequest::Info => ManagerAction::Info,
crate::daemon::ManagerRequest::Resolve { name } => ManagerAction::Resolve { name },
crate::daemon::ManagerRequest::Kill { name } => ManagerAction::Kill { name },
};
if let ManagerAction::Resolve { name: Some(name) }
| ManagerAction::Kill { name }
| ManagerAction::Create { name } = &action
&& let Err(error) = crate::ids::validate_workspace_name(name)
{
let reply = crate::daemon::ManagerReply::Failed {
message: error.to_string(),
};
return write_line(&mut writer, &reply).await;
}
let token = owner.token();
let (sender, receiver) = oneshot::channel();
owner.manager_replies.send((token, sender)).await?;
owner
.inbound
.send(Inbound::Manager { action, token })
.await?;
let outcome = match tokio::time::timeout(crate::daemon::MANAGER_DEADLINE, receiver).await {
Ok(Ok(outcome)) => outcome,
_ => ManagerOutcome::Failed("manager request was not answered".into()),
};
let reply = match outcome {
ManagerOutcome::Final(result) => crate::daemon::ManagerReply::Final { result },
ManagerOutcome::Names(names) => crate::daemon::ManagerReply::Names { names },
ManagerOutcome::Info(info) => crate::daemon::ManagerReply::Info { info },
ManagerOutcome::Failed(message) => crate::daemon::ManagerReply::Failed { message },
ManagerOutcome::Attach { name, stream, .. } => match (DESCRIPTOR_HOOK.get())(&name, stream)
{
Some(descriptor) => crate::daemon::ManagerReply::Attach { descriptor },
None => crate::daemon::ManagerReply::Failed {
message: "workspace descriptor unavailable".into(),
},
},
};
write_line(&mut writer, &reply).await
}
pub struct DescriptorLookup(std::sync::OnceLock<DescriptorHook>);
impl DescriptorLookup {
pub fn install(&self, hook: DescriptorHook) {
let _ = self.0.set(hook);
}
fn get(&self) -> DescriptorHook {
self.0
.get()
.cloned()
.unwrap_or_else(|| Arc::new(|_, _| None))
}
}
pub static DESCRIPTOR_HOOK: DescriptorLookup = DescriptorLookup(std::sync::OnceLock::new());
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn subscription_replay_deduplicates_queued_overlap_and_filters_both_sides()
-> anyhow::Result<()> {
use control::{CommandResult, Event, EventCursor, EventKind, SequencedEvent};
tokio::time::timeout(Duration::from_secs(5), async {
let (server, mut client) = UnixStream::pair()?;
let (inbound, mut inbound_rx) = mpsc::channel(1);
let (control_replies, mut control_rx) = mpsc::channel(1);
let (manager_replies, _manager_rx) = mpsc::channel(1);
let (viewer_outboxes, _viewer_rx) = mpsc::channel(1);
let owner = Owner {
instance: "current-server".into(), inbound,
tokens: Arc::new(AtomicU64::new(1)), control_replies,
manager_replies, viewer_outboxes,
viewer_ids: Arc::new(AtomicU64::new(1)),
};
let subscribers = Arc::new(Mutex::new(Vec::new()));
let serving = tokio::spawn(serve_control_connection(server, "default".into(), owner, subscribers.clone()));
client.write_all(CONTROL_PREFACE).await?;
let mut preface = [0; CONTROL_PREFACE.len()];
client.read_exact(&mut preface).await?;
let (client, mut writer) = client.into_split();
let mut client = BufReader::new(client);
let cursor = |sequence| EventCursor { stream: 7, sequence };
write_line(&mut writer, &Request::Subscribe {
id: 42, instance: Some("current-server".into()),
events: vec![EventKind::PaneTitle], after: Some(cursor(0)),
}).await?;
let (token, reply) = control_rx.recv().await.ok_or_else(|| anyhow::anyhow!("missing reply channel"))?;
let request = inbound_rx.recv().await.ok_or_else(|| anyhow::anyhow!("missing replay request"))?;
assert!(matches!(request, Inbound::ControlRequest { token: actual, request: Request::Events { after, .. }, .. } if actual == token && after == cursor(0)));
assert_eq!(crate::os::lock(&subscribers).len(), 1);
let title = |value: &str| Event::PaneTitle { id: 0, pane: crate::ids::PaneId(1), title: value.into() };
let publish = |event: Event, cursor: EventCursor| {
let size = crate::ecs::events::encoded_len(&SequencedEvent { cursor, event: event.clone() });
super::super::adapter::publish(&subscribers, &event, cursor, size);
};
publish(title("overlap"), cursor(2));
reply.send(Reply::Completed { id: 42, result: CommandResult::Events {
cursor: cursor(2), events: vec![
SequencedEvent { cursor: cursor(1), event: Event::WorkspaceChanged { id: 0 } },
SequencedEvent { cursor: cursor(2), event: title("overlap") },
],
}}).map_err(|_| anyhow::anyhow!("replay receiver closed"))?;
publish(Event::WorkspaceChanged { id: 0 }, cursor(3));
publish(title("later"), cursor(4));
let accepted = read_line(&mut client).await?.ok_or_else(|| anyhow::anyhow!("accept EOF"))?;
assert_eq!(serde_json::from_slice::<Reply>(&accepted)?, Reply::Accepted { id: 42 });
for (sequence, value) in [(2, "overlap"), (4, "later")] {
let line = read_line(&mut client).await?.ok_or_else(|| anyhow::anyhow!("event EOF"))?;
let entry: SequencedEvent = serde_json::from_slice(&line)?;
assert_eq!(entry.cursor, cursor(sequence));
assert_eq!(entry.event, title(value).with_id(42));
}
drop(client);
drop(writer);
serving.await??;
Ok::<_, anyhow::Error>(())
}).await??;
Ok(())
}
#[tokio::test]
async fn subscription_checks_incarnation_before_registering() -> anyhow::Result<()> {
tokio::time::timeout(Duration::from_secs(5), async {
let (server, client) = UnixStream::pair()?;
let (inbound, _inbound_rx) = mpsc::channel(1);
let (control_replies, _control_rx) = mpsc::channel(1);
let (manager_replies, _manager_rx) = mpsc::channel(1);
let (viewer_outboxes, _viewer_rx) = mpsc::channel(1);
let owner = Owner {
instance: "current-server".into(),
inbound,
tokens: Arc::new(AtomicU64::new(1)),
control_replies,
manager_replies,
viewer_outboxes,
viewer_ids: Arc::new(AtomicU64::new(1)),
};
let subscribers = Arc::new(Mutex::new(Vec::new()));
let serving = tokio::spawn(serve_control_connection(
server,
"default".into(),
owner,
subscribers.clone(),
));
let mut client = client;
client.write_all(CONTROL_PREFACE).await?;
let mut preface = [0; CONTROL_PREFACE.len()];
client.read_exact(&mut preface).await?;
assert_eq!(&preface, CONTROL_PREFACE);
let (client, mut writer) = client.into_split();
let mut client = BufReader::new(client);
for (instance, expected) in [
("old-server", ErrorCode::Conflict),
("bad instance", ErrorCode::InvalidRequest),
] {
write_line(
&mut writer,
&Request::Subscribe {
after: None,
id: 1,
instance: Some(instance.into()),
events: Vec::new(),
},
)
.await?;
let line = read_line(&mut client)
.await?
.ok_or_else(|| anyhow::anyhow!("reply EOF"))?;
let response: Reply = serde_json::from_slice(&line)?;
assert!(matches!(response, Reply::Failed { error, .. } if error.code == expected));
assert!(crate::os::lock(&subscribers).is_empty());
}
write_line(
&mut writer,
&Request::Subscribe {
after: None,
id: 2,
instance: Some("current-server".into()),
events: Vec::new(),
},
)
.await?;
let line = read_line(&mut client)
.await?
.ok_or_else(|| anyhow::anyhow!("reply EOF"))?;
assert_eq!(
serde_json::from_slice::<Reply>(&line)?,
Reply::Accepted { id: 2 }
);
assert_eq!(crate::os::lock(&subscribers).len(), 1);
drop(client);
drop(writer);
serving.await??;
Ok::<_, anyhow::Error>(())
})
.await??;
Ok(())
}
}