1 #![warn(rust_2018_idioms)]
2 #![cfg(not(target_os = "wasi"))] // Wasi doesn't support UDP
3
4 use tokio::net::UdpSocket;
5 use tokio_stream::StreamExt;
6 use tokio_util::codec::{Decoder, Encoder, LinesCodec};
7 use tokio_util::udp::UdpFramed;
8
9 use bytes::{BufMut, BytesMut};
10 use futures::future::try_join;
11 use futures::future::FutureExt;
12 use futures::sink::SinkExt;
13 use std::io;
14 use std::sync::Arc;
15
16 #[cfg_attr(any(target_os = "macos", target_os = "ios"), allow(unused_assignments))]
17 #[tokio::test]
send_framed_byte_codec() -> std::io::Result<()>18 async fn send_framed_byte_codec() -> std::io::Result<()> {
19 let mut a_soc = UdpSocket::bind("127.0.0.1:0").await?;
20 let mut b_soc = UdpSocket::bind("127.0.0.1:0").await?;
21
22 let a_addr = a_soc.local_addr()?;
23 let b_addr = b_soc.local_addr()?;
24
25 // test sending & receiving bytes
26 {
27 let mut a = UdpFramed::new(a_soc, ByteCodec);
28 let mut b = UdpFramed::new(b_soc, ByteCodec);
29
30 let msg = b"4567";
31
32 let send = a.send((msg, b_addr));
33 let recv = b.next().map(|e| e.unwrap());
34 let (_, received) = try_join(send, recv).await.unwrap();
35
36 let (data, addr) = received;
37 assert_eq!(msg, &*data);
38 assert_eq!(a_addr, addr);
39
40 a_soc = a.into_inner();
41 b_soc = b.into_inner();
42 }
43
44 #[cfg(not(any(target_os = "macos", target_os = "ios")))]
45 // test sending & receiving an empty message
46 {
47 let mut a = UdpFramed::new(a_soc, ByteCodec);
48 let mut b = UdpFramed::new(b_soc, ByteCodec);
49
50 let msg = b"";
51
52 let send = a.send((msg, b_addr));
53 let recv = b.next().map(|e| e.unwrap());
54 let (_, received) = try_join(send, recv).await.unwrap();
55
56 let (data, addr) = received;
57 assert_eq!(msg, &*data);
58 assert_eq!(a_addr, addr);
59 }
60
61 Ok(())
62 }
63
64 pub struct ByteCodec;
65
66 impl Decoder for ByteCodec {
67 type Item = Vec<u8>;
68 type Error = io::Error;
69
decode(&mut self, buf: &mut BytesMut) -> Result<Option<Vec<u8>>, io::Error>70 fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Vec<u8>>, io::Error> {
71 let len = buf.len();
72 Ok(Some(buf.split_to(len).to_vec()))
73 }
74 }
75
76 impl Encoder<&[u8]> for ByteCodec {
77 type Error = io::Error;
78
encode(&mut self, data: &[u8], buf: &mut BytesMut) -> Result<(), io::Error>79 fn encode(&mut self, data: &[u8], buf: &mut BytesMut) -> Result<(), io::Error> {
80 buf.reserve(data.len());
81 buf.put_slice(data);
82 Ok(())
83 }
84 }
85
86 #[tokio::test]
send_framed_lines_codec() -> std::io::Result<()>87 async fn send_framed_lines_codec() -> std::io::Result<()> {
88 let a_soc = UdpSocket::bind("127.0.0.1:0").await?;
89 let b_soc = UdpSocket::bind("127.0.0.1:0").await?;
90
91 let a_addr = a_soc.local_addr()?;
92 let b_addr = b_soc.local_addr()?;
93
94 let mut a = UdpFramed::new(a_soc, ByteCodec);
95 let mut b = UdpFramed::new(b_soc, LinesCodec::new());
96
97 let msg = b"1\r\n2\r\n3\r\n".to_vec();
98 a.send((&msg, b_addr)).await?;
99
100 assert_eq!(b.next().await.unwrap().unwrap(), ("1".to_string(), a_addr));
101 assert_eq!(b.next().await.unwrap().unwrap(), ("2".to_string(), a_addr));
102 assert_eq!(b.next().await.unwrap().unwrap(), ("3".to_string(), a_addr));
103
104 Ok(())
105 }
106
107 #[tokio::test]
framed_half() -> std::io::Result<()>108 async fn framed_half() -> std::io::Result<()> {
109 let a_soc = Arc::new(UdpSocket::bind("127.0.0.1:0").await?);
110 let b_soc = a_soc.clone();
111
112 let a_addr = a_soc.local_addr()?;
113 let b_addr = b_soc.local_addr()?;
114
115 let mut a = UdpFramed::new(a_soc, ByteCodec);
116 let mut b = UdpFramed::new(b_soc, LinesCodec::new());
117
118 let msg = b"1\r\n2\r\n3\r\n".to_vec();
119 a.send((&msg, b_addr)).await?;
120
121 let msg = b"4\r\n5\r\n6\r\n".to_vec();
122 a.send((&msg, b_addr)).await?;
123
124 assert_eq!(b.next().await.unwrap().unwrap(), ("1".to_string(), a_addr));
125 assert_eq!(b.next().await.unwrap().unwrap(), ("2".to_string(), a_addr));
126 assert_eq!(b.next().await.unwrap().unwrap(), ("3".to_string(), a_addr));
127
128 assert_eq!(b.next().await.unwrap().unwrap(), ("4".to_string(), a_addr));
129 assert_eq!(b.next().await.unwrap().unwrap(), ("5".to_string(), a_addr));
130 assert_eq!(b.next().await.unwrap().unwrap(), ("6".to_string(), a_addr));
131
132 Ok(())
133 }
134