Skip to main content

zlink_core/connection/
read_connection.rs

1//! Contains connection related API.
2
3use core::{fmt::Debug, str::from_utf8_unchecked};
4
5#[cfg(all(feature = "std", target_os = "linux"))]
6use crate::connection::PassedCredentials;
7use crate::{Result, varlink_service};
8
9use super::{
10    BUFFER_SIZE, Call, MAX_BUFFER_SIZE,
11    reply::{self, Reply},
12    socket::ReadHalf,
13};
14#[cfg(feature = "std")]
15use alloc::collections::VecDeque;
16use alloc::vec::Vec;
17use serde::Deserialize;
18use serde_json::Deserializer;
19
20#[cfg(feature = "std")]
21use std::os::fd::OwnedFd;
22
23// Type alias for receive methods - std returns FDs, no_std doesn't
24#[cfg(feature = "std")]
25type RecvResult<T> = (T, Vec<OwnedFd>);
26#[cfg(not(feature = "std"))]
27type RecvResult<T> = T;
28
29/// A connection that can only be used for reading.
30///
31/// # Cancel safety
32///
33/// All async methods of this type are cancel safe unless explicitly stated otherwise in its
34/// documentation.
35#[derive(Debug)]
36pub struct ReadConnection<Read: ReadHalf> {
37    socket: Read,
38    read_pos: usize,
39    msg_pos: usize,
40    buffer: Vec<u8>,
41    id: usize,
42    #[cfg(feature = "std")]
43    pending_fds: VecDeque<Vec<OwnedFd>>,
44    // Number of `recvmsg` calls that returned FDs. Used by `Connection` to drain
45    // `WriteConnection::held_fds` on macOS. See that field's comment for details.
46    #[cfg(all(feature = "std", target_os = "macos"))]
47    pub(super) fd_recvs: usize,
48    #[cfg(all(feature = "std", target_os = "linux"))]
49    received_credentials: Option<std::sync::Arc<PassedCredentials>>,
50}
51
52impl<Read: ReadHalf> ReadConnection<Read> {
53    /// Create a new connection.
54    pub(super) fn new(socket: Read, id: usize) -> Self {
55        Self {
56            socket,
57            read_pos: 0,
58            msg_pos: 0,
59            id,
60            buffer: alloc::vec![0; BUFFER_SIZE],
61            #[cfg(feature = "std")]
62            pending_fds: VecDeque::new(),
63            #[cfg(all(feature = "std", target_os = "macos"))]
64            fd_recvs: 0,
65            #[cfg(all(feature = "std", target_os = "linux"))]
66            received_credentials: None,
67        }
68    }
69
70    /// The unique identifier of the connection.
71    #[inline]
72    pub fn id(&self) -> usize {
73        self.id
74    }
75
76    /// Receives a method call reply.
77    ///
78    /// The generic parameters needs some explanation:
79    ///
80    /// * `ReplyParams` is the type of the successful reply. This should be a type that can
81    ///   deserialize itself from the `parameters` field of the reply.
82    /// * `ReplyError` is the type of the error reply. This should be a type that can deserialize
83    ///   itself from the whole reply object itself and must fail when there is no `error` field in
84    ///   the object. This can be easily achieved using the `zlink::ReplyError` derive:
85    ///
86    /// ```rust
87    /// use zlink_core::ReplyError;
88    ///
89    /// #[derive(Debug, ReplyError)]
90    /// #[zlink(
91    ///     interface = "org.example.ftl",
92    ///     // Not needed in the real code because you'll use `ReplyError` through `zlink` crate.
93    ///     crate = "zlink_core",
94    /// )]
95    /// enum MyError {
96    ///     Alpha { param1: u32, param2: String },
97    ///     Bravo,
98    ///     Charlie { param1: String },
99    /// }
100    /// ```
101    ///
102    /// Returns the reply and any file descriptors received (std only).
103    pub async fn receive_reply<'r, ReplyParams, ReplyError>(
104        &'r mut self,
105    ) -> Result<RecvResult<reply::Result<ReplyParams, ReplyError>>>
106    where
107        ReplyParams: Deserialize<'r> + Debug,
108        ReplyError: Deserialize<'r> + Debug,
109    {
110        #[derive(Debug, Deserialize)]
111        #[serde(untagged)]
112        enum ReplyMsg<'m, ReplyParams, ReplyError> {
113            #[serde(borrow)]
114            Varlink(varlink_service::Error<'m>),
115            Error(ReplyError),
116            Reply(Reply<ReplyParams>),
117        }
118
119        let recv_result = self
120            .read_message::<ReplyMsg<'_, ReplyParams, ReplyError>>()
121            .await?;
122
123        #[cfg(feature = "std")]
124        let (msg, fds) = recv_result;
125        #[cfg(not(feature = "std"))]
126        let msg = recv_result;
127
128        let result = match msg {
129            // Varlink service interface error need to be returned as the top-level error.
130            ReplyMsg::Varlink(e) => Err(crate::Error::VarlinkService(e.into())),
131            ReplyMsg::Error(e) => Ok(Err(e)),
132            ReplyMsg::Reply(reply) => Ok(Ok(reply)),
133        };
134
135        #[cfg(feature = "std")]
136        return result.map(|r| (r, fds));
137        #[cfg(not(feature = "std"))]
138        return result;
139    }
140
141    /// Receive a method call over the socket.
142    ///
143    /// The generic `Method` is the type of the method name and its input parameters. This should be
144    /// a type that can deserialize itself from a complete method call message, i-e an object
145    /// containing `method` and `parameter` fields. This can be easily achieved using the
146    /// `serde::Deserialize` derive (See the code snippet in [`super::WriteConnection::send_call`]
147    /// documentation for an example).
148    ///
149    /// Returns the call and any file descriptors received (std only).
150    pub async fn receive_call<'m, Method>(&'m mut self) -> Result<RecvResult<Call<Method>>>
151    where
152        Method: Deserialize<'m> + Debug,
153    {
154        self.read_message::<Call<Method>>().await
155    }
156
157    // Reads at least one full message from the socket and return a single message bytes.
158    async fn read_message<'m, M>(&'m mut self) -> Result<RecvResult<M>>
159    where
160        M: Deserialize<'m> + Debug,
161    {
162        self.read_from_socket().await?;
163
164        let mut stream = Deserializer::from_slice(&self.buffer[self.msg_pos..]).into_iter::<M>();
165        let msg = stream.next();
166        let null_index = self.msg_pos + stream.byte_offset();
167        let buffer = &self.buffer[self.msg_pos..null_index];
168        if self.buffer[null_index + 1] == b'\0' {
169            // This means we're reading the last message and can now reset the indices.
170            self.read_pos = 0;
171            self.msg_pos = 0;
172        } else {
173            self.msg_pos = null_index + 1;
174        }
175
176        match msg {
177            Some(Ok(msg)) => {
178                // SAFETY: Since the parsing from JSON already succeeded, we can be sure that the
179                // buffer contains a valid UTF-8 string.
180                trace!("connection {}: received a message: {}", self.id, unsafe {
181                    from_utf8_unchecked(buffer)
182                });
183
184                #[cfg(feature = "std")]
185                {
186                    let fds = self.pending_fds.pop_front().unwrap_or_default();
187                    Ok((msg, fds))
188                }
189                #[cfg(not(feature = "std"))]
190                Ok(msg)
191            }
192            Some(Err(e)) => Err(e.into()),
193            None => Err(crate::Error::UnexpectedEof),
194        }
195    }
196
197    // Reads at least one full message from the socket.
198    async fn read_from_socket(&mut self) -> Result<()> {
199        if self.msg_pos > 0 {
200            // This means we already have at least one message in the buffer so no need to read.
201            return Ok(());
202        }
203
204        loop {
205            let result = self.socket.read(&mut self.buffer[self.read_pos..]).await?;
206            #[cfg(all(feature = "std", target_os = "linux"))]
207            let mut result = result;
208
209            if result.bytes_read() == 0 {
210                return Err(crate::Error::UnexpectedEof);
211            }
212            self.read_pos += result.bytes_read();
213            #[cfg(all(feature = "std", target_os = "linux"))]
214            if let Some(creds) = result.take_credentials() {
215                self.received_credentials = Some(std::sync::Arc::new(creds));
216            }
217            #[cfg(feature = "std")]
218            if !result.fds().is_empty() {
219                self.pending_fds.push_back(result.into_fds());
220                // Track receipt so `Connection` can drain `WriteConnection::held_fds`.
221                #[cfg(target_os = "macos")]
222                {
223                    self.fd_recvs += 1;
224                }
225            }
226
227            if self.read_pos == self.buffer.len() {
228                if self.read_pos >= MAX_BUFFER_SIZE {
229                    return Err(crate::Error::BufferOverflow);
230                }
231
232                self.buffer.extend(core::iter::repeat_n(0, BUFFER_SIZE));
233            }
234
235            // This marks end of all messages. After this loop is finished, we'll have 2 consecutive
236            // null bytes at the end. This is then used by the callers to determine that they've
237            // read all messages and can now reset the `read_pos`.
238            self.buffer[self.read_pos] = b'\0';
239
240            if self.buffer[self.read_pos - 1] == b'\0' {
241                // One or more full messages were read.
242                break;
243            }
244        }
245
246        Ok(())
247    }
248
249    /// The underlying read half of the socket.
250    pub fn read_half(&self) -> &Read {
251        &self.socket
252    }
253
254    /// The credentials passed through over the socket with the latest message (if any).
255    #[cfg(all(feature = "std", target_os = "linux"))]
256    pub fn received_credentials(&self) -> Option<&std::sync::Arc<PassedCredentials>>
257    where
258        Read: super::socket::UnixSocket,
259    {
260        self.received_credentials.as_ref()
261    }
262}