From eeb89c28b026ed96baf2479491c071e527b92b58 Mon Sep 17 00:00:00 2001
From: John Ericson <John.Ericson@Obsidian.Systems>
Date: Mon, 27 May 2024 00:18:58 -0400
Subject: [PATCH] Worker proto use proper serialiser for `BuildMode`

Do this instead of an unchecked cast

I redid this to use the serialisation framework (including a unit test),
but I am keeping the reference to credit Jade for spotting the issue.

Change-Id: Icf6af7935e8f139bef36b40ad475e973aa48855c
(adapted from commit 2a7a824d83dc5fb33326b8b89625685f283a743b)

Co-Authored-By: Jade Lovelace <lix@jade.fyi>
---
 src/libstore/daemon.cc                        |   6 ++--
 src/libstore/store-api.hh                     |   2 +-
 src/libstore/worker-protocol.cc               |  28 ++++++++++++++++++
 src/libstore/worker-protocol.hh               |   3 ++
 .../data/worker-protocol/build-mode.bin       | Bin 0 -> 24 bytes
 tests/unit/libstore/worker-protocol.cc        |  11 +++++++
 6 files changed, 46 insertions(+), 4 deletions(-)
 create mode 100644 tests/unit/libstore/data/worker-protocol/build-mode.bin

diff --git a/src/libstore/daemon.cc b/src/libstore/daemon.cc
index 47d6d5541..eb6a4e690 100644
--- a/src/libstore/daemon.cc
+++ b/src/libstore/daemon.cc
@@ -531,7 +531,7 @@ static void performOp(TunnelLogger * logger, ref<Store> store,
         auto drvs = WorkerProto::Serialise<DerivedPaths>::read(*store, rconn);
         BuildMode mode = bmNormal;
         if (GET_PROTOCOL_MINOR(clientVersion) >= 15) {
-            mode = (BuildMode) readInt(from);
+            mode = WorkerProto::Serialise<BuildMode>::read(*store, rconn);
 
             /* Repairing is not atomic, so disallowed for "untrusted"
                clients.
@@ -555,7 +555,7 @@ static void performOp(TunnelLogger * logger, ref<Store> store,
     case WorkerProto::Op::BuildPathsWithResults: {
         auto drvs = WorkerProto::Serialise<DerivedPaths>::read(*store, rconn);
         BuildMode mode = bmNormal;
-        mode = (BuildMode) readInt(from);
+        mode = WorkerProto::Serialise<BuildMode>::read(*store, rconn);
 
         /* Repairing is not atomic, so disallowed for "untrusted"
            clients.
@@ -586,7 +586,7 @@ static void performOp(TunnelLogger * logger, ref<Store> store,
          * correctly.
          */
         readDerivation(from, *store, drv, Derivation::nameFromPath(drvPath));
-        BuildMode buildMode = (BuildMode) readInt(from);
+        auto buildMode = WorkerProto::Serialise<BuildMode>::read(*store, rconn);
         logger->startWork();
 
         auto drvType = drv.type();
diff --git a/src/libstore/store-api.hh b/src/libstore/store-api.hh
index 430d9a5ab..15712458c 100644
--- a/src/libstore/store-api.hh
+++ b/src/libstore/store-api.hh
@@ -92,7 +92,7 @@ enum SubstituteFlag : bool { NoSubstitute = false, Substitute = true };
 const uint32_t exportMagic = 0x4558494e;
 
 
-enum BuildMode { bmNormal, bmRepair, bmCheck };
+enum BuildMode : uint8_t { bmNormal, bmRepair, bmCheck };
 enum TrustedFlag : bool { NotTrusted = false, Trusted = true };
 
 struct BuildResult;
diff --git a/src/libstore/worker-protocol.cc b/src/libstore/worker-protocol.cc
index a50259d24..ed6367c38 100644
--- a/src/libstore/worker-protocol.cc
+++ b/src/libstore/worker-protocol.cc
@@ -14,6 +14,34 @@ namespace nix {
 
 /* protocol-specific definitions */
 
+BuildMode WorkerProto::Serialise<BuildMode>::read(const StoreDirConfig & store, WorkerProto::ReadConn conn)
+{
+    auto temp = readNum<uint8_t>(conn.from);
+    switch (temp) {
+    case bmNormal: return bmNormal;
+    case bmRepair: return bmRepair;
+    case bmCheck: return bmCheck;
+    default: throw Error("Invalid build mode");
+    }
+}
+
+void WorkerProto::Serialise<BuildMode>::write(const StoreDirConfig & store, WorkerProto::WriteConn conn, const BuildMode & buildMode)
+{
+    switch (buildMode) {
+    case bmNormal:
+        conn.to << uint8_t{bmNormal};
+        break;
+    case bmRepair:
+        conn.to << uint8_t{bmRepair};
+        break;
+    case bmCheck:
+        conn.to << uint8_t{bmCheck};
+        break;
+    default:
+        assert(false);
+    };
+}
+
 std::optional<TrustedFlag> WorkerProto::Serialise<std::optional<TrustedFlag>>::read(const StoreDirConfig & store, WorkerProto::ReadConn conn)
 {
     auto temp = readNum<uint8_t>(conn.from);
diff --git a/src/libstore/worker-protocol.hh b/src/libstore/worker-protocol.hh
index 91d277b77..34607b7af 100644
--- a/src/libstore/worker-protocol.hh
+++ b/src/libstore/worker-protocol.hh
@@ -35,6 +35,7 @@ struct BuildResult;
 struct KeyedBuildResult;
 struct ValidPathInfo;
 struct UnkeyedValidPathInfo;
+enum BuildMode : uint8_t;
 enum TrustedFlag : bool;
 
 
@@ -215,6 +216,8 @@ DECLARE_WORKER_SERIALISER(ValidPathInfo);
 template<>
 DECLARE_WORKER_SERIALISER(UnkeyedValidPathInfo);
 template<>
+DECLARE_WORKER_SERIALISER(BuildMode);
+template<>
 DECLARE_WORKER_SERIALISER(std::optional<TrustedFlag>);
 template<>
 DECLARE_WORKER_SERIALISER(std::optional<std::chrono::microseconds>);
diff --git a/tests/unit/libstore/data/worker-protocol/build-mode.bin b/tests/unit/libstore/data/worker-protocol/build-mode.bin
new file mode 100644
index 0000000000000000000000000000000000000000..51b239409bc815c19b00cb97f62996fb782f24c3
GIT binary patch
literal 24
PcmZQzfB;4)%><<Z05|{y

literal 0
HcmV?d00001

diff --git a/tests/unit/libstore/worker-protocol.cc b/tests/unit/libstore/worker-protocol.cc
index 2b2e559a9..89c3fe241 100644
--- a/tests/unit/libstore/worker-protocol.cc
+++ b/tests/unit/libstore/worker-protocol.cc
@@ -529,6 +529,17 @@ VERSIONED_CHARACTERIZATION_TEST(
         }),
     }))
 
+VERSIONED_CHARACTERIZATION_TEST(
+    WorkerProtoTest,
+    buildMode,
+    "build-mode",
+    defaultVersion,
+    (std::tuple<BuildMode, BuildMode, BuildMode> {
+        bmNormal,
+        bmRepair,
+        bmCheck,
+    }))
+
 VERSIONED_CHARACTERIZATION_TEST(
     WorkerProtoTest,
     optionalTrustedFlag,