From 1fb25829692e5455c0edec96226af295957d99b4 Mon Sep 17 00:00:00 2001
From: John Ericson <John.Ericson@Obsidian.Systems>
Date: Fri, 19 Jan 2024 18:42:27 -0500
Subject: [PATCH] Create unit tests for the serve proto handshake

Co-authored-by: Robert Hensing <roberth@users.noreply.github.com>
---
 .../serve-protocol/handshake-to-client.bin    | Bin 0 -> 16 bytes
 tests/unit/libstore/serve-protocol.cc         | 110 ++++++++++++++++++
 2 files changed, 110 insertions(+)
 create mode 100644 tests/unit/libstore/data/serve-protocol/handshake-to-client.bin

diff --git a/tests/unit/libstore/data/serve-protocol/handshake-to-client.bin b/tests/unit/libstore/data/serve-protocol/handshake-to-client.bin
new file mode 100644
index 0000000000000000000000000000000000000000..15ba4b5e3d96e388637107542f6eb9f7e94ac708
GIT binary patch
literal 16
RcmX^8E+~Wn1em}i0{|m{0%8CF

literal 0
HcmV?d00001

diff --git a/tests/unit/libstore/serve-protocol.cc b/tests/unit/libstore/serve-protocol.cc
index 8f256d1e6..597c0b570 100644
--- a/tests/unit/libstore/serve-protocol.cc
+++ b/tests/unit/libstore/serve-protocol.cc
@@ -1,3 +1,4 @@
+#include <thread>
 #include <regex>
 
 #include <nlohmann/json.hpp>
@@ -6,6 +7,7 @@
 #include "serve-protocol.hh"
 #include "serve-protocol-impl.hh"
 #include "build-result.hh"
+#include "file-descriptor.hh"
 #include "tests/protocol.hh"
 #include "tests/characterization.hh"
 
@@ -401,4 +403,112 @@ VERSIONED_CHARACTERIZATION_TEST(
         },
     }))
 
+TEST_F(ServeProtoTest, handshake_log)
+{
+    CharacterizationTest::writeTest("handshake-to-client", [&]() -> std::string {
+        StringSink toClientLog;
+
+        Pipe toClient, toServer;
+        toClient.create();
+        toServer.create();
+
+        ServeProto::Version clientResult, serverResult;
+
+        auto thread = std::thread([&]() {
+            FdSink out { toServer.writeSide.get() };
+            FdSource in0 { toClient.readSide.get() };
+            TeeSource in { in0, toClientLog };
+            clientResult = ServeProto::BasicClientConnection::handshake(
+                out, in, defaultVersion, "blah");
+        });
+
+        {
+            FdSink out { toClient.writeSide.get() };
+            FdSource in { toServer.readSide.get() };
+            serverResult = ServeProto::BasicServerConnection::handshake(
+                out, in, defaultVersion);
+        };
+
+        thread.join();
+
+        return std::move(toClientLog.s);
+    });
+}
+
+/// Has to be a `BufferedSink` for handshake.
+struct NullBufferedSink : BufferedSink {
+    void writeUnbuffered(std::string_view data) override { }
+};
+
+TEST_F(ServeProtoTest, handshake_client_replay)
+{
+    CharacterizationTest::readTest("handshake-to-client", [&](std::string toClientLog) {
+        NullBufferedSink nullSink;
+
+        StringSource in { toClientLog };
+        auto clientResult = ServeProto::BasicClientConnection::handshake(
+            nullSink, in, defaultVersion, "blah");
+
+        EXPECT_EQ(clientResult, defaultVersion);
+    });
+}
+
+TEST_F(ServeProtoTest, handshake_client_truncated_replay_throws)
+{
+    CharacterizationTest::readTest("handshake-to-client", [&](std::string toClientLog) {
+        for (size_t len = 0; len < toClientLog.size(); ++len) {
+            NullBufferedSink nullSink;
+            StringSource in {
+                // truncate
+                toClientLog.substr(0, len)
+            };
+            if (len < 8) {
+                EXPECT_THROW(
+                    ServeProto::BasicClientConnection::handshake(
+                        nullSink, in, defaultVersion, "blah"),
+                    EndOfFile);
+            } else {
+                // Not sure why cannot keep on checking for `EndOfFile`.
+                EXPECT_THROW(
+                    ServeProto::BasicClientConnection::handshake(
+                        nullSink, in, defaultVersion, "blah"),
+                    Error);
+            }
+        }
+    });
+}
+
+TEST_F(ServeProtoTest, handshake_client_corrupted_throws)
+{
+    CharacterizationTest::readTest("handshake-to-client", [&](const std::string toClientLog) {
+        for (size_t idx = 0; idx < toClientLog.size(); ++idx) {
+            // corrupt a copy
+            std::string toClientLogCorrupt = toClientLog;
+            toClientLogCorrupt[idx] *= 4;
+            ++toClientLogCorrupt[idx];
+
+            NullBufferedSink nullSink;
+            StringSource in { toClientLogCorrupt };
+
+            if (idx < 4 || idx == 9) {
+                // magic bytes don't match
+                EXPECT_THROW(
+                    ServeProto::BasicClientConnection::handshake(
+                        nullSink, in, defaultVersion, "blah"),
+                    Error);
+            } else if (idx < 8 || idx >= 12) {
+                // Number out of bounds
+                EXPECT_THROW(
+                    ServeProto::BasicClientConnection::handshake(
+                        nullSink, in, defaultVersion, "blah"),
+                    SerialisationError);
+            } else {
+                auto ver = ServeProto::BasicClientConnection::handshake(
+                    nullSink, in, defaultVersion, "blah");
+                EXPECT_NE(ver, defaultVersion);
+            }
+        }
+    });
+}
+
 }