2 files changed, 58 insertions, 23 deletions

diff --git a/src/Net/Connection.cpp b/src/Net/Connection.cpp
index d069fc9..8af02c6 100644
--- a/src/Net/Connection.cpp
+++ b/src/Net/Connection.cpp
@@ -37,7 +37,7 @@ void Connection::doHandshake() {
       return;
     
     // TODO: Error
-    disconnect();
+    doDisconnect();
     return;
   }
   
@@ -45,6 +45,25 @@ void Connection::doHandshake() {
   connectionHeader();
 }
 
+void Connection::doBye() {
+  if(state != BYE)
+    return;
+  
+  int ret = gnutls_bye(session, GNUTLS_SHUT_WR);
+  if(ret < 0) {
+    if(ret == GNUTLS_E_INTERRUPTED || ret == GNUTLS_E_AGAIN)
+      return;
+    
+    // TODO: Error
+    doDisconnect();
+    return;
+  }
+  
+  std::cout << "Bye!" << std::endl;
+  
+  doDisconnect();
+}
+
 void Connection::packetHeaderReceiveHandler(const void *data, unsigned long length) {
   if(state != PACKET_HEADER)
       return;
@@ -168,12 +187,36 @@ bool Connection::rawSend(const unsigned char *data, unsigned long length) {
   return true;
 }
 
+void Connection::sendReceive(short events) {
+  if(events & POLLHUP || events & POLLERR) {
+    doDisconnect();
+    return;
+  }
+  
+  if(state == HANDSHAKE) {
+    doHandshake();
+    return;
+  }
+  
+  if(state == BYE) {
+    doBye();
+    return;
+  }
+  
+  if(events & POLLIN)
+    doReceive();
+  
+  if(events & POLLOUT)
+    doSend();
+  
+  if(state == DISCONNECT && sendQueueEmpty()) 
+    bye();
+}
+
 void Connection::doDisconnect() {
   if(!isConnected())
     return;
   
-  gnutls_bye(session, GNUTLS_SHUT_WR);
-  
   shutdown(sock, SHUT_RDWR);
   close(sock);
   
diff --git a/src/Net/Connection.h b/src/Net/Connection.h
index a3670b0..01926e1 100644
--- a/src/Net/Connection.h
+++ b/src/Net/Connection.h
@@ -62,12 +62,23 @@ class Connection {
     void doReceive();
     void doSend();
     
+    void doBye();
+    
     void doDisconnect();
     
     bool receiveComplete() const {
       return (transR.length == transR.transmitted);
     }
     
+    void bye() {
+      if(state != DISCONNECT)
+        return;
+      
+      state = BYE;
+      
+      doBye();
+    }
+    
     // Prevent shallow copy
     Connection(const Connection &o);
     Connection& operator=(const Connection &o);
@@ -162,26 +173,7 @@ class Connection {
       return rawSend(reinterpret_cast<const unsigned char*>(packet.getRawData()), packet.getRawDataLength());
     }
     
-    void sendReceive(short events = POLLIN|POLLOUT) {
-      if(events & POLLHUP || events & POLLERR) {
-        doDisconnect();
-        return;
-      }
-      
-      if(state == HANDSHAKE) {
-        doHandshake();
-        return;
-      }
-      
-      if(events & POLLIN)
-        doReceive();
-      
-      if(events & POLLOUT)
-        doSend();
-      
-      if(state == DISCONNECT && sendQueueEmpty())
-        doDisconnect();
-    }
+    void sendReceive(short events = POLLIN|POLLOUT);
     
     bool sendQueueEmpty() const {return transS.empty();}