From c316734e8f34bc1eee5d57ead8ed0fc4cecbda04 Mon Sep 17 00:00:00 2001
From: Sergey Yagovtsev <yse.sey@gmail.com>
Date: Fri, 16 Sep 2016 22:51:15 +0300
Subject: [PATCH] Fix disconnect problem on downloading

---
 include/profiler/easy_socket.h | 14 +++++++++++++-
 profiler_gui/main_window.cpp   | 30 ++++++++++++++++++++++++++++--
 src/easy_socket.cpp            | 14 ++++++++------
 3 files changed, 49 insertions(+), 9 deletions(-)

diff --git a/include/profiler/easy_socket.h b/include/profiler/easy_socket.h
index ca5d5f4..71c7807 100644
--- a/include/profiler/easy_socket.h
+++ b/include/profiler/easy_socket.h
@@ -39,6 +39,15 @@ along with this program.If not, see <http://www.gnu.org/licenses/>.
 
 class PROFILER_API EasySocket
 {
+public:
+    enum ConnectionState
+    {
+        CONNECTION_STATE_UNKNOWN,
+        CONNECTION_STATE_SUCCESS,
+
+        CONNECTION_STATE_DISCONNECTED
+    };
+private:
 #ifdef _WIN32
     typedef SOCKET socket_t;
 #else
@@ -58,7 +67,7 @@ class PROFILER_API EasySocket
 #endif
 
 
-
+    ConnectionState m_state = CONNECTION_STATE_UNKNOWN;
 public:
     EasySocket();
     ~EasySocket();
@@ -71,6 +80,9 @@ public:
 
     bool setAddress(const char* serv, uint16_t port);
     int connect();
+
+    void setState(ConnectionState state){m_state=state;}
+    ConnectionState state() const{return m_state;}
 };
 
 #endif // EASY________SOCKET_________H
diff --git a/profiler_gui/main_window.cpp b/profiler_gui/main_window.cpp
index aa43502..9f45f16 100644
--- a/profiler_gui/main_window.cpp
+++ b/profiler_gui/main_window.cpp
@@ -335,10 +335,16 @@ void EasyMainWindow::listen()
     int bytes = 0;
 
     static auto timeBegin = std::chrono::system_clock::now();
-    while (true)
+    bool isListen = true;
+    while (isListen)
     {
         if ((bytes - seek) == 0){
             bytes = m_easySocket.receive(buffer, buffer_size);
+            if(bytes == -1)
+            {
+                isListen = false;
+                continue;
+            }
             seek = 0;
         }
             
@@ -399,7 +405,8 @@ void EasyMainWindow::listen()
                 m_receivedProfileData.clear();
                 //loadFile(QString(tempfilename.c_str()));
                 m_recFrames = false;
-                return;
+                isListen = false;
+                continue;
             }
             break;
             case profiler::net::MESSAGE_TYPE_REPLY_BLOCKS:
@@ -424,6 +431,14 @@ void EasyMainWindow::listen()
                 while (neededSize > 0)
                 {
                     bytes = m_easySocket.receive(buffer, buffer_size);
+
+                    if(bytes == -1)
+                    {
+                        isListen = false;
+                        neededSize = 0;
+                        continue;
+                    }
+
                     buf = &buffer[0];
                     int toWrite = std::min(bytes, neededSize);
                     m_receivedProfileData.write(buf, toWrite);
@@ -446,6 +461,8 @@ void EasyMainWindow::listen()
 
         }
     }
+    m_easySocket.setState(EasySocket::CONNECTION_STATE_DISCONNECTED);
+    delete [] buffer;
 }
 
 void TcpReceiverThread::run()
@@ -623,6 +640,8 @@ void EasyMainWindow::onCollapseAllClicked(bool)
 
 void EasyMainWindow::onConnectClicked(bool)
 {
+    if(m_isConnected)
+        return;
     int res = m_easySocket.setAddress(m_hostString->text().toStdString().c_str(), m_portString->text().toUShort());
 
     res = m_easySocket.connect();
@@ -721,6 +740,13 @@ void EasyMainWindow::onCaptureClicked(bool)
 
     m_downloading = false;
 
+    if(m_easySocket.state() == EasySocket::CONNECTION_STATE_DISCONNECTED)
+    {
+        QMessageBox::warning(this,"Warning" ,"Application was disconnected",QMessageBox::Close);
+        m_isConnected = false;
+        return;
+    }
+
     std::string tempfilename = "test_rec.prof";
     loadFile(QString(tempfilename.c_str()));
     /*m_isConnected = (m_server->state() == QTcpSocket::ConnectedState);
diff --git a/src/easy_socket.cpp b/src/easy_socket.cpp
index 75fcc27..70b91a5 100644
--- a/src/easy_socket.cpp
+++ b/src/easy_socket.cpp
@@ -36,12 +36,6 @@ int EasySocket::bind(uint16_t portno)
 EasySocket::EasySocket()
 {
     m_socket = socket(AF_INET, SOCK_STREAM, 0);
-    struct timeval tv;
-
-    tv.tv_sec = 1;
-    tv.tv_usec = 0;
-
-    //setsockopt(m_socket, SOL_SOCKET, SO_RCVTIMEO, (char *)&tv,sizeof(struct timeval));
 }
 
 EasySocket::~EasySocket()
@@ -98,6 +92,14 @@ int EasySocket::connect()
     }
     int res = ::connect(m_socket,(struct sockaddr *) &serv_addr,sizeof(serv_addr));
     if(res == 0){
+
+        struct timeval tv;
+
+        tv.tv_sec = 1;
+        tv.tv_usec = 0;
+
+        setsockopt(m_socket, SOL_SOCKET, SO_RCVTIMEO, (char *)&tv,sizeof(struct timeval));
+
         m_replySocket = m_socket;
     }
     return res;