From 58a7f03feac0e8ca23fac6529734f9ded14e724b Mon Sep 17 00:00:00 2001
From: Eloi DEMOLIS <eloi.demolis@clever-cloud.com>
Date: Thu, 23 Nov 2023 15:51:49 +0100
Subject: [PATCH] Fix: WouldBlock in SocketHandler::socket_write breaks
 properly

Signed-off-by: Eloi DEMOLIS <eloi.demolis@clever-cloud.com>
---
 lib/src/socket.rs | 43 +++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 41 insertions(+), 2 deletions(-)

diff --git a/lib/src/socket.rs b/lib/src/socket.rs
index 8e50a2225..a8d4a1ca5 100644
--- a/lib/src/socket.rs
+++ b/lib/src/socket.rs
@@ -44,6 +44,8 @@ pub enum TransportProtocol {
     Tls1_3,
 }
 
+const MAX_LOOP_ITERATION: usize = 100000;
+
 pub trait SocketHandler {
     fn socket_read(&mut self, buf: &mut [u8]) -> (usize, SocketResult);
     fn socket_write(&mut self, buf: &[u8]) -> (usize, SocketResult);
@@ -60,7 +62,13 @@ pub trait SocketHandler {
 impl SocketHandler for TcpStream {
     fn socket_read(&mut self, buf: &mut [u8]) -> (usize, SocketResult) {
         let mut size = 0usize;
+        let mut counter = 0;
         loop {
+            counter += 1;
+            if counter > MAX_LOOP_ITERATION {
+                error!("MAX_LOOP_ITERATION reached in TcpStream::socket_read");
+                incr!("socket.read.infinite_loop.error");
+            }
             if size == buf.len() {
                 return (size, SocketResult::Continue);
             }
@@ -83,7 +91,13 @@ impl SocketHandler for TcpStream {
 
     fn socket_write(&mut self, buf: &[u8]) -> (usize, SocketResult) {
         let mut size = 0usize;
+        let mut counter = 0;
         loop {
+            counter += 1;
+            if counter > MAX_LOOP_ITERATION {
+                error!("MAX_LOOP_ITERATION reached in TcpStream::socket_write");
+                incr!("socket.write.infinite_loop.error");
+            }
             if size == buf.len() {
                 return (size, SocketResult::Continue);
             }
@@ -165,7 +179,14 @@ impl SocketHandler for FrontRustls {
         let mut is_error = false;
         let mut is_closed = false;
 
+        let mut counter = 0;
         loop {
+            counter += 1;
+            if counter > MAX_LOOP_ITERATION {
+                error!("MAX_LOOP_ITERATION reached in FrontRustls::socket_read");
+                incr!("socket.read.infinite_loop.error");
+            }
+
             if size == buf.len() {
                 break;
             }
@@ -251,7 +272,13 @@ impl SocketHandler for FrontRustls {
         let mut is_error = false;
         let mut is_closed = false;
 
+        let mut counter = 0;
         loop {
+            counter += 1;
+            if counter > MAX_LOOP_ITERATION {
+                error!("MAX_LOOP_ITERATION reached in FrontRustls::socket_write");
+                incr!("socket.write.infinite_loop.error");
+            }
             if buffered_size == buf.len() {
                 break;
             }
@@ -297,7 +324,10 @@ impl SocketHandler for FrontRustls {
                     }
                     Ok(_sz) => {}
                     Err(e) => match e.kind() {
-                        ErrorKind::WouldBlock => can_write = false,
+                        ErrorKind::WouldBlock => {
+                            can_write = false;
+                            break;
+                        }
                         ErrorKind::ConnectionReset
                         | ErrorKind::ConnectionAborted
                         | ErrorKind::BrokenPipe => {
@@ -358,7 +388,13 @@ impl SocketHandler for FrontRustls {
             },
         }
 
+        let mut counter = 0;
         loop {
+            counter += 1;
+            if counter > MAX_LOOP_ITERATION {
+                error!("MAX_LOOP_ITERATION reached in FrontRustls::socket_write_vectored");
+                incr!("socket.write.infinite_loop.error");
+            }
             match self.session.write_tls(&mut self.stream) {
                 Ok(0) => {
                     //can_write = false;
@@ -366,7 +402,10 @@ impl SocketHandler for FrontRustls {
                 }
                 Ok(_sz) => {}
                 Err(e) => match e.kind() {
-                    ErrorKind::WouldBlock => can_write = false,
+                    ErrorKind::WouldBlock => {
+                        can_write = false;
+                        break;
+                    }
                     ErrorKind::ConnectionReset
                     | ErrorKind::ConnectionAborted
                     | ErrorKind::BrokenPipe => {