1 files changed, 25 insertions, 9 deletions

diff --git a/src/Net/Connection.h b/src/Net/Connection.h
index adb677a..a3670b0 100644
--- a/src/Net/Connection.h
+++ b/src/Net/Connection.h
@@ -44,7 +44,7 @@ class Connection {
     };
     
     enum State {
-      DISCONNECTED, HANDSHAKE, CONNECTION_HEADER, PACKET_HEADER, PACKET_DATA
+      DISCONNECTED, HANDSHAKE, CONNECTION_HEADER, PACKET_HEADER, PACKET_DATA, DISCONNECT, BYE
     } state;
     
     Transmission transR;
@@ -52,7 +52,7 @@ class Connection {
     
     Packet::Data header;
     
-    sigc::signal<void,const Connection*,const Packet&> signal;
+    sigc::signal<void,Connection*,const Packet&> signal;
     
     void doHandshake();
     
@@ -62,6 +62,8 @@ class Connection {
     void doReceive();
     void doSend();
     
+    void doDisconnect();
+    
     bool receiveComplete() const {
       return (transR.length == transR.transmitted);
     }
@@ -103,7 +105,7 @@ class Connection {
     bool rawSend(const unsigned char *data, unsigned long length);
     
     bool enterReceiveLoop() {
-      if(!isConnected())
+      if(!isConnected() || isDisconnecting())
         return false;
       
       state = PACKET_HEADER;
@@ -119,7 +121,7 @@ class Connection {
     
     virtual ~Connection() {
       if(isConnected())
-        disconnect();
+        doDisconnect();
       
       if(transR.data)
         delete [] transR.data;
@@ -131,19 +133,30 @@ class Connection {
     }
     
     bool isConnected() const {return (state != DISCONNECTED);}
-    bool isConnecting() {
+    bool isConnecting() const {
       return (state == HANDSHAKE || state == CONNECTION_HEADER);
     }
     
+    bool isDisconnecting() const {
+      return (state == DISCONNECT || state == BYE);
+    }
+    
     const IPAddress* getPeer() {return peer;}
     int getSocket() const {return sock;}
     
-    void disconnect();
+    void disconnect() {
+      if(isConnected() && !isDisconnecting()) {
+        state = DISCONNECT;
+        
+        if(sendQueueEmpty())
+          doDisconnect();
+      }
+    }
     
     struct pollfd getPollfd() const;
     
     bool send(const Packet &packet) {
-      if(!isConnected() || isConnecting())
+      if(!isConnected() || isConnecting() || isDisconnecting())
         return false;
       
       return rawSend(reinterpret_cast<const unsigned char*>(packet.getRawData()), packet.getRawDataLength());
@@ -151,7 +164,7 @@ class Connection {
     
     void sendReceive(short events = POLLIN|POLLOUT) {
       if(events & POLLHUP || events & POLLERR) {
-        disconnect();
+        doDisconnect();
         return;
       }
       
@@ -165,11 +178,14 @@ class Connection {
       
       if(events & POLLOUT)
         doSend();
+      
+      if(state == DISCONNECT && sendQueueEmpty())
+        doDisconnect();
     }
     
     bool sendQueueEmpty() const {return transS.empty();}
     
-    sigc::signal<void,const Connection*,const Packet&> signalReceive() const {return signal;}
+    sigc::signal<void,Connection*,const Packet&> signalReceive() const {return signal;}
     
     static void init() {
       gnutls_global_init();