From 71a3855f0990cbdadfa5ff109db62912e5f3e320 Mon Sep 17 00:00:00 2001
From: Florian Klink <flokli@flokli.de>
Date: Mon, 8 Apr 2024 00:46:26 +0300
Subject: [PATCH] feat(tvix/nix-compat/wire): have write_bytes accept
 AsRef<[u8]> payloads

This includes String, &str etc. An example testcase with &str is
provided.

Change-Id: I900186d6ceb52f52bd41ef4596524c1f5b52470b
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11376
Autosubmit: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI
Reviewed-by: picnoir picnoir <picnoir@alternativebit.fr>
---
 tvix/nix-compat/src/wire/bytes.rs     | 24 ++++++++++++++++++++----
 users/picnoir/tvix-daemon/src/main.rs |  2 +-
 2 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/tvix/nix-compat/src/wire/bytes.rs b/tvix/nix-compat/src/wire/bytes.rs
index 0a0930515..d299eea65 100644
--- a/tvix/nix-compat/src/wire/bytes.rs
+++ b/tvix/nix-compat/src/wire/bytes.rs
@@ -89,20 +89,25 @@ where
 
 /// Writes a "bytes wire packet" to a (hopefully buffered) [AsyncWriteExt].
 ///
+/// Accepts anything implementing AsRef<[u8]> as payload.
+///
 /// See [read_bytes] for a description of the format.
 ///
 /// Note: if performance matters to you, make sure your
 /// [AsyncWriteExt] handle is buffered. This function is quite
 /// write-intesive.
-pub async fn write_bytes<W: AsyncWriteExt + Unpin>(w: &mut W, b: &[u8]) -> std::io::Result<()> {
+pub async fn write_bytes<W: AsyncWriteExt + Unpin, B: AsRef<[u8]>>(
+    w: &mut W,
+    b: B,
+) -> std::io::Result<()> {
     // write the size packet.
-    primitive::write_u64(w, b.len() as u64).await?;
+    primitive::write_u64(w, b.as_ref().len() as u64).await?;
 
     // write the payload
-    w.write_all(b).await?;
+    w.write_all(b.as_ref()).await?;
 
     // write padding if needed
-    let padding_len = padding_len(b.len() as u64) as usize;
+    let padding_len = padding_len(b.as_ref().len() as u64) as usize;
     if padding_len != 0 {
         w.write_all(&EMPTY_BYTES[..padding_len]).await?;
     }
@@ -209,4 +214,15 @@ mod tests {
             .build();
         assert_ok!(write_bytes(&mut mock, &input).await)
     }
+
+    #[tokio::test]
+    async fn test_write_string() {
+        let input = "Hello, World!";
+        let len = input.len() as u64;
+        let mut mock = Builder::new()
+            .write(&len.to_le_bytes())
+            .write(&hex!("48656c6c6f2c20576f726c6421000000"))
+            .build();
+        assert_ok!(write_bytes(&mut mock, &input).await)
+    }
 }
diff --git a/users/picnoir/tvix-daemon/src/main.rs b/users/picnoir/tvix-daemon/src/main.rs
index 595991b31..d18ff2471 100644
--- a/users/picnoir/tvix-daemon/src/main.rs
+++ b/users/picnoir/tvix-daemon/src/main.rs
@@ -165,7 +165,7 @@ where
             // good idea.
             debug!("write version");
             // Plain str padded to 64 bits.
-            bytes::write_bytes(&mut conn, "2.3.17".as_bytes()).await?;
+            bytes::write_bytes(&mut conn, "2.3.17").await?;
             conn.flush().await?;
         }
         if protocol_minor >= 35 {