Skip to main content

statime_netptp/socket/
recv.rs

1use std::{
2    collections::HashMap,
3    future::poll_fn,
4    io::Result,
5    sync::Arc,
6    task::{Context, Poll, Waker},
7};
8
9use statime_base::UniversalTimestamp;
10use timestamped_socket::interface::InterfaceName;
11
12use crate::{
13    BoundInterface, ConnectedSocket, MAX_PACKET_SIZE, NetworkManagerData, OpenSocket,
14    PtpAddressFamily, SocketData, TimestampSource,
15};
16
17/// Result from a receive operation.
18///
19/// Asside from the message, also provides context such as who sent it, to
20/// which local address it was sent, and when it arrived.
21#[derive(Debug, Clone, PartialEq, Eq, Hash)]
22pub struct RecvResult<A> {
23    /// The actual received message
24    pub bytes_read: Arc<[u8]>,
25    /// The address of the remote which sent the message.
26    pub remote_addr: A,
27    /// The local address at which the message was targeted.
28    pub local_addr: A,
29    /// The moment at which the message was received, if available.
30    pub timestamp: Option<UniversalTimestamp>,
31}
32
33impl<A: PtpAddressFamily> NetworkManagerData<A> {
34    fn recv_for_socket(&self, socket_id: usize, cx: &mut Context<'_>) -> Poll<Result<()>> {
35        self.read_wakers.add_waker(cx.waker().clone());
36        let mut buf = [0u8; MAX_PACKET_SIZE];
37        // The mutex can only be poisoned from an earlier panic. It is ok for
38        // us to propagate that to all the threads.
39        let sockets = self.sockets.read().unwrap();
40        // The mutex can only be poisoned from an earlier panic. It is ok for
41        // us to propagate that to all the threads.
42        let interfaces = self.interfaces.read().unwrap();
43        // The socket state will always be present.
44        if let Some(interface_name) = sockets[&socket_id].interface_filter {
45            if let Poll::Ready(recv_result) = interfaces[&Some(interface_name)].interface.poll_recv(
46                &mut buf,
47                &mut Context::from_waker(&Waker::from(self.read_wakers.clone())),
48            ) {
49                Poll::Ready(handle_recv_result(
50                    &buf,
51                    &sockets,
52                    recv_result,
53                    Some(interface_name),
54                ))
55            } else {
56                Poll::Pending
57            }
58        } else {
59            for (&interface_name, interface) in interfaces.iter() {
60                if let Poll::Ready(recv_result) = interface.interface.poll_recv(
61                    &mut buf,
62                    &mut Context::from_waker(&Waker::from(self.read_wakers.clone())),
63                ) {
64                    return Poll::Ready(handle_recv_result(
65                        &buf,
66                        &sockets,
67                        recv_result,
68                        interface_name,
69                    ));
70                }
71            }
72
73            Poll::Pending
74        }
75    }
76}
77
78fn handle_recv_result<A: PtpAddressFamily>(
79    buf: &[u8],
80    sockets: &HashMap<usize, SocketData<A>>,
81    recv_result: Result<timestamped_socket::socket::RecvResult<A>>,
82    interface_name: Option<InterfaceName>,
83) -> Result<()> {
84    match recv_result {
85        Ok(recv_result) => {
86            let bytes_read: Arc<[u8]> = buf[..recv_result.bytes_read].into();
87
88            for socket in sockets.values() {
89                if socket
90                    .interface_filter
91                    .is_none_or(|socket_interface| Some(socket_interface) == interface_name)
92                    && socket
93                        .local_filter
94                        .is_none_or(|addr| addr == recv_result.local_addr)
95                    && socket
96                        .remote_filter
97                        .is_none_or(|addr| addr == recv_result.remote_addr)
98                {
99                    // UDP is non-guaranteed in delivery. This is just another
100                    // point where the message can be dropped if the buffers
101                    // are full.
102                    let _ = socket.message_channel.try_send(RecvResult {
103                        bytes_read: bytes_read.clone(),
104                        remote_addr: recv_result.remote_addr,
105                        local_addr: recv_result.local_addr,
106                        timestamp: match socket.timestamp_source {
107                            TimestampSource::System => recv_result
108                                .timestamp_data
109                                .software
110                                .map(UniversalTimestamp::from),
111                            TimestampSource::Hardware => recv_result
112                                .timestamp_data
113                                .hardware
114                                .map(UniversalTimestamp::from),
115                        },
116                    });
117                }
118            }
119
120            Ok(())
121        }
122        Err(e) => Err(e),
123    }
124}
125
126impl<A: PtpAddressFamily> OpenSocket<A> {
127    /// Receive an incoming message.
128    ///
129    /// # Errors
130    ///
131    /// Returns any IO error that occured trying to receive a packet.
132    pub async fn recv(&mut self) -> Result<RecvResult<A>> {
133        poll_fn(|cx| {
134            loop {
135                if let Poll::Ready(Some(recv_packet)) = self.packet_receiver.poll_recv(cx) {
136                    return Poll::Ready(Ok(recv_packet));
137                }
138
139                match self.state.recv_for_socket(self.socket_id, cx) {
140                    Poll::Ready(Ok(())) => {}
141                    Poll::Ready(Err(e)) => break Poll::Ready(Err(e)),
142                    Poll::Pending => break Poll::Pending,
143                }
144            }
145        })
146        .await
147    }
148}
149
150impl<A: PtpAddressFamily> ConnectedSocket<A> {
151    /// Receive an incoming message.
152    ///
153    /// # Errors
154    ///
155    /// Returns any IO error that occured trying to receive a packet.
156    pub async fn recv(&mut self) -> Result<RecvResult<A>> {
157        poll_fn(|cx| {
158            loop {
159                if let Poll::Ready(Some(recv_packet)) = self.packet_receiver.poll_recv(cx) {
160                    return Poll::Ready(Ok(recv_packet));
161                }
162
163                match self.state.recv_for_socket(self.socket_id, cx) {
164                    Poll::Ready(Ok(())) => {}
165                    Poll::Ready(Err(e)) => break Poll::Ready(Err(e)),
166                    Poll::Pending => break Poll::Pending,
167                }
168            }
169        })
170        .await
171    }
172}