From 3c16f782c36348fa3b817d0d0ba779e0d06b442b Mon Sep 17 00:00:00 2001
From: dec05eba <0xdec05eba@gmail.com>
Date: Sat, 9 Jun 2018 20:14:57 +0200
Subject: Fix receive

---
 src/BootstrapConnection.cpp |  1 +
 src/DirectConnection.cpp    | 20 ++++++++++++--------
 2 files changed, 13 insertions(+), 8 deletions(-)

(limited to 'src')

diff --git a/src/BootstrapConnection.cpp b/src/BootstrapConnection.cpp
index 361f13a..ccd2141 100644
--- a/src/BootstrapConnection.cpp
+++ b/src/BootstrapConnection.cpp
@@ -33,6 +33,7 @@ namespace sibs
     // TODO: This is vulnerable against MitM attack, replace with asymmetric cryptography, get data signed with server private key and verify against known server public key
     void BootstrapConnection::receiveDataFromServer(std::shared_ptr<DirectConnectionPeer> peer, const void *data, const usize size)
     {
+        Log::debug("BootstrapConnection: Received subscriber(s) from bootstrap node");
         sibs::SafeDeserializer deserializer((const u8*)data, size);
         PubsubKey pubsubKey;
         deserializer.extract(pubsubKey.data.data(), PUBSUB_KEY_LENGTH);
diff --git a/src/DirectConnection.cpp b/src/DirectConnection.cpp
index 37f5b34..ce63eaf 100644
--- a/src/DirectConnection.cpp
+++ b/src/DirectConnection.cpp
@@ -163,7 +163,7 @@ namespace sibs
         std::vector<char> data;
         data.reserve(MAX_RECEIVED_DATA_SIZE);
         
-        Log::debug("DirectConnections::recieveData(): waiting for events");
+        Log::debug("DirectConnections::receiveData(): waiting for events");
         std::set<UDTSOCKET> readfds;
         while(alive)
         {
@@ -188,7 +188,8 @@ namespace sibs
             for(UDTSOCKET receivedDataFromPeer : readfds)
             {
                 bool peerDisconnected = false;
-                int receivedDataStatus = receiveDataFromPeer(receivedDataFromPeer, data.data());
+                usize receivedTotalSize = 0;
+                int receivedDataStatus = receiveDataFromPeer(receivedDataFromPeer, data.data(), &receivedTotalSize);
                 if(receivedDataStatus == 0)
                 {
                     peersMutex.lock();
@@ -198,7 +199,7 @@ namespace sibs
                     {
                         Log::debug("DirectConnection: Received data from peer: (ip: %s, port: %d)", peer->address.getAddress().c_str(), peer->address.getPort());
                         if(peer->receiveDataCallbackFunc)
-                            peer->receiveDataCallbackFunc(peer, data.data(), data.size());
+                            peer->receiveDataCallbackFunc(peer, data.data(), receivedTotalSize);
                     }
                     catch(std::exception &e)
                     {
@@ -228,10 +229,10 @@ namespace sibs
         }
     }
     
-    int DirectConnections::receiveDataFromPeer(const int socket, char *output)
+    int DirectConnections::receiveDataFromPeer(const int socket, char *output, usize *receivedTotalSize)
     {
-        usize receivedTotalSize = 0;
-        while(receivedTotalSize < MAX_RECEIVED_DATA_SIZE)
+        *receivedTotalSize = 0;
+        while(*receivedTotalSize < MAX_RECEIVED_DATA_SIZE)
         {
             int dataAvailableSize;
             int receiveSizeDataTypeSize = sizeof(dataAvailableSize);
@@ -241,7 +242,10 @@ namespace sibs
                 return UDT::getlasterror_code();
             }
             
-            int receivedSize = UDT::recv(socket, &output[receivedTotalSize], MAX_RECEIVED_DATA_SIZE - receivedTotalSize, 0);
+            if(dataAvailableSize == 0)
+                return 0;
+            
+            int receivedSize = UDT::recv(socket, &output[*receivedTotalSize], MAX_RECEIVED_DATA_SIZE - *receivedTotalSize, 0);
             if(receivedSize == UDT::ERROR)
             {
                 Log::error("UDT: Failed to receive data, error: %s (%d)", UDT::getlasterror_desc(), UDT::getlasterror_code());
@@ -253,7 +257,7 @@ namespace sibs
                 return 0;
             }
             
-            receivedTotalSize += receivedSize;
+            (*receivedTotalSize) += receivedSize;
         }
         
         Log::error("UDT: Received too much data, ignoring...");
-- 
cgit v1.2.3