fix(tvix/nix-compat): fix writing padding partially r/7883

We forgot to only write the part of the padding that has not been
written already.

Unfortunately it seems the tokio mocking thing does buffer small writes,
so the tests succeeds even with the bug :-/

Change-Id: I5a6cf04212d559728639427c57207094d507ec75
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11384
Tested-by: BuildkiteCI
Reviewed-by: raitobezarius <tvl@lahfa.xyz>
Reviewed-by: Brian Olsen <me@griff.name>
Autosubmit: flokli <flokli@flokli.de>

1 files changed, 28 insertions, 1 deletions

diff --git a/tvix/nix-compat/src/wire/bytes_writer.rs b/tvix/nix-compat/src/wire/bytes_writer.rs
index 5941441f44c0..933b0e6efa6c 100644
--- a/tvix/nix-compat/src/wire/bytes_writer.rs
+++ b/tvix/nix-compat/src/wire/bytes_writer.rs
@@ -190,7 +190,7 @@ where
                         let bytes_written = ensure_nonzero_bytes_written(ready!(this
                             .inner
                             .as_mut()
-                            .poll_write(cx, &EMPTY_BYTES[..total_padding_len]))?)?;
+                            .poll_write(cx, &EMPTY_BYTES[pos..total_padding_len]))?)?;
                         *this.state = BytesPacketPosition::Padding(pos + bytes_written);
                     } else {
                         // everything written, break
@@ -235,6 +235,8 @@ where
 
 #[cfg(test)]
 mod tests {
+    use std::time::Duration;
+
     use crate::wire::bytes::write_bytes;
     use hex_literal::hex;
     use lazy_static::lazy_static;
@@ -358,6 +360,31 @@ mod tests {
         assert_ok!(w.shutdown().await);
     }
 
+    /// Write a 9 bytes packet, but cause the sink to only accept half of the
+    /// padding, ensuring we correctly write (only) the rest of the padding later.
+    /// We write another 2 bytes of "bait", where a faulty implementation (pre
+    /// cl/11384) would put too many null bytes.
+    #[tokio::test]
+    async fn write_9b_write_padding_2steps() {
+        let payload = &hex!("000102030405060708");
+        let exp_bytes = produce_exp_bytes(payload).await;
+
+        let mut mock = Builder::new()
+            .write(&exp_bytes[0..8]) // size
+            .write(&exp_bytes[8..17]) // payload
+            .write(&exp_bytes[17..19]) // padding (2 of 7 bytes)
+            // insert a wait to prevent Mock from merging the two writes into one
+            .wait(Duration::from_nanos(1))
+            .write(&hex!("0000000000ffff")) // padding (5 of 7 bytes, plus 2 bytes of "bait")
+            .build();
+
+        let mut w = BytesWriter::new(&mut mock, payload.len() as u64);
+        assert_ok!(w.write_all(&payload[..]).await);
+        assert_ok!(w.flush().await);
+        // Write bait
+        assert_ok!(mock.write_all(&hex!("ffff")).await);
+    }
+
     /// Write a larger bytes packet
     #[tokio::test]
     async fn write_1m() {