1 files changed, 30 insertions, 10 deletions

diff --git a/src/Net/Connection.h b/src/Net/Connection.h
index 381ed17..a5bda5a 100644
--- a/src/Net/Connection.h
+++ b/src/Net/Connection.h
@@ -42,6 +42,10 @@ class Connection {
       sigc::slot<void,const void*,unsigned long> notify;
     };
     
+    enum State {
+      DISCONNECTED, HANDSHAKE, CONNECTION_HEADER, PACKET_HEADER, PACKET_DATA
+    } state;
+    
     Transmission transR;
     std::queue<Transmission> transS;
     
@@ -72,23 +76,33 @@ class Connection {
       unsigned char protVerMax;
     };
     
-    virtual gnutls_session_t& getSession() = 0;
+    int sock;
+    gnutls_session_t session;
+    
+    IPAddress *peer;
     
-    bool rawReceive(unsigned long length,
-        const sigc::slot<void,const void*,unsigned long> &notify);
+    void handshake();
+    virtual void connectionHeader() = 0;
+    
+    bool rawReceive(unsigned long length, const sigc::slot<void,const void*,unsigned long> &notify);
     bool rawSend(const unsigned char *data, unsigned long length);
     
     bool enterReceiveLoop() {
+      state = PACKET_HEADER;
+      
       return rawReceive(sizeof(Packet::Data), sigc::mem_fun(this, &Connection::packetHeaderReceiveHandler));
     }
     
   public:
-    Connection() {
+    Connection() : state(DISCONNECTED), peer(0) {
       transR.length = transR.transmitted = 0;
       transR.data = 0;
     }
     
     virtual ~Connection() {
+      if(isConnected())
+        disconnect();
+      
       if(transR.data)
         delete [] transR.data;
       
@@ -96,24 +110,30 @@ class Connection {
         delete [] transS.front().data;
         transS.pop();
       }
-        
     }
     
-    virtual bool isConnected() const = 0;
-    virtual const IPAddress* getPeer() const = 0;
+    bool isConnected() const {return (state != DISCONNECTED);}
+    bool isConnecting() {
+      return (state == HANDSHAKE || state == CONNECTION_HEADER); 
+    }
     
-    virtual bool isConnecting() const = 0;
+    const IPAddress* getPeer() {return peer;}
     
-    virtual bool dataPending() const = 0;
+    void disconnect();
     
     bool send(const Packet &packet) {
-      if(isConnecting())
+      if(!isConnected() || isConnecting())
         return false;
       
       return rawSend(reinterpret_cast<const unsigned char*>(packet.getRawData()), packet.getRawDataLength());
     }
     
     void sendReceive() {
+      if(state == HANDSHAKE) {
+        handshake();
+        return;
+      }
+      
       doReceive();
       doSend();
     }