1 files changed, 50 insertions, 33 deletions
diff --git a/Swiften/TLS/UnitTest/ClientServerTest.cpp b/Swiften/TLS/UnitTest/ClientServerTest.cpp
index 24bd7c5..a356dcf 100644
--- a/Swiften/TLS/UnitTest/ClientServerTest.cpp
+++ b/Swiften/TLS/UnitTest/ClientServerTest.cpp
@@ -1,7 +1,7 @@
 /*
- * Copyright (c) 2010-2018 Isode Limited.
+ * Copyright (c) 2010-2019 Isode Limited.
  * All rights reserved.
  * See the COPYING file for more information.
  */
 
 #include <map>
@@ -406,11 +406,10 @@ class TLSEventToSafeByteArrayVisitor : public boost::static_visitor<SafeByteArra
         }
 
         SafeByteArray operator()(const TLSServerNameRequested&) const {
             return createSafeByteArray("");
         }
-
 };
 
 class TLSEventToStringVisitor : public boost::static_visitor<std::string> {
     public:
         std::string operator()(const TLSDataForNetwork& event) const {
@@ -457,10 +456,27 @@ class TLSClientServerEventHistory {
                 std::cout << event.first << " : " << boost::apply_visitor(TLSEventToStringVisitor(), event.second) << std::endl;
                 count++;
             }
         }
 
+        template<class TLSEventType>
+        boost::optional<TLSEventType> getEvent(const std::string& peer, size_t number = 0) {
+            for (const auto& pair : events) {
+                if (pair.first == peer) {
+                    if (pair.second.type() == typeid(TLSEventType)) {
+                        if (number == 0) {
+                            return boost::optional<TLSEventType>(boost::get<TLSEventType>(pair.second));
+                        }
+                        else {
+                            number--;
+                        }
+                    }
+                }
+            }
+            return {};
+        }
+
     private:
         void connectContext(const std::string& name, TLSContext* context) {
             connections_.push_back(context->onDataForNetwork.connect([=](const SafeByteArray& data) {
                 events.push_back(std::pair<std::string, TLSEvent>(name, TLSDataForNetwork{data}));
             }));
@@ -600,16 +616,16 @@ TEST(ClientServerTest, testClientServerBasicCommunication) {
     clientContext->connect();
 
     clientContext->handleDataFromApplication(createSafeByteArray("This is a test message from the client."));
     serverContext->handleDataFromApplication(createSafeByteArray("This is a test message from the server."));
 
-    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the client.")), safeByteArrayToString(boost::apply_visitor(TLSEventToSafeByteArrayVisitor(), std::find_if(events.events.begin(), events.events.end(), [](std::pair<std::string, TLSEvent>& event){
-        return event.first == "server" && (event.second.type() == typeid(TLSDataForApplication));
-    })->second)));
-    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the server.")), safeByteArrayToString(boost::apply_visitor(TLSEventToSafeByteArrayVisitor(), std::find_if(events.events.begin(), events.events.end(), [](std::pair<std::string, TLSEvent>& event){
-        return event.first == "client" && (event.second.type() == typeid(TLSDataForApplication));
-    })->second)));
+    auto firstMessageFromClient = events.getEvent<TLSDataForApplication>("server");
+    ASSERT_EQ(true, firstMessageFromClient.is_initialized());
+    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the client.")), safeByteArrayToString(firstMessageFromClient->data));
+    auto firstMessageFromServer = events.getEvent<TLSDataForApplication>("client");
+    ASSERT_EQ(true, firstMessageFromServer.is_initialized());
+    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the server.")), safeByteArrayToString(firstMessageFromServer->data));
 }
 
 TEST(ClientServerTest, testClientServerBasicCommunicationEncryptedPrivateKeyRightPassword) {
     auto clientContext = createTLSContext(TLSContext::Mode::Client);
     auto serverContext = createTLSContext(TLSContext::Mode::Server);
@@ -630,16 +646,16 @@ TEST(ClientServerTest, testClientServerBasicCommunicationEncryptedPrivateKeyRigh
     clientContext->connect();
 
     clientContext->handleDataFromApplication(createSafeByteArray("This is a test message from the client."));
     serverContext->handleDataFromApplication(createSafeByteArray("This is a test message from the server."));
 
-    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the client.")), safeByteArrayToString(boost::apply_visitor(TLSEventToSafeByteArrayVisitor(), std::find_if(events.events.begin(), events.events.end(), [](std::pair<std::string, TLSEvent>& event){
-        return event.first == "server" && (event.second.type() == typeid(TLSDataForApplication));
-    })->second)));
-    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the server.")), safeByteArrayToString(boost::apply_visitor(TLSEventToSafeByteArrayVisitor(), std::find_if(events.events.begin(), events.events.end(), [](std::pair<std::string, TLSEvent>& event){
-        return event.first == "client" && (event.second.type() == typeid(TLSDataForApplication));
-    })->second)));
+    auto firstMessageFromClient = events.getEvent<TLSDataForApplication>("server");
+    ASSERT_EQ(true, firstMessageFromClient.is_initialized());
+    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the client.")), safeByteArrayToString(firstMessageFromClient->data));
+    auto firstMessageFromServer = events.getEvent<TLSDataForApplication>("client");
+    ASSERT_EQ(true, firstMessageFromServer.is_initialized());
+    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the server.")), safeByteArrayToString(firstMessageFromServer->data));
 }
 
 TEST(ClientServerTest, testClientServerBasicCommunicationWithChainedCert) {
     auto clientContext = createTLSContext(TLSContext::Mode::Client);
     auto serverContext = createTLSContext(TLSContext::Mode::Server);
@@ -737,18 +753,19 @@ TEST(ClientServerTest, testClientServerSNIRequestedHostAvailable) {
     serverContext->accept();
     clientContext->connect("montague.example");
 
     clientContext->handleDataFromApplication(createSafeByteArray("This is a test message from the client."));
     serverContext->handleDataFromApplication(createSafeByteArray("This is a test message from the server."));
-    ASSERT_EQ("This is a test message from the client.", safeByteArrayToString(boost::apply_visitor(TLSEventToSafeByteArrayVisitor(), std::find_if(events.events.begin(), events.events.end(), [](std::pair<std::string, TLSEvent>& event){
-        return event.first == "server" && (event.second.type() == typeid(TLSDataForApplication));
-    })->second)));
-    ASSERT_EQ("This is a test message from the server.", safeByteArrayToString(boost::apply_visitor(TLSEventToSafeByteArrayVisitor(), std::find_if(events.events.begin(), events.events.end(), [](std::pair<std::string, TLSEvent>& event){
-        return event.first == "client" && (event.second.type() == typeid(TLSDataForApplication));
-    })->second)));
 
-    ASSERT_EQ("/CN=montague.example", boost::get<TLSConnected>(events.events[5].second).chain[0]->getSubjectName());
+    auto firstMessageFromClient = events.getEvent<TLSDataForApplication>("server");
+    ASSERT_EQ(true, firstMessageFromClient.is_initialized());
+    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the client.")), safeByteArrayToString(firstMessageFromClient->data));
+    auto firstMessageFromServer = events.getEvent<TLSDataForApplication>("client");
+    ASSERT_EQ(true, firstMessageFromServer.is_initialized());
+    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the server.")), safeByteArrayToString(firstMessageFromServer->data));
+
+    ASSERT_EQ("/CN=montague.example", events.getEvent<TLSConnected>("client")->chain[0]->getSubjectName());
 }
 
 TEST(ClientServerTest, testClientServerSNIRequestedHostUnavailable) {
     auto tlsFactories = std::make_shared<PlatformTLSFactories>();
     auto clientContext = createTLSContext(TLSContext::Mode::Client);
@@ -823,16 +840,16 @@ TEST(ClientServerTest, testClientServerBasicCommunicationWith2048BitDHParams) {
     clientContext->connect();
 
     clientContext->handleDataFromApplication(createSafeByteArray("This is a test message from the client."));
     serverContext->handleDataFromApplication(createSafeByteArray("This is a test message from the server."));
 
-    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the client.")), safeByteArrayToString(boost::apply_visitor(TLSEventToSafeByteArrayVisitor(), std::find_if(events.events.begin(), events.events.end(), [](std::pair<std::string, TLSEvent>& event){
-        return event.first == "server" && (event.second.type() == typeid(TLSDataForApplication));
-    })->second)));
-    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the server.")), safeByteArrayToString(boost::apply_visitor(TLSEventToSafeByteArrayVisitor(), std::find_if(events.events.begin(), events.events.end(), [](std::pair<std::string, TLSEvent>& event){
-        return event.first == "client" && (event.second.type() == typeid(TLSDataForApplication));
-    })->second)));
+    auto firstMessageFromClient = events.getEvent<TLSDataForApplication>("server");
+    ASSERT_EQ(true, firstMessageFromClient.is_initialized());
+    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the client.")), safeByteArrayToString(firstMessageFromClient->data));
+    auto firstMessageFromServer = events.getEvent<TLSDataForApplication>("client");
+    ASSERT_EQ(true, firstMessageFromServer.is_initialized());
+    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the server.")), safeByteArrayToString(firstMessageFromServer->data));
 }
 
 TEST(ClientServerTest, testClientServerBasicCommunicationWith1024BitDHParams) {
     auto clientContext = createTLSContext(TLSContext::Mode::Client);
     auto serverContext = createTLSContext(TLSContext::Mode::Server);
@@ -855,12 +872,12 @@ TEST(ClientServerTest, testClientServerBasicCommunicationWith1024BitDHParams) {
     clientContext->connect();
 
     clientContext->handleDataFromApplication(createSafeByteArray("This is a test message from the client."));
     serverContext->handleDataFromApplication(createSafeByteArray("This is a test message from the server."));
 
-    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the client.")), safeByteArrayToString(boost::apply_visitor(TLSEventToSafeByteArrayVisitor(), std::find_if(events.events.begin(), events.events.end(), [](std::pair<std::string, TLSEvent>& event){
-        return event.first == "server" && (event.second.type() == typeid(TLSDataForApplication));
-    })->second)));
-    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the server.")), safeByteArrayToString(boost::apply_visitor(TLSEventToSafeByteArrayVisitor(), std::find_if(events.events.begin(), events.events.end(), [](std::pair<std::string, TLSEvent>& event){
-        return event.first == "client" && (event.second.type() == typeid(TLSDataForApplication));
-    })->second)));
+    auto firstMessageFromClient = events.getEvent<TLSDataForApplication>("server");
+    ASSERT_EQ(true, firstMessageFromClient.is_initialized());
+    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the client.")), safeByteArrayToString(firstMessageFromClient->data));
+    auto firstMessageFromServer = events.getEvent<TLSDataForApplication>("client");
+    ASSERT_EQ(true, firstMessageFromServer.is_initialized());
+    ASSERT_EQ(safeByteArrayToString(createSafeByteArray("This is a test message from the server.")), safeByteArrayToString(firstMessageFromServer->data));
 }