From c2049e530873dce4ffbdc94573e8ab9c671c13a1 Mon Sep 17 00:00:00 2001
From: Ryan Huffman <ryanhuffman@gmail.com>
Date: Fri, 4 Dec 2015 16:25:22 -0800
Subject: [PATCH] Fix udt-test to work with ReceivedMessage updates

---
 tools/udt-test/src/UDTTest.cpp | 42 +++++++++++++++++++++++++++-------
 tools/udt-test/src/UDTTest.h   | 11 ++++++++-
 2 files changed, 44 insertions(+), 9 deletions(-)

diff --git a/tools/udt-test/src/UDTTest.cpp b/tools/udt-test/src/UDTTest.cpp
index 3d0d03b8d2..d4054527f9 100644
--- a/tools/udt-test/src/UDTTest.cpp
+++ b/tools/udt-test/src/UDTTest.cpp
@@ -176,10 +176,36 @@ UDTTest::UDTTest(int& argc, char** argv) :
     } else {
         // this is a receiver - in case there are ordered packets (messages) being sent to us make sure that we handle them
         // so that they can be verified
-        // TODO Fix support for message testing
-        //_socket.setMessageHandler(
-            //[this](std::unique_ptr<udt::PacketList> packetList) { handlePacketList(std::move(packetList)); });
+        _socket.setMessageHandler(
+            [this](std::unique_ptr<udt::Packet> packet) {
+                auto messageNumber = packet->getMessageNumber();
+                auto it = _pendingMessages.find(messageNumber);
+
+                if (it == _pendingMessages.end()) {
+                    auto message = std::unique_ptr<Message>(new Message { messageNumber, packet->readAll() });
+                    message->data.reserve(_messageSize);
+                    if (packet->getPacketPosition() == udt::Packet::ONLY) {
+                        handleMessage(std::move(message));
+                    } else {
+                        _pendingMessages[messageNumber] = std::move(message);
+                    }
+                } else {
+                    auto& message = it->second;
+                    message->data.append(packet->readAll());
+
+                    if (packet->getPacketPosition() == udt::Packet::LAST) {
+                        handleMessage(std::move(message));
+                        _pendingMessages.erase(it);
+                    }
+                }
+
+        });
     }
+    _socket.setMessageFailureHandler(
+        [this](HifiSockAddr from, udt::Packet::MessageNumber messageNumber) {
+            _pendingMessages.erase(messageNumber);
+        }
+    );
     
     // the sender reports stats every 100 milliseconds, unless passed a custom value
     
@@ -309,11 +335,11 @@ void UDTTest::sendPacket() {
     
 }
 
-void UDTTest::handlePacketList(std::unique_ptr<udt::PacketList> packetList) {
+void UDTTest::handleMessage(std::unique_ptr<Message> message) {
     // generate the byte array that should match this message - using the same seed the sender did
     
     int packetSize = udt::Packet::maxPayloadSize(true);
-    int messageSize = packetList->getMessageSize();
+    int messageSize = message->data.size();
     
     QByteArray messageData(messageSize, 0);
    
@@ -324,13 +350,13 @@ void UDTTest::handlePacketList(std::unique_ptr<udt::PacketList> packetList) {
         messageData.replace(i, sizeof(randomInt), reinterpret_cast<char*>(&randomInt), sizeof(randomInt));
     }
     
-    bool dataMatch = messageData == packetList->getMessage();
+    bool dataMatch = messageData == message->data;
     
-    Q_ASSERT_X(dataMatch, "UDTTest::handlePacketList",
+    Q_ASSERT_X(dataMatch, "UDTTest::handleMessage",
                "received message did not match expected message (from seeded random number generation).");
     
     if (!dataMatch) {
-        qCritical() << "UDTTest::handlePacketList" << "received message did not match expected message"
+        qCritical() << "UDTTest::handleMessage" << "received message did not match expected message"
             << "(from seeded random number generation).";
     }
 }
diff --git a/tools/udt-test/src/UDTTest.h b/tools/udt-test/src/UDTTest.h
index fda57cc183..3a4b1b897d 100644
--- a/tools/udt-test/src/UDTTest.h
+++ b/tools/udt-test/src/UDTTest.h
@@ -23,6 +23,13 @@
 #include <udt/Constants.h>
 #include <udt/Socket.h>
 
+#include <ReceivedMessage.h>
+
+struct Message {
+    udt::MessageNumber messageNumber;
+    QByteArray data;
+};
+
 class UDTTest : public QCoreApplication {
     Q_OBJECT
 public:
@@ -34,7 +41,7 @@ public slots:
     
 private:
     void parseArguments();
-    void handlePacketList(std::unique_ptr<udt::PacketList> packetList);
+    void handleMessage(std::unique_ptr<Message> message);
     
     void sendInitialPackets(); // fills the queue with packets to start
     void sendPacket(); // constructs and sends a packet according to the test parameters
@@ -53,6 +60,8 @@ private:
     bool _sendOrdered { false }; // whether to send ordered packets
     
     int _messageSize { 10000000 }; // number of bytes per message while sending ordered
+
+    std::unordered_map<udt::Packet::MessageNumber, std::unique_ptr<Message>> _pendingMessages;
     
     std::random_device _randomDevice;
     std::mt19937 _generator { _randomDevice() }; // random number generator for ordered data testing