[antlir2][sendstream_parser] write a tokio codec

Summary:
The code we have to feed the streaming parser is very complicated and has some
subtle bugs around large commands / things that are split across boundaries.
In practice we usually do not read too deep into sendstreams so this doesn't
often cause a problem, but sometimes it does (S584606)

By implementing `tokio_util::codec`, we can let `tokio` take care of those
tricky boundary cases for us, and just focus on parsing bytes.

Test Plan:
```
❯ buck2 test fbcode//antlir/antlir2/sendstream_parser:
Buck UI: https://www.internalfb.com/buck2/0614533a-88ac-4ccd-baab-a093e505c07d
Tests finished: Pass 4. Fail 0. Fatal 0. Skip 0. Build failure 0
```
https://www.internalfb.com/intern/testinfra/testrun/3940649996430413

Reviewed By: wujj123456

Differential Revision: D86209569

fbshipit-source-id: 4e0d55590efb29418e1611da318003824023614f

Author: vmagro

antlir/antlir2/sendstream_parser/BUCK

@@ -22,12 +22,14 @@ rust_library(
     deps = [
         "bytes",
         "derive_more",
+        "futures",
         "hex",
         "nix",
         "nom",
         "serde",
         "thiserror",
         "tokio",
+        "tokio-util",
         "uuid",
     ],
 )

antlir/antlir2/sendstream_parser/Cargo.toml

@@ -17,12 +17,14 @@ path = "src/lib.rs"
 [dependencies]
 bytes = { version = "1.10", features = ["serde"] }
 derive_more = { version = "1.0.0", features = ["full"] }
+futures = { version = "0.3.31", features = ["async-await", "compat"] }
 hex = { version = "0.4.3", features = ["alloc"] }
 nix = { version = "0.30.1", features = ["dir", "event", "hostname", "inotify", "ioctl", "mman", "mount", "net", "poll", "ptrace", "reboot", "resource", "sched", "signal", "term", "time", "user", "zerocopy"] }
 nom = "8"
 serde = { version = "1.0.219", features = ["derive", "rc"], optional = true }
 thiserror = "2.0.12"
 tokio = { version = "1.47.1", features = ["full", "test-util", "tracing"] }
+tokio-util = { version = "0.7.15", features = ["full"] }
 uuid = { version = "1.17", features = ["rng-getrandom", "serde", "v4", "v5", "v6", "v7", "v8"] }
 
 [dev-dependencies]

antlir/antlir2/sendstream_parser/src/lib.rs

@@ -40,9 +40,9 @@ pub mod wire;
 #[derive(Debug, thiserror::Error)]
 pub enum Error {
     #[error(
-        "Sendstream had unexpected trailing data. This probably means the parser is broken: {0:?}"
+        "Sendstream had unexpected trailing data ({0} bytes). This probably means the parser is broken"
     )]
-    TrailingData(Vec<u8>),
+    TrailingData(usize),
     #[error("Sendstream is incomplete")]
     Incomplete,
     #[error("IO error: {0:#}")]

antlir/antlir2/sendstream_parser/src/wire/framed.rs

@@ -0,0 +1,65 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+use bytes::Buf;
+use bytes::BytesMut;
+use nom::IResult;
+use nom::Parser;
+use tokio_util::codec::Decoder;
+
+use crate::Command;
+use crate::Error;
+use crate::wire::NomBytes;
+pub(super) struct SendstreamDecoder;
+
+pub(super) enum Item {
+    /// Magic header that starts a sendstream - the only data here is the
+    /// sendstream version
+    SendstreamStart(#[allow(dead_code)] u32),
+    Command(Command),
+}
+
+static MAGIC_HEADER: &[u8] = b"btrfs-stream\0";
+
+/// Parse a chunk of bytes to see if we can extract the header expected atop each sendstream.
+fn sendstream_header(input: NomBytes) -> IResult<NomBytes, u32> {
+    let (remainder, (_magic, version)) = (
+        nom::bytes::streaming::tag::<&[u8], NomBytes, nom::error::Error<NomBytes>>(MAGIC_HEADER),
+        nom::number::streaming::le_u32,
+    )
+        .parse(input)?;
+    Ok((remainder, version))
+}
+
+impl Decoder for SendstreamDecoder {
+    type Item = Item;
+    type Error = Error;
+
+    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
+        // TODO: make a NomBytes for BytesMut too? This copy feels bad
+        let parsable: NomBytes = src.clone().into();
+        let starting_len = parsable.len();
+        match nom::branch::alt((
+            sendstream_header.map(Item::SendstreamStart),
+            Command::parse.map(Item::Command),
+        ))
+        .parse(parsable)
+        {
+            Ok((remaining, item)) => {
+                src.advance(starting_len - remaining.len());
+                Ok(Some(item))
+            }
+            Err(nom::Err::Incomplete(needed)) => {
+                if let nom::Needed::Size(s) = needed {
+                    src.reserve(s.into());
+                }
+                Ok(None)
+            }
+            Err(e) => Err(Error::Unparsable(e.to_string())),
+        }
+    }
+}

antlir/antlir2/sendstream_parser/src/wire/mod.rs

@@ -5,15 +5,12 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-use bytes::BytesMut;
-use nom::IResult;
-use nom::Parser as _;
+use futures::StreamExt;
 use tokio::io::AsyncRead;
-use tokio::io::AsyncReadExt;
-
-static MAGIC_HEADER: &[u8] = b"btrfs-stream\0";
+use tokio_util::codec::FramedRead;
 
 pub(crate) mod cmd;
+mod framed;
 mod nombytes;
 mod tlv;
 pub use nombytes::NomBytes;
@@ -24,16 +21,6 @@ pub enum ParserControl {
     Enough,
 }
 
-/// Parse a chunk of bytes to see if we can extract the header expected atop each sendstream.
-fn parse_header<'a>(input: &'a [u8]) -> IResult<&'a [u8], u32> {
-    let (remainder, (_magic, version)) = (
-        nom::bytes::streaming::tag::<&[u8], &[u8], nom::error::Error<&[u8]>>(MAGIC_HEADER),
-        nom::number::streaming::le_u32,
-    )
-        .parse(input)?;
-    Ok((remainder, version))
-}
-
 /// Parse an async source of bytes, expecting to find it to contain one or more sendstreams.
 /// Because the parsed commands reference data owned by the source, we do not collect the commands.
 /// Instead, we allow the caller to process them via `f`, which can instruct the processing to
@@ -46,61 +33,27 @@ fn parse_header<'a>(input: &'a [u8]) -> IResult<&'a [u8], u32> {
 /// Returns number of commands parsed.
 ///
 /// See https://btrfs.readthedocs.io/en/latest/dev/dev-send-stream.html for reference.
-pub async fn parse<R, F>(mut reader: R, mut f: F) -> crate::Result<u128>
+pub async fn parse<R, F>(reader: R, mut f: F) -> crate::Result<u128>
 where
     R: AsyncRead + Unpin + Send,
     F: FnMut(&crate::Command) -> ParserControl + Send,
 {
-    let mut unparsed = BytesMut::with_capacity(1000);
+    let mut reader = FramedRead::new(reader, framed::SendstreamDecoder);
     let mut command_count = 0;
-    let mut header: Option<u32> = None;
-    'read_bytes: loop {
-        let bytes_read = reader.read_buf(&mut unparsed).await?;
-        if bytes_read != 0 || !unparsed.is_empty() {
-            while header.is_some() {
-                match crate::Command::parse(unparsed.clone().into()) {
-                    Ok((remainder, command)) => {
-                        command_count += 1;
-                        if let ParserControl::Enough = f(&command) {
-                            // caller got what they needed, no need to continue parsing
-                            return Ok(command_count);
-                        }
-                        if let crate::Command::End = command {
-                            unparsed = remainder.into();
-                            header = None;
-                            continue 'read_bytes;
-                        }
-                        unparsed = remainder.into();
-                    }
-                    Err(nom::Err::Error(err)) | Err(nom::Err::Failure(err)) => {
-                        Err(crate::Error::Unparsable(format!("{err:?}")))?
-                    }
-                    Err(nom::Err::Incomplete(_)) => {
-                        if bytes_read == 0 {
-                            // we've found extra data that cannot be parsed w/nothing more to read
-                            Err(crate::Error::TrailingData(unparsed.clone().into()))?
-                        }
-                        continue 'read_bytes;
-                    }
+    while let Some(item_res) = reader.next().await {
+        match item_res {
+            Ok(framed::Item::Command(command)) => {
+                command_count += 1;
+                if let ParserControl::Enough = f(&command) {
+                    // caller got what they needed, no need to continue parsing
+                    break;
                 }
             }
-            match parse_header(&unparsed) {
-                Ok((remainder, _version)) => {
-                    header = Some(_version);
-                    unparsed = remainder.into();
-                }
-                Err(nom::Err::Error(err)) | Err(nom::Err::Failure(err)) => {
-                    Err(crate::Error::Unparsable(format!("{err:?}")))?
-                }
-                Err(nom::Err::Incomplete(_)) => continue,
-            };
-        } else {
-            break 'read_bytes;
+            Ok(framed::Item::SendstreamStart(_)) => {}
+            Err(e) => {
+                return Err(e);
+            }
         }
     }
-    if header.is_some() {
-        // We've found the end of the data but not the end of the last sendstream
-        Err(crate::Error::Incomplete)?;
-    }
     Ok(command_count)
 }