1 files changed, 36 insertions, 20 deletions

diff --git a/light-client/src/network.rs b/light-client/src/network.rs
index 997a3344..ed160773 100644
--- a/light-client/src/network.rs
+++ b/light-client/src/network.rs
@@ -17,13 +17,14 @@
 */
 use anyhow::Result;
 use hurrycurry_protocol::{PacketC, PacketS, BINCODE_CONFIG};
-use log::warn;
+use log::{debug, warn};
 use std::{collections::VecDeque, net::TcpStream};
 use tungstenite::{
     client::{uri_mode, IntoClientRequest},
     client_tls_with_config,
     handshake::client::Request,
     stream::{MaybeTlsStream, Mode},
+    util::NonBlockingError,
     Message, WebSocket,
 };
 
@@ -74,34 +75,49 @@ impl Network {
         })
     }
     pub fn poll(&mut self) {
-        self.queue_in.extend(match self.sock.read() {
-            Ok(Message::Text(packet)) => match serde_json::from_str(&packet) {
-                Ok(p) => Some(p),
-                Err(e) => {
-                    warn!("invalid json packet: {e:?}");
-                    None
-                }
-            },
-            Ok(Message::Binary(packet)) => {
-                match bincode::decode_from_slice(&packet, BINCODE_CONFIG) {
-                    Ok((p, _)) => Some(p),
+        loop {
+            self.queue_in.extend(match self.sock.read() {
+                Ok(Message::Text(packet)) => match serde_json::from_str(&packet) {
+                    Ok(packet) => {
+                        debug!("<- {packet:?}");
+                        Some(packet)
+                    }
                     Err(e) => {
-                        warn!("invalid bincode packet: {e:?}");
+                        warn!("invalid json packet: {e:?}");
                         None
                     }
+                },
+                Ok(Message::Binary(packet)) => {
+                    match bincode::decode_from_slice(&packet, BINCODE_CONFIG) {
+                        Ok((packet, _)) => {
+                            debug!("<- {packet:?}");
+                            Some(packet)
+                        }
+                        Err(e) => {
+                            warn!("invalid bincode packet: {e:?}");
+                            None
+                        }
+                    }
                 }
-            }
-            Ok(_) => None,
-            Err(e) => {
-                warn!("{e:?}");
-                None
-            }
-        });
+                Ok(_) => None,
+                Err(e) => {
+                    if let Some(e) = e.into_non_blocking() {
+                        warn!("{e:?}");
+                        None
+                    } else {
+                        break;
+                    }
+                }
+            });
+        }
 
         for packet in self.queue_out.drain(..) {
+            debug!("-> {packet:?}");
             self.sock
                 .write(Message::Text(serde_json::to_string(&packet).unwrap()))
                 .unwrap();
         }
+
+        self.sock.flush().unwrap();
     }
 }