mavlink_core/connection/
udp.rs

1//! UDP MAVLink connection
2
3use std::collections::VecDeque;
4
5use crate::connectable::{UdpConnectable, UdpMode};
6use crate::connection::MavConnection;
7use crate::peek_reader::PeekReader;
8use crate::{MavHeader, MavlinkVersion, Message};
9use core::ops::DerefMut;
10use std::io::{self, Read};
11use std::net::{SocketAddr, UdpSocket};
12use std::sync::Mutex;
13
14use super::{get_socket_addr, Connectable};
15
16#[cfg(not(feature = "signing"))]
17use crate::{read_versioned_msg, write_versioned_msg};
18
19#[cfg(feature = "signing")]
20use crate::{read_versioned_msg_signed, write_versioned_msg_signed, SigningConfig, SigningData};
21
22struct UdpRead {
23    socket: UdpSocket,
24    buffer: VecDeque<u8>,
25    last_recv_address: Option<SocketAddr>,
26}
27
28const MTU_SIZE: usize = 1500;
29impl Read for UdpRead {
30    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
31        if !self.buffer.is_empty() {
32            self.buffer.read(buf)
33        } else {
34            let mut read_buffer = [0u8; MTU_SIZE];
35            let (n_buffer, address) = self.socket.recv_from(&mut read_buffer)?;
36            let n = (&read_buffer[0..n_buffer]).read(buf)?;
37            self.buffer.extend(&read_buffer[n..n_buffer]);
38
39            self.last_recv_address = Some(address);
40            Ok(n)
41        }
42    }
43}
44
45struct UdpWrite {
46    socket: UdpSocket,
47    dest: Option<SocketAddr>,
48    sequence: u8,
49}
50
51pub struct UdpConnection {
52    reader: Mutex<PeekReader<UdpRead>>,
53    writer: Mutex<UdpWrite>,
54    protocol_version: MavlinkVersion,
55    server: bool,
56    #[cfg(feature = "signing")]
57    signing_data: Option<SigningData>,
58}
59
60impl UdpConnection {
61    fn new(socket: UdpSocket, server: bool, dest: Option<SocketAddr>) -> io::Result<Self> {
62        Ok(Self {
63            server,
64            reader: Mutex::new(PeekReader::new(UdpRead {
65                socket: socket.try_clone()?,
66                buffer: VecDeque::new(),
67                last_recv_address: None,
68            })),
69            writer: Mutex::new(UdpWrite {
70                socket,
71                dest,
72                sequence: 0,
73            }),
74            protocol_version: MavlinkVersion::V2,
75            #[cfg(feature = "signing")]
76            signing_data: None,
77        })
78    }
79}
80
81impl<M: Message> MavConnection<M> for UdpConnection {
82    fn recv(&self) -> Result<(MavHeader, M), crate::error::MessageReadError> {
83        let mut reader = self.reader.lock().unwrap();
84
85        loop {
86            #[cfg(not(feature = "signing"))]
87            let result = read_versioned_msg(reader.deref_mut(), self.protocol_version);
88            #[cfg(feature = "signing")]
89            let result = read_versioned_msg_signed(
90                reader.deref_mut(),
91                self.protocol_version,
92                self.signing_data.as_ref(),
93            );
94            if self.server {
95                if let addr @ Some(_) = reader.reader_ref().last_recv_address {
96                    self.writer.lock().unwrap().dest = addr;
97                }
98            }
99            if let ok @ Ok(..) = result {
100                return ok;
101            }
102        }
103    }
104
105    fn send(&self, header: &MavHeader, data: &M) -> Result<usize, crate::error::MessageWriteError> {
106        let mut guard = self.writer.lock().unwrap();
107        let state = &mut *guard;
108
109        let header = MavHeader {
110            sequence: state.sequence,
111            system_id: header.system_id,
112            component_id: header.component_id,
113        };
114
115        state.sequence = state.sequence.wrapping_add(1);
116
117        let len = if let Some(addr) = state.dest {
118            let mut buf = Vec::new();
119            #[cfg(not(feature = "signing"))]
120            write_versioned_msg(&mut buf, self.protocol_version, header, data)?;
121            #[cfg(feature = "signing")]
122            write_versioned_msg_signed(
123                &mut buf,
124                self.protocol_version,
125                header,
126                data,
127                self.signing_data.as_ref(),
128            )?;
129            state.socket.send_to(&buf, addr)?
130        } else {
131            0
132        };
133
134        Ok(len)
135    }
136
137    fn set_protocol_version(&mut self, version: MavlinkVersion) {
138        self.protocol_version = version;
139    }
140
141    fn protocol_version(&self) -> MavlinkVersion {
142        self.protocol_version
143    }
144
145    #[cfg(feature = "signing")]
146    fn setup_signing(&mut self, signing_data: Option<SigningConfig>) {
147        self.signing_data = signing_data.map(SigningData::from_config)
148    }
149}
150
151impl Connectable for UdpConnectable {
152    fn connect<M: Message>(&self) -> io::Result<Box<dyn MavConnection<M> + Sync + Send>> {
153        let (addr, server, dest): (&str, _, _) = match self.mode {
154            UdpMode::Udpin => (&self.address, true, None),
155            _ => ("0.0.0.0:0", false, Some(get_socket_addr(&self.address)?)),
156        };
157        let socket = UdpSocket::bind(addr)?;
158        if matches!(self.mode, UdpMode::Udpcast) {
159            socket.set_broadcast(true)?;
160        }
161        Ok(Box::new(UdpConnection::new(socket, server, dest)?))
162    }
163}
164
165#[cfg(test)]
166mod tests {
167    use super::*;
168
169    #[test]
170    fn test_datagram_buffering() {
171        let receiver_socket = UdpSocket::bind("127.0.0.1:5000").unwrap();
172        let mut udp_reader = UdpRead {
173            socket: receiver_socket.try_clone().unwrap(),
174            buffer: VecDeque::new(),
175            last_recv_address: None,
176        };
177        let sender_socket = UdpSocket::bind("0.0.0.0:0").unwrap();
178        sender_socket.connect("127.0.0.1:5000").unwrap();
179
180        let datagram: Vec<u8> = (0..50).collect::<Vec<_>>();
181
182        let mut n_sent = sender_socket.send(&datagram).unwrap();
183        assert_eq!(n_sent, datagram.len());
184        n_sent = sender_socket.send(&datagram).unwrap();
185        assert_eq!(n_sent, datagram.len());
186
187        let mut buf = [0u8; 30];
188
189        let mut n_read = udp_reader.read(&mut buf).unwrap();
190        assert_eq!(n_read, 30);
191        assert_eq!(&buf[0..n_read], (0..30).collect::<Vec<_>>().as_slice());
192
193        n_read = udp_reader.read(&mut buf).unwrap();
194        assert_eq!(n_read, 20);
195        assert_eq!(&buf[0..n_read], (30..50).collect::<Vec<_>>().as_slice());
196
197        n_read = udp_reader.read(&mut buf).unwrap();
198        assert_eq!(n_read, 30);
199        assert_eq!(&buf[0..n_read], (0..30).collect::<Vec<_>>().as_slice());
200
201        n_read = udp_reader.read(&mut buf).unwrap();
202        assert_eq!(n_read, 20);
203        assert_eq!(&buf[0..n_read], (30..50).collect::<Vec<_>>().as_slice());
204    }
205}