1 files changed, 28 insertions, 9 deletions
diff --git a/src/Net/Connection.cpp b/src/Net/Connection.cpp
index 4e3fee4..19e7bf1 100644
--- a/src/Net/Connection.cpp
+++ b/src/Net/Connection.cpp
@@ -25,15 +25,17 @@
 namespace Mad {
 namespace Net {
 
-void Connection::handshake() {
-  state = HANDSHAKE;
+void Connection::doHandshake() {
+  if(state != HANDSHAKE)
+    return;
   
   int ret = gnutls_handshake(session);
   if(ret < 0) {
     if(ret == GNUTLS_E_INTERRUPTED || ret == GNUTLS_E_AGAIN)
       return;
     
-    // Error... disconnect
+    // TODO: Error
+    disconnect();
     return;
   }
   
@@ -42,8 +44,14 @@ void Connection::handshake() {
 }
 
 void Connection::packetHeaderReceiveHandler(const void *data, unsigned long length) {
-  if(length != sizeof(Packet::Data))
-    return; // Error... disconnect?
+  if(state != PACKET_HEADER)
+      return;
+  
+  if(length != sizeof(Packet::Data)) {
+    // TODO: Error
+    disconnect();
+    return;
+  }
   
   header = *reinterpret_cast<const Packet::Data*>(data);
   
@@ -59,8 +67,14 @@ void Connection::packetHeaderReceiveHandler(const void *data, unsigned long leng
 }
 
 void Connection::packetDataReceiveHandler(const void *data, unsigned long length) {
-  if(length != header.length)
-    return; // Error... disconnect?
+  if(state != PACKET_DATA)
+    return;
+  
+  if(length != header.length) {
+    // TODO: Error
+    disconnect();
+    return;
+  }
   
   signal(this, Packet(header.type, header.requestId, data, length));
   
@@ -80,7 +94,8 @@ void Connection::doReceive() {
     if(ret == GNUTLS_E_INTERRUPTED || ret == GNUTLS_E_AGAIN)
       return;
     
-    // Error... disconnect?
+    // TODO: Error
+    disconnect();
     return;
   }
   
@@ -126,7 +141,8 @@ void Connection::doSend() {
       if(ret == GNUTLS_E_INTERRUPTED || ret == GNUTLS_E_AGAIN)
         return;
       
-      // Error... disconnect?
+      // TODO: Error
+      disconnect();
       return;
     }
     
@@ -171,6 +187,9 @@ void Connection::disconnect() {
 struct pollfd Connection::getPollfd() const {
   struct pollfd fd = {sock, (receiveComplete() ? 0 : POLLIN) | (sendQueueEmpty() ? 0 : POLLOUT), 0};
   
+  if(state == HANDSHAKE)
+    fd.events = ((gnutls_record_get_direction(session) == 0) ? POLLIN : POLLOUT);
+  
   return fd;
 }