From cbfd65833f6cf90f9b12bf2f7b76161d628a8bb2 Mon Sep 17 00:00:00 2001
From: Piotr Duszynski <2052966+drk1wi@users.noreply.github.com>
Date: Tue, 29 Aug 2023 10:20:50 +0200
Subject: [PATCH] Update connection.cpp

a small fix to address a potential bypass
---
 src/connection.cpp | 30 ++++++++++++++++++++----------
 1 file changed, 20 insertions(+), 10 deletions(-)

diff --git a/src/connection.cpp b/src/connection.cpp
index c4a1685..4294420 100644
--- a/src/connection.cpp
+++ b/src/connection.cpp
@@ -38,6 +38,7 @@
 #include <sys/types.h>
 #include <sys/socket.h>
 #include <sys/ioctl.h>
+#include <linux/sockios.h>
 #include <time.h> 
 #include "Threads.h"
 #include "connection.h"
@@ -127,15 +128,24 @@ void* process_connection(void *arg)
   				    
   				    select_return = select(threads[tid].clients[i], &read_mask, (fd_set *)0, (fd_set *)0, &tv);
 
-					if(select_return <= 0) /* [timeout=0, -1= ERROR] is returned */
-					{
-						n=1;
-					}
-					else
-					{
-						buffer_size=configuration->mapPort2Buffer(original_port);
-						n = recv(threads[tid].clients[i],buffer,buffer_size, 0);	
-					}
+					 if (select_return < 0) /* [timeout=0, -1= ERROR] is returned */
+                {
+                    n = -1;
+                } 
+           else {
+                    n = 0;
+                    int data_to_be_read_size = 0;
+
+                    if (ioctl(threads[tid].clients[i], FIONREAD, &data_to_be_read_size) < 0) {
+                        perror("ioctl failed");
+                    }
+
+                    if (data_to_be_read_size > 0) {
+                        buffer_size = data_to_be_read_size;
+                        n = recv(threads[tid].clients[i], buffer, buffer_size, 0);
+                    }
+
+                }
 				}
 			
 				// deal with different recv buffer size
@@ -306,4 +316,4 @@ void* process_connection(void *arg)
 	}
 
 	return 0;
-}
\ No newline at end of file
+}