mavlink_core/connection/udp/
async.rs

1//! Async UDP MAVLink connection
2
3use core::task::Poll;
4use std::io;
5use std::sync::{Arc, Mutex as StdMutex};
6
7use futures::lock::{Mutex, MutexGuard};
8use tokio::{
9    io::{AsyncRead, AsyncWrite, ReadBuf},
10    net::UdpSocket,
11};
12
13use crate::connection::r#async::{
14    AsyncConnectable, AsyncDialectConnection, AsyncTransport, Transport,
15};
16use crate::connection::get_socket_addr;
17use crate::connection::udp::config::{UdpConfig, UdpMode};
18use crate::{AsyncMavlinkReader, Dialect};
19
20pub(crate) struct UdpRead {
21    socket: Arc<UdpSocket>,
22    reply_destination: Option<Arc<StdMutex<Option<std::net::SocketAddr>>>>,
23}
24
25impl AsyncRead for UdpRead {
26    fn poll_read(
27        self: core::pin::Pin<&mut Self>,
28        cx: &mut core::task::Context<'_>,
29        buf: &mut ReadBuf<'_>,
30    ) -> Poll<io::Result<()>> {
31        match self.socket.poll_recv_from(cx, buf) {
32            Poll::Ready(Ok(address)) => {
33                if let Some(reply_destination) = &self.reply_destination {
34                    *reply_destination.lock().unwrap() = Some(address);
35                }
36                Poll::Ready(Ok(()))
37            }
38            Poll::Ready(Err(err)) => Poll::Ready(Err(err)),
39            Poll::Pending => Poll::Pending,
40        }
41    }
42}
43
44pub(crate) struct UdpWrite {
45    socket: Arc<UdpSocket>,
46    destination: Arc<StdMutex<Option<std::net::SocketAddr>>>,
47    connected: bool,
48}
49
50impl AsyncWrite for UdpWrite {
51    fn poll_write(
52        self: core::pin::Pin<&mut Self>,
53        cx: &mut core::task::Context<'_>,
54        buf: &[u8],
55    ) -> Poll<io::Result<usize>> {
56        let this = self.get_mut();
57        let addr = this
58            .destination
59            .lock()
60            .unwrap()
61            .expect("`dest` is checked before write");
62
63        let result = if this.connected {
64            this.socket.poll_send(cx, buf)
65        } else {
66            this.socket.poll_send_to(cx, buf, addr)
67        };
68        match result {
69            Poll::Ready(Ok(written)) if written == buf.len() => Poll::Ready(Ok(written)),
70            Poll::Ready(Ok(_)) => Poll::Ready(Err(io::Error::new(
71                io::ErrorKind::WriteZero,
72                "failed to send complete UDP datagram",
73            ))),
74            Poll::Ready(Err(err)) => Poll::Ready(Err(err)),
75            Poll::Pending => Poll::Pending,
76        }
77    }
78
79    fn poll_flush(
80        self: core::pin::Pin<&mut Self>,
81        _cx: &mut core::task::Context<'_>,
82    ) -> Poll<io::Result<()>> {
83        Poll::Ready(Ok(()))
84    }
85
86    fn poll_shutdown(
87        self: core::pin::Pin<&mut Self>,
88        _cx: &mut core::task::Context<'_>,
89    ) -> Poll<io::Result<()>> {
90        Poll::Ready(Ok(()))
91    }
92}
93
94pub struct AsyncUdpConnection {
95    reader: Mutex<AsyncMavlinkReader<UdpRead>>,
96    writer: Mutex<UdpWrite>,
97}
98
99impl AsyncUdpConnection {
100    fn new(
101        socket: UdpSocket,
102        server: bool,
103        dest: Option<std::net::SocketAddr>,
104        connected: bool,
105    ) -> io::Result<Self> {
106        let socket = Arc::new(socket);
107        let destination = Arc::new(StdMutex::new(dest));
108        Ok(Self {
109            reader: Mutex::new(AsyncMavlinkReader::new(UdpRead {
110                socket: socket.clone(),
111                reply_destination: server.then(|| destination.clone()),
112            })),
113            writer: Mutex::new(UdpWrite {
114                socket,
115                destination,
116                connected,
117            }),
118        })
119    }
120}
121
122impl AsyncTransport for AsyncUdpConnection {
123    type Reader = UdpRead;
124    type Writer = UdpWrite;
125
126    fn reader(&self) -> &Mutex<AsyncMavlinkReader<Self::Reader>> {
127        &self.reader
128    }
129
130    async fn lock_writer(&self) -> Option<MutexGuard<'_, Self::Writer>> {
131        let writer = self.writer.lock().await;
132        let has_destination = writer.destination.lock().unwrap().is_some();
133        has_destination.then_some(writer)
134    }
135
136    fn retry_receive(&self, _: &crate::error::MessageReadError) -> bool {
137        true
138    }
139}
140
141impl UdpConfig {
142    pub(crate) async fn open_async(&self) -> io::Result<AsyncUdpConnection> {
143        let (addr, server, dest): (&str, _, _) = match self.mode {
144            UdpMode::Udpin => (&self.address, true, None),
145            _ => ("0.0.0.0:0", false, Some(get_socket_addr(&self.address)?)),
146        };
147        let (socket, connected) = match self.take_socket()? {
148            Some(socket) => {
149                socket.set_nonblocking(true)?;
150                (UdpSocket::from_std(socket)?, !server)
151            }
152            None => (UdpSocket::bind(addr).await?, false),
153        };
154        if matches!(self.mode, UdpMode::UdpBroadcast) {
155            socket.set_broadcast(true)?;
156        }
157        AsyncUdpConnection::new(socket, server, dest, connected)
158    }
159}
160
161impl AsyncConnectable for UdpConfig {
162    async fn connect_async_with_dialect<D>(
163        &self,
164        dialect: D,
165    ) -> io::Result<AsyncDialectConnection<D>>
166    where
167        D: Dialect + Send + Sync,
168        D::Message: Send + Sync,
169    {
170        Ok(AsyncDialectConnection::new(
171            Transport::Udp(self.open_async().await?),
172            dialect,
173        ))
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180    use tokio::io::AsyncReadExt;
181
182    #[tokio::test]
183    async fn test_datagram_reading() {
184        let receiver_socket = Arc::new(UdpSocket::bind("127.0.0.1:5001").await.unwrap());
185        let mut udp_reader = UdpRead {
186            socket: receiver_socket.clone(),
187            reply_destination: None,
188        };
189        let sender_socket = UdpSocket::bind("0.0.0.0:0").await.unwrap();
190        sender_socket.connect("127.0.0.1:5001").await.unwrap();
191
192        let datagram: Vec<u8> = (0..50).collect::<Vec<_>>();
193
194        let mut n_sent = sender_socket.send(&datagram).await.unwrap();
195        assert_eq!(n_sent, datagram.len());
196        n_sent = sender_socket.send(&datagram).await.unwrap();
197        assert_eq!(n_sent, datagram.len());
198
199        let mut buf = [0u8; 50];
200
201        let mut n_read = udp_reader.read(&mut buf).await.unwrap();
202        assert_eq!(n_read, 50);
203        assert_eq!(&buf[0..n_read], datagram.as_slice());
204
205        n_read = udp_reader.read(&mut buf).await.unwrap();
206        assert_eq!(n_read, 50);
207        assert_eq!(&buf[0..n_read], datagram.as_slice());
208    }
209}
mavlink_core/connection/udp/async.rs

mavlink_core/connection/udp/
async.rs
