1 files changed, 46 insertions, 3 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);