use super::{config::Seat, handshake::Session, tools::Identity};
use serde_json::{Value, json};
use tokio::io::{AsyncBufRead, AsyncBufReadExt as _, AsyncWriteExt as _, BufReader};
use tokio::net::UnixStream;
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
enum Pending {
Frame(Value),
Flush(tokio::sync::oneshot::Sender<()>),
}
#[derive(Clone)]
pub(super) struct Output {
sender: mpsc::Sender<Pending>,
bound: usize,
}
impl Output {
pub async fn flush(&self) -> anyhow::Result<()> {
let (sent, received) = tokio::sync::oneshot::channel();
self.sender
.send(Pending::Flush(sent))
.await
.map_err(|_| anyhow::anyhow!("door_writer_closed"))?;
received
.await
.map_err(|_| anyhow::anyhow!("door_flush_failed"))
}
pub fn send(&self, value: Value) -> anyhow::Result<()> {
anyhow::ensure!(
serde_json::to_vec(&value)?.len().saturating_add(1) <= self.bound,
"door_frame_bytes_exceeded"
);
self.sender
.try_send(Pending::Frame(value))
.map_err(|error| match error {
mpsc::error::TrySendError::Full(_) => anyhow::anyhow!("door_write_queue_full"),
mpsc::error::TrySendError::Closed(_) => anyhow::anyhow!("door_writer_closed"),
})
}
}
pub(super) struct Link {
pub person: String,
pub input: mpsc::Receiver<anyhow::Result<Option<Value>>>,
pub output: Output,
pub reader: JoinHandle<()>,
pub writer: JoinHandle<anyhow::Result<()>>,
}
impl Link {
pub async fn connect(
seat: &Seat,
identity: &Identity,
session: &Session,
) -> anyhow::Result<Self> {
use std::os::unix::fs::{FileTypeExt as _, MetadataExt as _, PermissionsExt as _};
let metadata = std::fs::symlink_metadata(&seat.door_socket)?;
anyhow::ensure!(
metadata.file_type().is_socket()
&& metadata.permissions().mode() & 0o777 == 0o600
&& metadata.uid() == rustix::process::geteuid().as_raw(),
"assistant_socket_requires_owner_and_mode_0600"
);
let stream = UnixStream::connect(&seat.door_socket).await?;
anyhow::ensure!(
stream.peer_cred()?.uid() == rustix::process::geteuid().as_raw(),
"assistant_socket_owner_differs"
);
let (reader, mut writer) = stream.into_split();
let bound = seat.harness.acp_frame_bytes;
let (outgoing, mut pending) = mpsc::channel::<Pending>(seat.harness.write_queue);
let writer = tokio::spawn(async move {
while let Some(value) = pending.recv().await {
let value = match value {
Pending::Frame(value) => value,
Pending::Flush(receipt) => {
writer.flush().await?;
receipt
.send(())
.map_err(|()| anyhow::anyhow!("door_flush_receiver_closed"))?;
continue;
}
};
let mut bytes = serde_json::to_vec(&value)?;
bytes.push(b'\n');
writer.write_all(&bytes).await?;
writer.flush().await?;
}
Ok(())
});
let (events, input) = mpsc::channel(1);
let reader = tokio::spawn(async move {
let mut reader = BufReader::new(reader);
loop {
let value = read(&mut reader, bound).await;
let ended = !matches!(value, Ok(Some(_)));
if events.send(value).await.is_err() || ended {
return;
}
}
});
let link = Self {
person: String::new(),
input,
output: Output {
sender: outgoing,
bound,
},
reader,
writer,
};
link.output.send(json!({"hello":{"agent":identity.agent,"pass":identity.pass,"store":seat.store,
"resumable":session.resumable,"harness":{"id":seat.harness.id,"protocol":session.version.number(),
"version":session.info["version"].as_str().unwrap_or("not-advertised")}}}))?;
Ok(link)
}
#[cfg(test)]
pub async fn open(seat: &Seat, identity: &Identity, session: &Session) -> anyhow::Result<Self> {
let mut link = Self::connect(seat, identity, session).await?;
link.welcome(seat, identity).await?;
Ok(link)
}
pub async fn welcome(&mut self, seat: &Seat, identity: &Identity) -> anyhow::Result<()> {
let welcome = self
.input
.recv()
.await
.ok_or_else(|| anyhow::anyhow!("assistant_door_closed"))??
.ok_or_else(|| anyhow::anyhow!("assistant_door_closed"))?;
if let Some(refused) = welcome.get("refused") {
let code = refused["code"]
.as_str()
.filter(|code| {
!code.is_empty()
&& code.len() <= 96
&& !code.contains(&identity.pass)
&& code
.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || byte == b'_')
})
.unwrap_or("unnamed_refusal");
anyhow::bail!("assistant_door_refused: {code}");
}
let person = welcome["welcome"]["person"]
.as_str()
.filter(|person| !person.is_empty() && person.len() <= seat.harness.content_bytes)
.ok_or_else(|| anyhow::anyhow!("assistant_door_welcome_invalid"))?;
person.clone_into(&mut self.person);
Ok(())
}
}
impl Drop for Link {
fn drop(&mut self) {
self.reader.abort();
self.writer.abort();
}
}
async fn read<R: AsyncBufRead + Unpin>(
reader: &mut R,
bound: usize,
) -> anyhow::Result<Option<Value>> {
let mut line = Vec::new();
loop {
let buffer = reader.fill_buf().await?;
if buffer.is_empty() {
anyhow::ensure!(line.is_empty(), "door_frame_missing_newline");
return Ok(None);
}
let count = buffer
.iter()
.position(|byte| *byte == b'\n')
.map_or(buffer.len(), |index| index + 1);
anyhow::ensure!(
count <= bound.saturating_sub(line.len()),
"door_frame_bytes_exceeded"
);
let ended = buffer[count - 1] == b'\n';
line.extend_from_slice(&buffer[..count]);
reader.consume(count);
if ended {
break;
}
}
let value: Value =
serde_json::from_slice(&line).map_err(|_| anyhow::anyhow!("door_frame_invalid_json"))?;
anyhow::ensure!(value.is_object(), "door_frame_requires_object");
Ok(Some(value))
}