From 6939f053c9752dae586854bec31789beda9ba870 Mon Sep 17 00:00:00 2001
From: Sergey Yagovtsev <yse.sey@gmail.com>
Date: Sun, 18 Sep 2016 16:43:43 +0300
Subject: [PATCH] Add checkResult function for socket

---
 include/profiler/easy_socket.h |  3 +++
 profiler_gui/main_window.cpp   | 22 ++++++++++++++++-----
 src/easy_socket.cpp            | 36 +++++++++++++++++++++++++++++++---
 3 files changed, 53 insertions(+), 8 deletions(-)

diff --git a/include/profiler/easy_socket.h b/include/profiler/easy_socket.h
index 71c7807..f7d8f4f 100644
--- a/include/profiler/easy_socket.h
+++ b/include/profiler/easy_socket.h
@@ -48,6 +48,9 @@ public:
         CONNECTION_STATE_DISCONNECTED
     };
 private:
+
+    void checkResult(int result);
+
 #ifdef _WIN32
     typedef SOCKET socket_t;
 #else
diff --git a/profiler_gui/main_window.cpp b/profiler_gui/main_window.cpp
index 9f45f16..924b0af 100644
--- a/profiler_gui/main_window.cpp
+++ b/profiler_gui/main_window.cpp
@@ -121,7 +121,6 @@ EasyMainWindow::EasyMainWindow() : Parent(), m_treeWidget(nullptr), m_graphicsVi
     m_hostString->setInputMask("000.000.000.000;");
     m_hostString->setValidator(&regValidator);
     m_hostString->setText("127.0.0.1");
-    //m_hostString->setText("192.224.4.109");
 
     fileToolBar->addWidget(m_hostString);
 
@@ -342,7 +341,12 @@ void EasyMainWindow::listen()
             bytes = m_easySocket.receive(buffer, buffer_size);
             if(bytes == -1)
             {
-                isListen = false;
+                if(m_easySocket.state() == EasySocket::CONNECTION_STATE_DISCONNECTED)
+                {
+                    isListen = false;
+                }
+                seek = 0;
+                bytes = 0;
                 continue;
             }
             seek = 0;
@@ -350,6 +354,11 @@ void EasyMainWindow::listen()
             
         char *buf = &buffer[seek];
 
+        if(bytes == 0){
+            isListen = false;
+            continue;
+        }
+
         if (bytes > 0)
         {
             profiler::net::Message* message = (profiler::net::Message*)buf;
@@ -434,8 +443,11 @@ void EasyMainWindow::listen()
 
                     if(bytes == -1)
                     {
-                        isListen = false;
-                        neededSize = 0;
+                        if(m_easySocket.state() == EasySocket::CONNECTION_STATE_DISCONNECTED)
+                        {
+                            isListen = false;
+                            neededSize = 0;
+                        }
                         continue;
                     }
 
@@ -461,7 +473,7 @@ void EasyMainWindow::listen()
 
         }
     }
-    m_easySocket.setState(EasySocket::CONNECTION_STATE_DISCONNECTED);
+
     delete [] buffer;
 }
 
diff --git a/src/easy_socket.cpp b/src/easy_socket.cpp
index 70b91a5..855a493 100644
--- a/src/easy_socket.cpp
+++ b/src/easy_socket.cpp
@@ -21,6 +21,7 @@ along with this program.If not, see <http://www.gnu.org/licenses/>.
 
 #ifndef _WIN32
 #include <strings.h>
+#include <errno.h>
 
 int EasySocket::bind(uint16_t portno)
 {
@@ -45,25 +46,34 @@ EasySocket::~EasySocket()
 int EasySocket::send(const void *buf, size_t nbyte)
 {
     if(m_replySocket <= 0)  return -1;
-    return ::write(m_replySocket,buf,nbyte);
+    int res = ::write(m_replySocket,buf,nbyte);
+    checkResult(res);
+    return res;
 }
 
 int EasySocket::receive(void *buf, size_t nbyte)
 {
     if(m_replySocket <= 0) return -1;
-    return ::read(m_replySocket,buf,nbyte);
+    int res = ::read(m_replySocket,buf,nbyte);
+    checkResult(res);
+    return res;
 }
 
 int EasySocket::listen(int count)
 {
     if(m_socket < 0 ) return -1;
-    return ::listen(m_socket,count);
+    int res = ::listen(m_socket,count);
+    checkResult(res);
+    return res;
 }
 
 int EasySocket::accept()
 {
     if(m_socket < 0 ) return -1;
     m_replySocket = ::accept(m_socket,nullptr,nullptr);
+
+    checkResult(m_replySocket);
+
     return m_replySocket;
 }
 
@@ -104,6 +114,26 @@ int EasySocket::connect()
     }
     return res;
 }
+#include <string.h>
+void EasySocket::checkResult(int result)
+{
+    if(result >= 0){
+        m_state = CONNECTION_STATE_SUCCESS;
+        return;
+    }else if(result == -1){
+        //printf("Errno: %s\n", strerror(errno));
+        switch(errno){
+        case ECONNABORTED:
+        case ECONNRESET:
+            m_state = CONNECTION_STATE_SUCCESS;
+            break;
+        default:
+            break;
+        }
+
+    }
+}
+
 #else
 
 #pragma comment (lib, "Ws2_32.lib")