Fix receive

3 files changed, 14 insertions, 9 deletions

diff --git a/include/sibs/DirectConnection.hpp b/include/sibs/DirectConnection.hpp
index 134eade..10359f9 100644
--- a/include/sibs/DirectConnection.hpp
+++ b/include/sibs/DirectConnection.hpp
@@ -57,7 +57,7 @@ namespace sibs
     private:
         void connect(const Ipv4 &address, bool rendezvous, bool reuseAddr, PubSubConnectCallback connectCallbackFunc, PubSubReceiveDataCallback receiveDataCallbackFunc);
         void receiveData();
-        int receiveDataFromPeer(const int socket, char *output);
+        int receiveDataFromPeer(const int socket, char *output, usize *receivedTotalSize);
         bool removePeer(int peerSocket);
     private:
         u16 port;
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...");