Fix: WouldBlock in SocketHandler::socket_write breaks properly

Signed-off-by: Eloi DEMOLIS <eloi.demolis@clever-cloud.com>
sozu-proxy · Nov 23, 2023 · 58a7f03 · 58a7f03
1 parent af17580
commit 58a7f03
Showing 1 changed file with 41 additions and 2 deletions.
diff --git 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,15 +388,24 @@ 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;
                     break;
                 }
                 Ok(_sz) => {}
                 Err(e) => match e.kind() {
-                    ErrorKind::WouldBlock => can_write = false,
+                    ErrorKind::WouldBlock => {
+                        can_write = false;
+                        break;
+                    }
                     ErrorKind::ConnectionReset
                     | ErrorKind::ConnectionAborted
                     | ErrorKind::BrokenPipe => {