feat(nix-compat/wire): add low-level wire format primitives code r/7695

This brings some initial Nix wire format parsing code, used in the nix
daemon protocol, remote store/builder protocol, as well as the NAR
format itself (note we already have more specialized code for the last
one).

Thanks to embr, this code already exists, in
https://codeberg.org/gorgon/gorgon/src/branch/main/nix-daemon/src/wire.rs,
and we can vendor it into here, as EUPL is compatible with GPL (in that
direction).

The code uses the tokio::io Reader and Writer traits, not the ones from
the `futures` crate, as they provide some more convenient `read_u64_le`
functions.

More application-specific parsing code, as well as code to read strings,
or bytes are left out for now, as we want to be be more restrictive
w.r.t allowed max sizes, and need to parse bytes, not strings.

The code slightly diverges, as we have clippy looped into CI.
`Ok(…?)` can be turned into just the inner expression, and
some .and_then can be expressed in a simpler fashion.

Change-Id: Ie3adcb485e9d66786673b1962a08d4e5df3781d9
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11148
Autosubmit: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI
Reviewed-by: picnoir picnoir <picnoir@alternativebit.fr>

6 files changed, 157 insertions, 4 deletions

diff --git a/tvix/Cargo.lock b/tvix/Cargo.lock
index 9f59ebe54f73..a10b4b35fa4d 100644
--- a/tvix/Cargo.lock
+++ b/tvix/Cargo.lock
@@ -1723,6 +1723,8 @@ dependencies = [
  "test-case",
  "test-generator",
  "thiserror",
+ "tokio",
+ "tokio-test",
  "zstd",
 ]
 
@@ -3271,6 +3273,19 @@ dependencies = [
 ]
 
 [[package]]
+name = "tokio-test"
+version = "0.4.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e89b3cbabd3ae862100094ae433e1def582cf86451b4e9bf83aa7ac1d8a7d719"
+dependencies = [
+ "async-stream",
+ "bytes",
+ "futures-core",
+ "tokio",
+ "tokio-stream",
+]
+
+[[package]]
 name = "tokio-util"
 version = "0.7.10"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/tvix/Cargo.nix b/tvix/Cargo.nix
index 6fc42a71b59f..258461ffe427 100644
--- a/tvix/Cargo.nix
+++ b/tvix/Cargo.nix
@@ -5149,6 +5149,12 @@ rec {
             name = "thiserror";
             packageId = "thiserror";
           }
+          {
+            name = "tokio";
+            packageId = "tokio";
+            optional = true;
+            features = [ "io-util" "macros" ];
+          }
         ];
         devDependencies = [
           {
@@ -5187,15 +5193,20 @@ rec {
             packageId = "test-generator";
           }
           {
+            name = "tokio-test";
+            packageId = "tokio-test";
+          }
+          {
             name = "zstd";
             packageId = "zstd";
           }
         ];
         features = {
-          "async" = [ "futures-util" ];
+          "async" = [ "futures-util" "tokio" ];
           "futures-util" = [ "dep:futures-util" ];
+          "tokio" = [ "dep:tokio" ];
         };
-        resolvedDefaultFeatures = [ "async" "futures-util" ];
+        resolvedDefaultFeatures = [ "async" "futures-util" "tokio" ];
       };
       "nom" = rec {
         crateName = "nom";
@@ -9687,7 +9698,7 @@ rec {
           "tracing" = [ "dep:tracing" ];
           "windows-sys" = [ "dep:windows-sys" ];
         };
-        resolvedDefaultFeatures = [ "bytes" "default" "fs" "io-std" "io-util" "libc" "macros" "mio" "net" "num_cpus" "rt" "rt-multi-thread" "signal" "signal-hook-registry" "socket2" "sync" "time" "tokio-macros" "windows-sys" ];
+        resolvedDefaultFeatures = [ "bytes" "default" "fs" "io-std" "io-util" "libc" "macros" "mio" "net" "num_cpus" "rt" "rt-multi-thread" "signal" "signal-hook-registry" "socket2" "sync" "test-util" "time" "tokio-macros" "windows-sys" ];
       };
       "tokio-io-timeout" = rec {
         crateName = "tokio-io-timeout";
@@ -9995,6 +10006,46 @@ rec {
         };
         resolvedDefaultFeatures = [ "default" "xattr" ];
       };
+      "tokio-test" = rec {
+        crateName = "tokio-test";
+        version = "0.4.3";
+        edition = "2021";
+        sha256 = "06fplzcc2ymahfzykd2ickw2qn7g3lz47bll00865s1spnx3r6z8";
+        authors = [
+          "Tokio Contributors <team@tokio.rs>"
+        ];
+        dependencies = [
+          {
+            name = "async-stream";
+            packageId = "async-stream";
+          }
+          {
+            name = "bytes";
+            packageId = "bytes";
+          }
+          {
+            name = "futures-core";
+            packageId = "futures-core";
+          }
+          {
+            name = "tokio";
+            packageId = "tokio";
+            features = [ "rt" "sync" "time" "test-util" ];
+          }
+          {
+            name = "tokio-stream";
+            packageId = "tokio-stream";
+          }
+        ];
+        devDependencies = [
+          {
+            name = "tokio";
+            packageId = "tokio";
+            features = [ "full" ];
+          }
+        ];
+
+      };
       "tokio-util" = rec {
         crateName = "tokio-util";
         version = "0.7.10";
diff --git a/tvix/nix-compat/Cargo.toml b/tvix/nix-compat/Cargo.toml
index c4672ff9fb9e..181eb9428915 100644
--- a/tvix/nix-compat/Cargo.toml
+++ b/tvix/nix-compat/Cargo.toml
@@ -6,7 +6,7 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [features]
-async = ["futures-util"]
+async = ["futures-util", "tokio"]
 
 [dependencies]
 bitflags = "2.4.1"
@@ -22,6 +22,11 @@ serde_json = "1.0"
 sha2 = "0.10.6"
 thiserror = "1.0.38"
 
+[dependencies.tokio]
+optional = true
+version = "1.32.0"
+features = ["io-util", "macros"]
+
 [dev-dependencies]
 futures = { version = "0.3.30", default-features = false, features = ["executor"] }
 lazy_static = "1.4.0"
@@ -30,6 +35,7 @@ test-case = "3.3.1"
 criterion = { version = "0.5", features = ["html_reports"] }
 hex-literal = "0.4.1"
 pretty_assertions = "1.4.0"
+tokio-test = "0.4.3"
 zstd = "^0.13.0"
 
 [dev-dependencies.test-generator]
diff --git a/tvix/nix-compat/src/lib.rs b/tvix/nix-compat/src/lib.rs
index dd161cc1f944..60dcbdf25caa 100644
--- a/tvix/nix-compat/src/lib.rs
+++ b/tvix/nix-compat/src/lib.rs
@@ -5,3 +5,4 @@ pub mod narinfo;
 pub mod nixbase32;
 pub mod nixhash;
 pub mod store_path;
+mod wire;
diff --git a/tvix/nix-compat/src/wire/mod.rs b/tvix/nix-compat/src/wire/mod.rs
new file mode 100644
index 000000000000..e0b184c78aec
--- /dev/null
+++ b/tvix/nix-compat/src/wire/mod.rs
@@ -0,0 +1,5 @@
+//! Module parsing and emitting the wire format used by Nix, both in the
+//! nix-daemon protocol as well as in the NAR format.
+
+#[cfg(feature = "async")]
+pub mod primitive;
diff --git a/tvix/nix-compat/src/wire/primitive.rs b/tvix/nix-compat/src/wire/primitive.rs
new file mode 100644
index 000000000000..54f00b15a0c0
--- /dev/null
+++ b/tvix/nix-compat/src/wire/primitive.rs
@@ -0,0 +1,75 @@
+// SPDX-FileCopyrightText: 2023 embr <git@liclac.eu>
+//
+// SPDX-License-Identifier: EUPL-1.2
+
+use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
+
+#[allow(dead_code)]
+/// Read a u64 from the stream (little endian).
+pub async fn read_u64<R: AsyncReadExt + Unpin>(r: &mut R) -> std::io::Result<u64> {
+    r.read_u64_le().await
+}
+
+#[allow(dead_code)]
+/// Write a u64 from the stream (little endian).
+pub async fn write_u64<W: AsyncWrite + Unpin>(w: &mut W, v: u64) -> std::io::Result<()> {
+    w.write_u64_le(v).await
+}
+
+#[allow(dead_code)]
+/// Read a boolean from the stream, encoded as u64 (>0 is true).
+pub async fn read_bool<R: AsyncRead + Unpin>(r: &mut R) -> std::io::Result<bool> {
+    Ok(read_u64(r).await? > 0)
+}
+
+#[allow(dead_code)]
+/// Write a boolean to the stream, encoded as u64 (>0 is true).
+pub async fn write_bool<W: AsyncWrite + Unpin>(w: &mut W, v: bool) -> std::io::Result<()> {
+    write_u64(w, if v { 1u64 } else { 0u64 }).await
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use tokio_test::io::Builder;
+
+    // Integers.
+    #[tokio::test]
+    async fn test_read_u64() {
+        let mut mock = Builder::new().read(&1234567890u64.to_le_bytes()).build();
+        assert_eq!(1234567890u64, read_u64(&mut mock).await.unwrap());
+    }
+    #[tokio::test]
+    async fn test_write_u64() {
+        let mut mock = Builder::new().write(&1234567890u64.to_le_bytes()).build();
+        write_u64(&mut mock, 1234567890).await.unwrap();
+    }
+
+    // Booleans.
+    #[tokio::test]
+    async fn test_read_bool_0() {
+        let mut mock = Builder::new().read(&0u64.to_le_bytes()).build();
+        assert!(!read_bool(&mut mock).await.unwrap());
+    }
+    #[tokio::test]
+    async fn test_read_bool_1() {
+        let mut mock = Builder::new().read(&1u64.to_le_bytes()).build();
+        assert!(read_bool(&mut mock).await.unwrap());
+    }
+    #[tokio::test]
+    async fn test_read_bool_2() {
+        let mut mock = Builder::new().read(&2u64.to_le_bytes()).build();
+        assert!(read_bool(&mut mock).await.unwrap());
+    }
+
+    #[tokio::test]
+    async fn test_write_bool_false() {
+        let mut mock = Builder::new().write(&0u64.to_le_bytes()).build();
+        write_bool(&mut mock, false).await.unwrap();
+    }
+    #[tokio::test]
+    async fn test_write_bool_true() {
+        let mut mock = Builder::new().write(&1u64.to_le_bytes()).build();
+        write_bool(&mut mock, true).await.unwrap();
+    }
+}