zlink_core/connection/

mod.rs

//! Contains connection related API.

#[cfg(feature = "std")]
mod credentials;
mod read_connection;
#[cfg(feature = "std")]
pub use credentials::Credentials;
pub use read_connection::ReadConnection;
#[cfg(feature = "std")]
pub use rustix::{process::Pid, process::Uid};
pub mod chain;
pub mod socket;
#[cfg(test)]
mod tests;
mod write_connection;
use crate::{
    reply::{self, Reply},
    Call, Result,
};
#[cfg(feature = "std")]
use alloc::vec;
pub use chain::Chain;
use core::{fmt::Debug, sync::atomic::AtomicUsize};
#[cfg(feature = "std")]
use socket::FetchPeerCredentials;
pub use write_connection::WriteConnection;

use serde::{Deserialize, Serialize};
pub use socket::Socket;

// Type alias for receive methods - std returns FDs, no_std doesn't
#[cfg(feature = "std")]
type RecvResult<T> = (T, Vec<std::os::fd::OwnedFd>);
#[cfg(not(feature = "std"))]
type RecvResult<T> = T;

/// A connection.
///
/// The low-level API to send and receive messages.
///
/// Each connection gets a unique identifier when created that can be queried using
/// [`Connection::id`]. This ID is shared betwen the read and write halves of the connection. It
/// can be used to associate the read and write halves of the same connection.
///
/// # Cancel safety
///
/// All async methods of this type are cancel safe unless explicitly stated otherwise in its
/// documentation.
#[derive(Debug)]
pub struct Connection<S: Socket> {
    read: ReadConnection<S::ReadHalf>,
    write: WriteConnection<S::WriteHalf>,
    #[cfg(feature = "std")]
    credentials: Option<std::sync::Arc<Credentials>>,
}

impl<S> Connection<S>
where
    S: Socket,
{
    /// Create a new connection.
    pub fn new(socket: S) -> Self {
        let (read, write) = socket.split();
        let id = NEXT_ID.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
        Self {
            read: ReadConnection::new(read, id),
            write: WriteConnection::new(write, id),
            #[cfg(feature = "std")]
            credentials: None,
        }
    }

    /// The reference to the read half of the connection.
    pub fn read(&self) -> &ReadConnection<S::ReadHalf> {
        &self.read
    }

    /// The mutable reference to the read half of the connection.
    pub fn read_mut(&mut self) -> &mut ReadConnection<S::ReadHalf> {
        &mut self.read
    }

    /// The reference to the write half of the connection.
    pub fn write(&self) -> &WriteConnection<S::WriteHalf> {
        &self.write
    }

    /// The mutable reference to the write half of the connection.
    pub fn write_mut(&mut self) -> &mut WriteConnection<S::WriteHalf> {
        &mut self.write
    }

    /// Split the connection into read and write halves.
    ///
    /// Note: This consumes any cached credentials. If you need the credentials after splitting,
    /// call [`Connection::peer_credentials`] before splitting.
    pub fn split(self) -> (ReadConnection<S::ReadHalf>, WriteConnection<S::WriteHalf>) {
        (self.read, self.write)
    }

    /// Join the read and write halves into a connection (the opposite of [`Connection::split`]).
    pub fn join(read: ReadConnection<S::ReadHalf>, write: WriteConnection<S::WriteHalf>) -> Self {
        Self {
            read,
            write,
            #[cfg(feature = "std")]
            credentials: None,
        }
    }

    /// The unique identifier of the connection.
    pub fn id(&self) -> usize {
        assert_eq!(self.read.id(), self.write.id());
        self.read.id()
    }

    /// Sends a method call.
    ///
    /// Convenience wrapper around [`WriteConnection::send_call`].
    pub async fn send_call<Method>(
        &mut self,
        call: &Call<Method>,
        #[cfg(feature = "std")] fds: Vec<std::os::fd::OwnedFd>,
    ) -> Result<()>
    where
        Method: Serialize + Debug,
    {
        #[cfg(feature = "std")]
        {
            self.write.send_call(call, fds).await
        }
        #[cfg(not(feature = "std"))]
        {
            self.write.send_call(call).await
        }
    }

    /// Receives a method call reply.
    ///
    /// Convenience wrapper around [`ReadConnection::receive_reply`].
    pub async fn receive_reply<'r, ReplyParams, ReplyError>(
        &'r mut self,
    ) -> Result<RecvResult<reply::Result<ReplyParams, ReplyError>>>
    where
        ReplyParams: Deserialize<'r> + Debug,
        ReplyError: Deserialize<'r> + Debug,
    {
        self.read.receive_reply().await
    }

    /// Call a method and receive a reply.
    ///
    /// This is a convenience method that combines [`Connection::send_call`] and
    /// [`Connection::receive_reply`].
    pub async fn call_method<'r, Method, ReplyParams, ReplyError>(
        &'r mut self,
        call: &Call<Method>,
        #[cfg(feature = "std")] fds: Vec<std::os::fd::OwnedFd>,
    ) -> Result<RecvResult<reply::Result<ReplyParams, ReplyError>>>
    where
        Method: Serialize + Debug,
        ReplyParams: Deserialize<'r> + Debug,
        ReplyError: Deserialize<'r> + Debug,
    {
        #[cfg(feature = "std")]
        self.send_call(call, fds).await?;
        #[cfg(not(feature = "std"))]
        self.send_call(call).await?;

        self.receive_reply().await
    }

    /// Receive a method call over the socket.
    ///
    /// Convenience wrapper around [`ReadConnection::receive_call`].
    pub async fn receive_call<'m, Method>(&'m mut self) -> Result<RecvResult<Call<Method>>>
    where
        Method: Deserialize<'m> + Debug,
    {
        self.read.receive_call().await
    }

    /// Send a reply over the socket.
    ///
    /// Convenience wrapper around [`WriteConnection::send_reply`].
    pub async fn send_reply<ReplyParams>(
        &mut self,
        reply: &Reply<ReplyParams>,
        #[cfg(feature = "std")] fds: Vec<std::os::fd::OwnedFd>,
    ) -> Result<()>
    where
        ReplyParams: Serialize + Debug,
    {
        #[cfg(feature = "std")]
        {
            self.write.send_reply(reply, fds).await
        }
        #[cfg(not(feature = "std"))]
        {
            self.write.send_reply(reply).await
        }
    }

    /// Send an error reply over the socket.
    ///
    /// Convenience wrapper around [`WriteConnection::send_error`].
    pub async fn send_error<ReplyError>(
        &mut self,
        error: &ReplyError,
        #[cfg(feature = "std")] fds: Vec<std::os::fd::OwnedFd>,
    ) -> Result<()>
    where
        ReplyError: Serialize + Debug,
    {
        #[cfg(feature = "std")]
        {
            self.write.send_error(error, fds).await
        }
        #[cfg(not(feature = "std"))]
        {
            self.write.send_error(error).await
        }
    }

    /// Enqueue a call to the server.
    ///
    /// Convenience wrapper around [`WriteConnection::enqueue_call`].
    pub fn enqueue_call<Method>(&mut self, method: &Call<Method>) -> Result<()>
    where
        Method: Serialize + Debug,
    {
        #[cfg(feature = "std")]
        {
            self.write.enqueue_call(method, vec![])
        }
        #[cfg(not(feature = "std"))]
        {
            self.write.enqueue_call(method)
        }
    }

    /// Flush the connection.
    ///
    /// Convenience wrapper around [`WriteConnection::flush`].
    pub async fn flush(&mut self) -> Result<()> {
        self.write.flush().await
    }

    /// Start a chain of method calls.
    ///
    /// This allows batching multiple calls together and sending them in a single write operation.
    ///
    /// # Examples
    ///
    /// ## Basic Usage with Sequential Access
    ///
    /// ```no_run
    /// use zlink_core::{Connection, Call, reply};
    /// use serde::{Serialize, Deserialize};
    /// use serde_prefix_all::prefix_all;
    /// use futures_util::{pin_mut, stream::StreamExt};
    ///
    /// # async fn example() -> zlink_core::Result<()> {
    /// # let mut conn: Connection<zlink_core::connection::socket::impl_for_doc::Socket> = todo!();
    ///
    /// #[prefix_all("org.example.")]
    /// #[derive(Debug, Serialize, Deserialize)]
    /// #[serde(tag = "method", content = "parameters")]
    /// enum Methods {
    ///     GetUser { id: u32 },
    ///     GetProject { id: u32 },
    /// }
    ///
    /// #[derive(Debug, Deserialize)]
    /// struct User { name: String }
    ///
    /// #[derive(Debug, Deserialize)]
    /// struct Project { title: String }
    ///
    /// #[derive(Debug, zlink_core::ReplyError)]
    /// #[zlink(
    ///     interface = "org.example",
    ///     // Not needed in the real code because you'll use `ReplyError` through `zlink` crate.
    ///     crate = "zlink_core",
    /// )]
    /// enum ApiError {
    ///     UserNotFound { code: i32 },
    ///     ProjectNotFound { code: i32 },
    /// }
    ///
    /// let get_user = Call::new(Methods::GetUser { id: 1 });
    /// let get_project = Call::new(Methods::GetProject { id: 2 });
    ///
    /// // Chain calls and send them in a batch
    /// # #[cfg(feature = "std")]
    /// let replies = conn
    ///     .chain_call::<Methods, User, ApiError>(&get_user, vec![])?
    ///     .append(&get_project, vec![])?
    ///     .send().await?;
    /// # #[cfg(not(feature = "std"))]
    /// # let replies = conn
    /// #     .chain_call::<Methods, User, ApiError>(&get_user)?
    /// #     .append(&get_project)?
    /// #     .send().await?;
    /// pin_mut!(replies);
    ///
    /// // Access replies sequentially - types are now fixed by the chain
    /// # #[cfg(feature = "std")]
    /// # {
    /// let (user_reply, _fds) = replies.next().await.unwrap()?;
    /// let (project_reply, _fds) = replies.next().await.unwrap()?;
    ///
    /// match user_reply {
    ///     Ok(user) => println!("User: {}", user.parameters().unwrap().name),
    ///     Err(error) => println!("User error: {:?}", error),
    /// }
    /// # }
    /// # #[cfg(not(feature = "std"))]
    /// # {
    /// # let user_reply = replies.next().await.unwrap()?;
    /// # let _project_reply = replies.next().await.unwrap()?;
    /// #
    /// # match user_reply {
    /// #     Ok(user) => println!("User: {}", user.parameters().unwrap().name),
    /// #     Err(error) => println!("User error: {:?}", error),
    /// # }
    /// # }
    /// # Ok(())
    /// # }
    /// ```
    ///
    /// ## Arbitrary Number of Calls
    ///
    /// ```no_run
    /// # use zlink_core::{Connection, Call, reply};
    /// # use serde::{Serialize, Deserialize};
    /// # use futures_util::{pin_mut, stream::StreamExt};
    /// # use serde_prefix_all::prefix_all;
    /// # async fn example() -> zlink_core::Result<()> {
    /// # let mut conn: Connection<zlink_core::connection::socket::impl_for_doc::Socket> = todo!();
    /// # #[prefix_all("org.example.")]
    /// # #[derive(Debug, Serialize, Deserialize)]
    /// # #[serde(tag = "method", content = "parameters")]
    /// # enum Methods {
    /// #     GetUser { id: u32 },
    /// # }
    /// # #[derive(Debug, Deserialize)]
    /// # struct User { name: String }
    /// # #[derive(Debug, zlink_core::ReplyError)]
    /// #[zlink(
    ///     interface = "org.example",
    ///     // Not needed in the real code because you'll use `ReplyError` through `zlink` crate.
    ///     crate = "zlink_core",
    /// )]
    /// # enum ApiError {
    /// #     UserNotFound { code: i32 },
    /// #     ProjectNotFound { code: i32 },
    /// # }
    /// # let get_user = Call::new(Methods::GetUser { id: 1 });
    ///
    /// // Chain many calls (no upper limit)
    /// # #[cfg(feature = "std")]
    /// let mut chain = conn.chain_call::<Methods, User, ApiError>(&get_user, vec![])?;
    /// # #[cfg(not(feature = "std"))]
    /// # let mut chain = conn.chain_call::<Methods, User, ApiError>(&get_user)?;
    /// # #[cfg(feature = "std")]
    /// for i in 2..100 {
    ///     chain = chain.append(&Call::new(Methods::GetUser { id: i }), vec![])?;
    /// }
    /// # #[cfg(not(feature = "std"))]
    /// # for i in 2..100 {
    /// #     chain = chain.append(&Call::new(Methods::GetUser { id: i }))?;
    /// # }
    ///
    /// let replies = chain.send().await?;
    /// pin_mut!(replies);
    ///
    /// // Process all replies sequentially - types are fixed by the chain
    /// # #[cfg(feature = "std")]
    /// while let Some(result) = replies.next().await {
    ///     let (user_reply, _fds) = result?;
    ///     // Handle each reply...
    ///     match user_reply {
    ///         Ok(user) => println!("User: {}", user.parameters().unwrap().name),
    ///         Err(error) => println!("Error: {:?}", error),
    ///     }
    /// }
    /// # #[cfg(not(feature = "std"))]
    /// # while let Some(result) = replies.next().await {
    /// #     let user_reply = result?;
    /// #     // Handle each reply...
    /// #     match user_reply {
    /// #         Ok(user) => println!("User: {}", user.parameters().unwrap().name),
    /// #         Err(error) => println!("Error: {:?}", error),
    /// #     }
    /// # }
    /// # Ok(())
    /// # }
    /// ```
    ///
    /// # Performance Benefits
    ///
    /// Instead of multiple write operations, the chain sends all calls in a single
    /// write operation, reducing context switching and therefore minimizing latency.
    pub fn chain_call<'c, Method, ReplyParams, ReplyError>(
        &'c mut self,
        call: &Call<Method>,
        #[cfg(feature = "std")] fds: alloc::vec::Vec<std::os::fd::OwnedFd>,
    ) -> Result<Chain<'c, S, ReplyParams, ReplyError>>
    where
        Method: Serialize + Debug,
        ReplyParams: Deserialize<'c> + Debug,
        ReplyError: Deserialize<'c> + Debug,
    {
        Chain::new(
            self,
            call,
            #[cfg(feature = "std")]
            fds,
        )
    }

    /// Get the peer credentials.
    ///
    /// This method caches the credentials on the first call.
    #[cfg(feature = "std")]
    pub async fn peer_credentials(&mut self) -> std::io::Result<&std::sync::Arc<Credentials>>
    where
        S::ReadHalf: socket::FetchPeerCredentials,
    {
        if self.credentials.is_none() {
            let creds = self.read.read_half().fetch_peer_credentials().await?;
            self.credentials = Some(std::sync::Arc::new(creds));
        }

        // Safety: `unwrap` won't panic because we ensure above that it's set correctly if the
        // method doesn't error out.
        Ok(self.credentials.as_ref().unwrap())
    }
}

impl<S> From<S> for Connection<S>
where
    S: Socket,
{
    fn from(socket: S) -> Self {
        Self::new(socket)
    }
}

pub(crate) const BUFFER_SIZE: usize = 256;
const MAX_BUFFER_SIZE: usize = 100 * 1024 * 1024; // Don't allow buffers over 100MB.

static NEXT_ID: AtomicUsize = AtomicUsize::new(0);