From a47277bb66f1c6bbf15e0ca129fce05c42eccb72 Mon Sep 17 00:00:00 2001
From: Dmitry Zolotukhin <zlogic@gmail.com>
Date: Thu, 12 Dec 2024 21:37:05 +0100
Subject: [PATCH] Fixed a bug where not all sockets would be polled.

This effectively caused packet loss.
---
 src/ikev2/mod.rs | 27 ++++++++++++++++-----------
 1 file changed, 16 insertions(+), 11 deletions(-)

diff --git a/src/ikev2/mod.rs b/src/ikev2/mod.rs
index 2646797..41e0b48 100644
--- a/src/ikev2/mod.rs
+++ b/src/ikev2/mod.rs
@@ -272,13 +272,15 @@ impl Server {
 }
 
 struct Sockets {
-    sockets: HashMap<SocketAddr, UdpSocket>,
+    sockets: HashMap<SocketAddr, Arc<UdpSocket>>,
+    socket_list: Vec<(SocketAddr, Arc<UdpSocket>)>,
     nat_port: u16,
 }
 
 impl Sockets {
     async fn new(listen_ips: &[IpAddr], port: u16, nat_port: u16) -> Result<Sockets, IKEv2Error> {
         let mut sockets = HashMap::new();
+        let mut socket_list = vec![];
         for listen_ip in listen_ips {
             for listen_port in [port, nat_port] {
                 let socket = match UdpSocket::bind((*listen_ip, listen_port)).await {
@@ -288,12 +290,18 @@ impl Sockets {
                         return Err(err.into());
                     }
                 };
+                let socket = Arc::new(socket);
                 let listen_addr = socket.local_addr()?;
                 info!("Started server on {}", listen_addr);
-                sockets.insert(listen_addr, socket);
+                sockets.insert(listen_addr, socket.clone());
+                socket_list.push((listen_addr, socket));
             }
         }
-        Ok(Sockets { sockets, nat_port })
+        Ok(Sockets {
+            sockets,
+            socket_list,
+            nat_port,
+        })
     }
 
     fn poll_recv(
@@ -302,14 +310,11 @@ impl Sockets {
         seed: usize,
         buf: &mut ReadBuf<'_>,
     ) -> Poll<Result<UdpDatagramOrigin, IKEv2Error>> {
-        let sockets_len = self.sockets.len();
-        for (listen_addr, socket) in self
-            .sockets
-            .iter()
-            .skip(seed % sockets_len)
-            .cycle()
-            .take(sockets_len)
-        {
+        // Split socket list into two parts, then combine the second half with the first one.
+        let seed = seed % self.socket_list.len();
+        let (left, right) = self.socket_list.split_at(seed % self.socket_list.len());
+        let it = right.into_iter().chain(left.into_iter());
+        for (listen_addr, socket) in it {
             let is_nat_port = listen_addr.port() == self.nat_port;
             buf.clear();
             match socket.poll_recv_from(cx, buf) {