Checkpoint

This commit is contained in:
2026-08-02 18:30:34 +00:00
parent 596af4ed80
commit 6c33bcfb27
41 changed files with 6360 additions and 1 deletions
+12
View File
@@ -0,0 +1,12 @@
[package]
name = "nye-wire-protocol"
version = "0.1.0"
edition = "2024"
[dependencies]
anyhow = { version = "1.0.104", features = ["backtrace"] }
bitcode = "0.6.9"
papaya = "0.2.4"
rustix = { version = "1.1.4", features = ["linux_latest", "net", "system"] }
tokio = { version = "1.53.1", features = ["io-util", "net"] }
tracing = { version = "0.1.44", features = ["async-await"] }
+550
View File
@@ -0,0 +1,550 @@
use std::path::Path;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use anyhow::Context;
use papaya::HashMap;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::unix::{OwnedReadHalf, OwnedWriteHalf, UCred};
use tokio::net::{UnixListener, UnixStream};
use tokio::sync::{Notify, mpsc};
use crate::messages::{Message, MessageRequest, MessageResponse, StreamMessage};
pub mod messages;
/// Binds a new [`UnixListener`] to the specified path.
///
/// Arguments:
/// * `path` - The path to bind the UnixListener to.
///
/// Returns:
/// * `Ok(Self)` - If the binding was successful, returns a new instance of [`Listener`].
/// * `Err(anyhow::Error)` - If the binding failed, returns an error.
pub async fn listen(path: impl AsRef<Path>) -> anyhow::Result<Listener> {
let listener = UnixListener::bind(path.as_ref()).context(format!(
"Could not listen for incoming connections at {}.",
path.as_ref().display()
))?;
tracing::info!(
"Listening for incoming connections at {}",
path.as_ref().display()
);
Ok(Listener { listener })
}
/// Connects to a Unix socket at the specified path and returns a new instance of [`Connection`].
///
/// Arguments:
/// * `path` - The path to the Unix socket to connect to.
///
/// Returns:
/// * `Ok(Connection)` - If the connection was successful, returns a new instance of
/// [`Connection`].
/// * `Err(anyhow::Error)` - If the connection failed, returns an error.
pub async fn connect(path: impl AsRef<Path>) -> anyhow::Result<Connection> {
let stream = UnixStream::connect(path.as_ref()).await.context(format!(
"Could not connect to the Unix socket at {}.",
path.as_ref().display()
))?;
tracing::info!(
"Connected to the Unix socket at {}",
path.as_ref().display()
);
Connection::new(stream, PartyType::Client)
}
/// A wrapper around a [`UnixListener`] to accept incoming connections on a Unix socket.
pub struct Listener {
listener: UnixListener,
}
impl Listener {
/// Accepts an incoming connection on the Unix socket.
///
/// Returns:
/// * `Ok(Connection)` - If the connection was accepted successfully, returns a new instance of
/// [`Connection`].
/// * `Err(anyhow::Error)` - If the connection could not be accepted, returns an error.
pub async fn accept(&self) -> anyhow::Result<Connection> {
let (stream, _) = self.listener.accept().await?;
Connection::new(stream, PartyType::Server)
}
}
/// An enum representing the type of party (client or server) in a connection.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum PartyType {
Client,
Server,
}
impl PartyType {
/// Creates a new instance of [`PartyType`] from a given ID.
///
/// Arguments:
/// * `id` - The ID to derive the party type from.
///
/// Returns:
/// [`PartyType`] - A new instance of [`PartyType`] derived from the given ID.
fn from_id(id: u64) -> Self {
// The most significant bit of the ID indicates the party type.
if id & (1 << 63) != 0 {
PartyType::Server
} else {
PartyType::Client
}
}
/// Converts a given ID to a party-specific ID based on the party type.
///
/// Arguments:
/// * `id` - The ID to convert to a party-specific ID.
///
/// Returns:
/// [`u64`] - The party-specific ID derived from the given ID.
fn type_id(&self, id: u64) -> u64 {
match self {
PartyType::Client => id & !(1 << 63),
PartyType::Server => id | (1 << 63),
}
}
}
type StreamsMap = Arc<HashMap<u64, mpsc::UnboundedSender<Message>>>;
type PendingStreamsTx = mpsc::UnboundedSender<Stream>;
type PendingStreamsRx = mpsc::UnboundedReceiver<Stream>;
/// A wrapper around a [`UnixStream`] to represent a connection to a Unix socket.
pub struct Connection {
/// The type of party (client or server) of the current party.
party_type: PartyType,
/// A shared map of protocol streams to route incoming messages to.
streams: StreamsMap,
/// The peer credentials of the other party in the connection.
peer_credentials: UCred,
/// An atomic counter to generate unique stream IDs for each protocol stream.
stream_id_counter: AtomicU64,
/// Notifies when the connection is closing, allowing tasks to clean up resources.
closing: Arc<Notify>,
/// A channel to receive pending streams that need to be processed by this party.
pending_streams: PendingStreamsRx,
/// A channel to send outgoing messages to the other party.
outgoing_tx: mpsc::UnboundedSender<StreamMessage>,
}
impl Connection {
/// Creates a new instance of [`Connection`] from an existing [`UnixStream`].
///
/// Arguments:
/// * `party_type` - The type of party (client or server) that this connection represents.
/// * `stream` - The existing [`UnixStream`] to wrap in a [`Connection`].
///
/// Returns:
/// * `Self` - A new instance of [`Connection`] wrapping the provided [`UnixStream`].
fn new(stream: UnixStream, party_type: PartyType) -> anyhow::Result<Self> {
let peer_credentials = stream
.peer_cred()
.context("Could not get the credentials of the connection's peer.")?;
let (stream_rx, stream_tx) = stream.into_split();
let streams: StreamsMap = Default::default();
let (outgoing_tx, outgoing_rx) = mpsc::unbounded_channel();
let (pending_streams_tx, pending_streams_rx) = mpsc::unbounded_channel();
let closing = Arc::new(Notify::new());
tokio::spawn({
let stream_rx = stream_rx;
let streams = streams.clone();
let pending_streams_tx = pending_streams_tx.clone();
let outgoing_tx = outgoing_tx.clone();
let closing = closing.clone();
async move {
let result = Self::route_incoming(
stream_rx,
streams,
pending_streams_tx,
party_type,
outgoing_tx,
closing,
)
.await;
if let Err(e) = &result {
tracing::warn!("{e:?}");
}
result
}
});
tokio::spawn(async move {
let result = Self::route_outgoing(stream_tx, outgoing_rx).await;
if let Err(e) = &result {
tracing::warn!("{e:?}");
}
result
});
Ok(Self {
party_type,
peer_credentials,
streams,
stream_id_counter: AtomicU64::new(0),
closing,
pending_streams: pending_streams_rx,
outgoing_tx,
})
}
/// Receives a pending stream that needs to be processed by this party.
///
/// Returns:
/// * `Ok(Stream)` - If a pending stream was received successfully, returns the stream.
/// * `Err(anyhow::Error)` - If the pending streams channel was closed, returns an error.
pub async fn recv_stream(&mut self) -> anyhow::Result<Stream> {
let stream = self
.pending_streams
.recv()
.await
.context("The pending streams channel was closed, which should not happen.")?;
Ok(stream)
}
/// Creates a new protocol stream with a unique ID and returns it.
///
/// Returns:
/// * `Ok(Stream)` - If the stream was created successfully, returns the new stream.
/// * `Err(anyhow::Error)` - If an error occurred while creating the stream, returns an error.
pub async fn send_stream(&self) -> anyhow::Result<Stream> {
let id = self.stream_id_counter.fetch_add(1, Ordering::Relaxed);
let id = self.party_type.type_id(id);
let (incoming_tx, incoming_rx) = mpsc::unbounded_channel();
let stream = Stream::new(
id,
self.streams.clone(),
incoming_tx,
incoming_rx,
self.outgoing_tx.clone(),
);
self.streams.pin().insert(id, stream.incoming_tx.clone());
Ok(stream)
}
/// Routes incoming messages from the Unix stream to the appropriate protocol streams.
///
/// Arguments:
/// * `stream_rx` - The read half of the Unix stream to read incoming messages from.
/// * `streams` - A shared map of protocol streams to route incoming messages to.
/// * `pending` - A channel to send pending streams that need to be processed by this party.
/// * `party_type` - The party type of the current party.
/// * `outgoing_tx` - A channel to send outgoing messages to the other party.
/// * `closing` - A notification to signal when the connection is closing.
///
/// Returns:
/// * `!` - This function runs indefinitely and does not return.
/// * `Err(anyhow::Error)` - If an error occurs while routing incoming messages, returns an
/// error.
async fn route_incoming(
mut stream_rx: OwnedReadHalf,
streams: StreamsMap,
pending: PendingStreamsTx,
party_type: PartyType,
outgoing_tx: mpsc::UnboundedSender<StreamMessage>,
closing: Arc<Notify>,
) -> anyhow::Result<()> {
loop {
// Only close on message boundaries.
let message_size = tokio::select! {
_ = closing.notified() => {
// The connection is closing, so we should stop routing incoming messages and
// return.
return Ok(());
}
size = stream_rx.read_u32() => {
size.context("Could not read the size of the incoming message.")?
}
};
let mut message_data = vec![0u8; message_size as usize];
stream_rx
.read_exact(&mut message_data)
.await
.context("Could not read the data of the incoming message.")?;
let message: StreamMessage =
bitcode::decode(&message_data).context("Could not decode the incoming message.")?;
tracing::info!(
stream.id = message.stream_id,
message.kind = ?message.data.kind(),
"Received a message"
);
if let Some(stream) = streams.pin().get(&message.stream_id) {
tracing::trace!(
stream.id = message.stream_id,
message.kind = ?message.data.kind(),
"Routing the message to the existing stream"
);
stream.send(message.data).context(
"Could not send the message through the stream's incoming messages channel.",
)?;
} else if PartyType::from_id(message.stream_id) != party_type {
// The other party initiated this stream, so we need to create a new stream and
// send it to the pending channel.
tracing::trace!(
stream.id = message.stream_id,
message.kind = ?message.data.kind(),
"Creating a new stream for the incoming message"
);
let (incoming_tx, incoming_rx) = mpsc::unbounded_channel();
let stream = Stream::new(
message.stream_id,
streams.clone(),
incoming_tx.clone(),
incoming_rx,
outgoing_tx.clone(),
);
stream.incoming_tx.send(message.data).context(concat!(
"Could not send the first message of the stream through the incoming ",
"messages channel."
))?;
streams.pin().insert(message.stream_id, incoming_tx);
pending
.send(stream)
.context("Could not send the stream through the pending streams channel.")?;
} else {
anyhow::bail!(concat!(
"A message was received from the other party for a stream that does not ",
"exist. The ID said we created the stream, yet we did not or it was already ",
"closed. This is a bug, either on this side or the other side (daemon or ",
"client)."
));
}
}
}
/// Routes outgoing messages from the protocol streams to the Unix stream.
///
/// Arguments:
/// * `stream_tx` - The write half of the Unix stream to write outgoing messages to.
/// * `outgoing_rx` - A channel to receive outgoing messages from the protocol streams.
///
/// Returns:
/// * `!` - This function runs indefinitely and does not return.
/// * `Err(anyhow::Error)` - If an error occurs while routing outgoing messages, returns an
/// error.
async fn route_outgoing(
mut stream_tx: OwnedWriteHalf,
mut outgoing_rx: mpsc::UnboundedReceiver<StreamMessage>,
) -> anyhow::Result<()> {
loop {
// If closed, the connection was dropped.
let Some(message) = outgoing_rx.recv().await else {
return Ok(());
};
let message_data = bitcode::encode(&message);
let message_size = message_data.len() as u32;
stream_tx
.write_u32(message_size)
.await
.context("Could not send outgoing message's size through the socket.")?;
stream_tx
.write_all(&message_data)
.await
.context("Could not send outgoing message's data through the socket.")?;
tracing::info!(
stream.id = message.stream_id,
message.kind = ?message.data.kind(),
"Sent a message"
);
}
}
/// Returns the peer credentials of the other party in the connection.
///
/// Returns:
/// [`UCred`] - A reference to the peer credentials of the other party in the connection.
pub fn peer_credentials(&self) -> UCred {
self.peer_credentials
}
}
impl Drop for Connection {
fn drop(&mut self) {
let peer_credentials = self.peer_credentials();
tracing::debug!(
user.uid = peer_credentials.uid(),
user.gid = peer_credentials.gid(),
user.pid = peer_credentials.pid(),
"Closing the connection"
);
self.closing.notify_waiters();
}
}
/// A representation of a protocol stream that can send and receive messages over a connection.
pub struct Stream {
id: u64,
streams: StreamsMap,
incoming_tx: mpsc::UnboundedSender<Message>,
incoming_rx: mpsc::UnboundedReceiver<Message>,
outgoing_tx: mpsc::UnboundedSender<StreamMessage>,
}
impl Stream {
/// Creates a new instance of [`Stream`] with the specified ID.
///
/// Arguments:
/// * `id` - The ID to assign to the new stream.
/// * `streams` - A shared map of protocol streams to route incoming messages to.
/// * `incoming_rx` - A channel to receive incoming messages from the other party.
/// * `incoming_tx` - A channel to receive incoming messages from the other party.
/// * `outgoing_tx` - A channel to send outgoing messages to the other party.
///
/// Returns:
/// [`Stream`] - A new instance of [`Stream`] with the specified ID.
fn new(
id: u64,
streams: StreamsMap,
incoming_tx: mpsc::UnboundedSender<Message>,
incoming_rx: mpsc::UnboundedReceiver<Message>,
outgoing_tx: mpsc::UnboundedSender<StreamMessage>,
) -> Self {
Self {
id,
streams,
incoming_tx,
incoming_rx,
outgoing_tx,
}
}
/// Sends a message over the stream to the other party.
///
/// Arguments:
/// * `message` - The message to send over the stream.
///
/// Returns:
/// * `Ok(())` - If the message was sent successfully.
/// * `Err(anyhow::Error)` - If an error occurred while sending the message.
pub async fn send(&self, message: impl Into<Message>) -> anyhow::Result<()> {
let stream_message = StreamMessage {
stream_id: self.id,
data: message.into(),
};
self.outgoing_tx
.send(stream_message)
.context("The outgoing messages channel was closed.")?;
Ok(())
}
/// Receives a message from the stream sent by the other party.
///
/// Returns:
/// * `Ok(Message)` - If a message was received successfully.
/// * `Err(anyhow::Error)` - If an error occurred while receiving the message.
pub async fn recv(&mut self) -> anyhow::Result<Message> {
let message = self.incoming_rx.recv().await.context(
"The incoming messages channel for this stream was closed, which should not happen.",
)?;
Ok(message)
}
/// Receives a request message from the stream sent by the other party.
///
/// Returns:
/// * `Ok(MessageRequest)` - If a request message was received successfully.
/// * `Err(anyhow::Error)` - If an error occurred while receiving the request message or the
/// message received was a response.
pub async fn recv_request(&mut self) -> anyhow::Result<MessageRequest> {
let message = self.recv().await?;
match message {
Message::Request(request) => Ok(request),
_ => anyhow::bail!(concat!(
"The message received from the other party was not a request. ",
"This is a bug, either on this side or the other side (daemon or client)."
)),
}
}
/// Receives a response message from the stream sent by the other party.
///
/// Returns:
/// * `Ok(MessageResponse)` - If a response message was received successfully.
/// * `Err(anyhow::Error)` - If an error occurred while receiving the response message or the
/// message received was a request.
pub async fn recv_response(&mut self) -> anyhow::Result<MessageResponse> {
let message = self.recv().await?;
match message {
Message::Response(response) => Ok(response),
_ => anyhow::bail!(concat!(
"The message received from the other party was not a response. ",
"This is a bug, either on this side or the other side (daemon or client)."
)),
}
}
/// Returns the ID of the stream.
///
/// Returns:
/// [`u64`] - The ID of the stream.
pub fn id(&self) -> u64 {
self.id
}
}
impl Drop for Stream {
fn drop(&mut self) {
let id = self.id;
tracing::debug!(
stream.id = id,
"Closing the stream and removing it from the streams map."
);
self.streams.pin().remove(&id);
}
}
+105
View File
@@ -0,0 +1,105 @@
use bitcode::{Decode, Encode};
/// Represents a message in the Nye wire protocol.
#[derive(Debug, Clone, Encode, Decode)]
pub struct StreamMessage {
/// The unique identifier of the stream, used to correlate requests and responses.
pub stream_id: u64,
/// The type of the message, represented as a string.
pub data: Message,
}
/// Represents the data contained in a message in the Nye wire protocol.
#[derive(Debug, Clone, Encode, Decode)]
pub enum Message {
/// Represents a request message in the Nye wire protocol.
Request(MessageRequest),
/// Represents a response message in the Nye wire protocol.
Response(MessageResponse),
}
impl Message {
/// Returns the kind of the message.
///
/// Returns:
/// [`MessageKind`] - the kind of the message.
pub fn kind(&self) -> MessageKind {
match self {
Message::Request(request) => MessageKind::Request(request.kind()),
Message::Response(response) => MessageKind::Response(response.kind()),
}
}
}
/// Represents the kind of a message in the Nye wire protocol.
#[derive(Debug, Clone, Encode, Decode, PartialEq, Eq)]
pub enum MessageKind {
Request(MessageRequestKind),
Response(MessageResponseKind),
}
/// Represents a request message in the Nye wire protocol.
///
/// Only sent from the client to the daemon.
#[derive(Debug, Clone, Encode, Decode)]
pub enum MessageRequest {
Setup,
}
impl MessageRequest {
/// Returns the kind of the message request.
///
/// Returns:
/// [`MessageRequestKind`] - the kind of the message request.
pub fn kind(&self) -> MessageRequestKind {
match self {
MessageRequest::Setup => MessageRequestKind::Setup,
}
}
}
impl From<MessageRequest> for Message {
fn from(request: MessageRequest) -> Self {
Message::Request(request)
}
}
/// Represents the kind of a message request in the Nye wire protocol.
#[derive(Debug, Clone, Encode, Decode, PartialEq, Eq)]
pub enum MessageRequestKind {
Setup,
}
/// Represents a response message in the Nye wire protocol.
///
/// Only sent from the daemon to the client.
#[derive(Debug, Clone, Encode, Decode)]
pub enum MessageResponse {
Setup,
}
impl MessageResponse {
/// Returns the kind of the message response.
///
/// Returns:
/// [`MessageResponseKind`] - the kind of the message response.
pub fn kind(&self) -> MessageResponseKind {
match self {
MessageResponse::Setup => MessageResponseKind::SetupFinished,
}
}
}
impl From<MessageResponse> for Message {
fn from(response: MessageResponse) -> Self {
Message::Response(response)
}
}
/// Represents the kind of a message response in the Nye wire protocol.
#[derive(Debug, Clone, Encode, Decode, PartialEq, Eq)]
pub enum MessageResponseKind {
SetupFinished,
}