diff options
author | Florian Klink <flokli@flokli.de> | 2024-04-07T21·46+0300 |
---|---|---|
committer | clbot <clbot@tvl.fyi> | 2024-04-08T07·02+0000 |
commit | 71a3855f0990cbdadfa5ff109db62912e5f3e320 (patch) | |
tree | 42cd484a21689bb16416d99bbbd748b3d96e10db /tvix/nix-compat/src/wire | |
parent | acee48986699a6f0dd567b0f1176fcc32875ca45 (diff) |
feat(tvix/nix-compat/wire): have write_bytes accept AsRef<[u8]> payloads r/7873
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>
Diffstat (limited to 'tvix/nix-compat/src/wire')
-rw-r--r-- | tvix/nix-compat/src/wire/bytes.rs | 24 |
1 files changed, 20 insertions, 4 deletions
diff --git a/tvix/nix-compat/src/wire/bytes.rs b/tvix/nix-compat/src/wire/bytes.rs index 0a0930515429..d299eea6523c 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) + } } |