hex/wstunnel
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix(udp): correctly multiplex udp incoming stream

Browse source
This commit is contained in:
Σrebe - Romain GERARD 2023-10-31 22:51:49 +01:00
parent 3f21c4316c
commit f9701a7879
No known key found for this signature in database
GPG key ID: 7A42B4B97E0332F4
1 changed files with 31 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

40
src/udp.rs
View file

@ -142,7 +142,8 @@ impl AsyncRead for UdpStream {
project.pending_notification.as_mut().set(None); project.pending_notification.as_mut().set(None);
} }
let _ = ready!(project.socket.poll_recv(cx, obuf)); let peer = ready!(project.socket.poll_recv_from(cx, obuf))?;
debug_assert_eq!(peer, *project.peer);
let notified: Notified<'static> = unsafe { std::mem::transmute(project.io.has_data_to_read.notified()) }; let notified: Notified<'static> = unsafe { std::mem::transmute(project.io.has_data_to_read.notified()) };
project.pending_notification.as_mut().set(Some(notified)); project.pending_notification.as_mut().set(Some(notified));
project.io.has_read_data.notify_one(); project.io.has_read_data.notify_one();
@ -179,7 +180,14 @@ pub async fn run_server(
.with_context(|| format!("Cannot create UDP server {:?}", bind))?; .with_context(|| format!("Cannot create UDP server {:?}", bind))?;
let udp_server = UdpServer::new(Arc::new(listener), timeout); let udp_server = UdpServer::new(Arc::new(listener), timeout);
let stream = stream::unfold(udp_server, |mut server| async { let stream = stream::unfold((udp_server, None), |(mut server, peer_with_data)| async move {
// New returned peer hasn't read its data yet, await for it.
if let Some(await_peer) = peer_with_data {
if let Some(peer) = server.peers.get(&await_peer) {
peer.has_read_data.notified().await;
}
};
loop { loop {
server.clean_dead_keys(); server.clean_dead_keys();
let peer_addr = match server.listener.peek_sender().await { let peer_addr = match server.listener.peek_sender().await {
@ -192,8 +200,8 @@ pub async fn run_server(
match server.peers.get(&peer_addr) { match server.peers.get(&peer_addr) {
Some(io) => { Some(io) => {
io.has_read_data.notified().await;
io.has_data_to_read.notify_one(); io.has_data_to_read.notify_one();
io.has_read_data.notified().await;
} }
None => { None => {
let (udp_client, io) = UdpStream::new( let (udp_client, io) = UdpStream::new(
@ -206,7 +214,7 @@ pub async fn run_server(
Arc::downgrade(&server.keys_to_delete), Arc::downgrade(&server.keys_to_delete),
); );
server.peers.insert(peer_addr, io); server.peers.insert(peer_addr, io);
return Some((Ok(udp_client), (server))); return Some((Ok(udp_client), (server, Some(peer_addr))));
} }
} }
} }
@ -360,8 +368,7 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn test_multiple_client() { async fn test_multiple_client() {
let server_addr: SocketAddr = "[::1]:1235".parse().unwrap(); let server_addr: SocketAddr = "[::1]:1235".parse().unwrap();
let server = run_server(server_addr, None).await.unwrap(); let mut server = Box::pin(run_server(server_addr, None).await.unwrap());
pin_mut!(server);
// Send some data to the server // Send some data to the server
let client = UdpSocket::bind("[::1]:0").await.unwrap(); let client = UdpSocket::bind("[::1]:0").await.unwrap();
@ -383,21 +390,28 @@ mod tests {
assert!(matches!(ret, Ok(5))); assert!(matches!(ret, Ok(5)));
assert_eq!(&buf[..6], b"aaaaa\0"); assert_eq!(&buf[..6], b"aaaaa\0");
// make the server make progress
let fut2 = timeout(Duration::from_millis(100), server.next()).await; let fut2 = timeout(Duration::from_millis(100), server.next()).await;
assert!(matches!(fut2, Ok(Some(Ok(_))))); assert!(matches!(fut2, Ok(Some(Ok(_)))));
let stream2 = fut2.unwrap().unwrap().unwrap(); let stream2 = fut2.unwrap().unwrap().unwrap();
pin_mut!(stream2); pin_mut!(stream2);
// let the server make progress
tokio::spawn(async move {
loop {
let _ = server.next().await;
}
});
let ret = stream2.read(&mut buf).await; let ret = stream2.read(&mut buf).await;
assert!(matches!(ret, Ok(5))); assert!(matches!(ret, Ok(5)));
assert_eq!(&buf[..6], b"bbbbb\0"); assert_eq!(&buf[..6], b"bbbbb\0");
assert!(client.send_to(b"ccccc".as_ref(), server_addr).await.is_ok()); assert!(client.send_to(b"ccccc".as_ref(), server_addr).await.is_ok());
assert!(client2.send_to(b"ddddd".as_ref(), server_addr).await.is_ok()); assert!(client2.send_to(b"ddddd".as_ref(), server_addr).await.is_ok());
assert!(client2.send_to(b"eeeee".as_ref(), server_addr).await.is_ok());
// Server need to be polled to feed the stream with need data assert!(client.send_to(b"fffff".as_ref(), server_addr).await.is_ok());
let _ = timeout(Duration::from_millis(100), server.next()).await;
let ret = stream.read(&mut buf).await; let ret = stream.read(&mut buf).await;
assert!(matches!(ret, Ok(5))); assert!(matches!(ret, Ok(5)));
@ -406,6 +420,14 @@ mod tests {
let ret = stream2.read(&mut buf).await; let ret = stream2.read(&mut buf).await;
assert!(matches!(ret, Ok(5))); assert!(matches!(ret, Ok(5)));
assert_eq!(&buf[..6], b"ddddd\0"); assert_eq!(&buf[..6], b"ddddd\0");
let ret = stream2.read(&mut buf).await;
assert!(matches!(ret, Ok(5)));
assert_eq!(&buf[..6], b"eeeee\0");
let ret = stream.read(&mut buf).await;
assert!(matches!(ret, Ok(5)));
assert_eq!(&buf[..6], b"fffff\0");
} }
#[tokio::test] #[tokio::test]