Merge pull request #150 from chenxiaolong/escape

Switch to bstr for escaping mostly UTF-8 binary data

Cargo.lock

@@ -114,6 +114,7 @@ dependencies = [
  "anyhow",
  "assert_matches",
  "base64",
+ "bstr",
  "byteorder",
  "bzip2",
  "clap",
@@ -193,6 +194,17 @@ dependencies = [
  "generic-array",
 ]
 
+[[package]]
+name = "bstr"
+version = "1.6.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4c2f7349907b712260e64b0afe2f84692af14a454be26187d9df565c7f69266a"
+dependencies = [
+ "memchr",
+ "regex-automata",
+ "serde",
+]
+
 [[package]]
 name = "bumpalo"
 version = "3.13.0"

avbroot/Cargo.toml

@@ -11,6 +11,7 @@ publish = false
 [dependencies]
 anyhow = "1.0.75"
 base64 = "0.21.3"
+bstr = "1.6.2"
 byteorder = "1.4.3"
 clap = { version = "4.4.1", features = ["derive"] }
 clap_complete = "4.4.0"

avbroot/src/boot.rs

@@ -14,6 +14,7 @@ use std::{
     sync::atomic::AtomicBool,
 };
 
+use bstr::ByteSlice;
 use regex::bytes::Regex;
 use ring::digest::Context;
 use rsa::RsaPrivateKey;
@@ -34,7 +35,6 @@ use crate::{
         cpio::{self, CpioEntryNew},
     },
     stream::{self, FromReader, HashingWriter, SectionReader, ToWriter},
-    util::EscapedString,
 };
 
 #[derive(Debug, Error)]
@@ -488,8 +488,8 @@ impl BootImagePatcher for OtaCertPatcher {
         // out of future updates if the OTA certificate mechanism has changed.
         if !patched_any {
             return Err(Error::Validation(format!(
-                "No ramdisk contains {}",
-                EscapedString::new(Self::OTACERTS_PATH),
+                "No ramdisk contains {:?}",
+                Self::OTACERTS_PATH.as_bstr(),
             )));
         }
 

avbroot/src/cli/ramdisk.rs

@@ -12,12 +12,9 @@ use std::{
 use anyhow::{Context, Result};
 use clap::{Parser, Subcommand};
 
-use crate::{
-    format::{
-        compression::{CompressedFormat, CompressedReader, CompressedWriter},
-        cpio::{self, CpioEntryNew},
-    },
-    util::EscapedString,
+use crate::format::{
+    compression::{CompressedFormat, CompressedReader, CompressedWriter},
+    cpio::{self, CpioEntryNew},
 };
 
 static CONTENT_BEGIN: &str = "----- BEGIN UTF-8 CONTENT -----";
@@ -59,10 +56,10 @@ fn print_content(data: &[u8], truncate: bool) {
     println!("{BINARY_BEGIN}");
 
     if data.len() > 512 && truncate {
-        println!("{}", EscapedString::new_unquoted(&data[..512]));
+        println!("{}", data[..512].escape_ascii());
         println!("{BINARY_END_TRUNCATED}");
     } else {
-        println!("{}", EscapedString::new_unquoted(&data));
+        println!("{}", data.escape_ascii());
         println!("{BINARY_END}");
     }
 }

avbroot/src/escape.rs

@@ -0,0 +1,163 @@
+/*
+ * SPDX-FileCopyrightText: 2023 Andrew Gunnerson
+ * SPDX-License-Identifier: GPL-3.0-only
+ */
+
+use std::{fmt, marker::PhantomData};
+
+use bstr::{ByteSlice, ByteVec};
+use serde::{de::Visitor, Deserializer, Serializer};
+use thiserror::Error;
+
+#[derive(Clone, Debug, Error)]
+pub enum Error {
+    #[error("Decoded string size ({actual}) does not match expected size ({expected})")]
+    BadLength { expected: usize, actual: usize },
+}
+
+pub trait FromEscaped: Sized {
+    type Error;
+
+    fn from_escaped(data: &str) -> Result<Self, Self::Error>;
+}
+
+impl FromEscaped for Vec<u8> {
+    type Error = Error;
+
+    fn from_escaped(data: &str) -> Result<Self, Self::Error> {
+        Ok(Self::unescape_bytes(data))
+    }
+}
+
+impl<const N: usize> FromEscaped for [u8; N] {
+    type Error = Error;
+
+    fn from_escaped(data: &str) -> Result<Self, Self::Error> {
+        // Wasteful allocation, but bstr doesn't expose its decoder iterator in
+        // its public API.
+        let decoded = Vec::<u8>::from_escaped(data)?;
+        let mut buf = [0u8; N];
+
+        if decoded.len() != buf.len() {
+            return Err(Error::BadLength {
+                expected: buf.len(),
+                actual: decoded.len(),
+            });
+        }
+
+        buf.copy_from_slice(&decoded);
+
+        Ok(buf)
+    }
+}
+
+pub fn serialize<S, T>(data: T, serializer: S) -> Result<S::Ok, S::Error>
+where
+    S: Serializer,
+    T: AsRef<[u8]>,
+{
+    let s = data.as_ref().escape_bytes().to_string();
+    serializer.serialize_str(&s)
+}
+
+pub fn deserialize<'de, D, T>(deserializer: D) -> Result<T, D::Error>
+where
+    D: Deserializer<'de>,
+    T: FromEscaped,
+    <T as FromEscaped>::Error: fmt::Display,
+{
+    struct EscapedStrVisitor<T>(PhantomData<T>);
+
+    impl<'de, T> Visitor<'de> for EscapedStrVisitor<T>
+    where
+        T: FromEscaped,
+        <T as FromEscaped>::Error: fmt::Display,
+    {
+        type Value = T;
+
+        fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
+            write!(f, "an escaped string")
+        }
+
+        fn visit_str<E>(self, data: &str) -> Result<Self::Value, E>
+        where
+            E: serde::de::Error,
+        {
+            FromEscaped::from_escaped(data).map_err(serde::de::Error::custom)
+        }
+    }
+
+    deserializer.deserialize_str(EscapedStrVisitor(PhantomData))
+}
+
+#[cfg(test)]
+mod tests {
+    use assert_matches::assert_matches;
+    use serde::{Deserialize, Serialize};
+
+    use super::*;
+
+    #[test]
+    fn decode_vec() {
+        for s in ["", "abc", "你好", "💩", "\x00", "\t\r\n"] {
+            let escaped_s = s.as_bytes().escape_bytes().to_string();
+            let decoded = Vec::<u8>::from_escaped(&escaped_s).unwrap();
+
+            assert_eq!(decoded, s.as_bytes());
+        }
+    }
+
+    #[test]
+    fn decode_array() {
+        assert_matches!(
+            <[u8; 4]>::from_escaped(r"\t\r\n"),
+            Err(Error::BadLength {
+                expected: 4,
+                actual: 3,
+            })
+        );
+        assert_matches!(
+            <[u8; 4]>::from_escaped(r"\t\r\n\x00\x00"),
+            Err(Error::BadLength {
+                expected: 4,
+                actual: 5,
+            })
+        );
+        assert_matches!(
+            <[u8; 4]>::from_escaped(r"\t\r\n\x00"),
+            Ok(data) if data == *b"\t\r\n\x00"
+        );
+
+        assert_matches!(
+            <[u8; 0]>::from_escaped(r"\t"),
+            Err(Error::BadLength {
+                expected: 0,
+                actual: 1,
+            })
+        );
+        assert_matches!(
+            <[u8; 0]>::from_escaped(r""),
+            Ok(data) if data.is_empty()
+        );
+    }
+
+    #[test]
+    fn round_trip_serde() {
+        #[derive(Deserialize, Serialize)]
+        struct TestData {
+            #[serde(with = "super")]
+            contents: Vec<u8>,
+        }
+
+        let mut contents = b"foo\xffbar".to_vec();
+        contents.extend("💩".as_bytes());
+
+        let data = TestData { contents };
+        let serialized = toml_edit::ser::to_string(&data).unwrap();
+
+        assert_eq!(serialized, "contents = 'foo\\xFFbar💩'\n");
+
+        let new_data: TestData = toml_edit::de::from_str(&serialized).unwrap();
+        assert_eq!(data.contents, new_data.contents);
+    }
+}

avbroot/src/format/avb.rs

@@ -10,6 +10,7 @@ use std::{
     sync::atomic::AtomicBool,
 };
 
+use bstr::ByteSlice;
 use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
 use num_bigint_dig::{ModInverse, ToBigInt};
 use num_traits::{Pow, ToPrimitive};
@@ -28,7 +29,7 @@ use crate::{
         self, CountingReader, FromReader, ReadDiscardExt, ReadSeek, ReadStringExt, ToWriter,
         WriteStringExt, WriteZerosExt,
     },
-    util::{self, EscapedString},
+    util,
 };
 
 pub const VERSION_MAJOR: u32 = 1;
@@ -232,7 +233,7 @@ impl fmt::Debug for PropertyDescriptor {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("PropertyDescriptor")
             .field("key", &self.key)
-            .field("value", &EscapedString::new(&self.value))
+            .field("value", &self.value.as_bstr())
             .finish()
     }
 }

avbroot/src/format/cpio.rs

@@ -9,13 +9,14 @@ use std::{
     io::{self, Read, Write},
 };
 
+use bstr::ByteSlice;
 use num_traits::ToPrimitive;
 use thiserror::Error;
 
 use crate::{
     format::padding,
     stream::{CountingReader, CountingWriter, FromReader, ToWriter, WriteZerosExt},
-    util::{EscapedString, NumBytes},
+    util::NumBytes,
 };
 
 const MAGIC_NEW: &[u8; 6] = b"070701";
@@ -38,8 +39,8 @@ const IO_BLOCK_SIZE: u64 = 512;
 pub enum Error {
     #[error("Unknown magic: {0:?}")]
     UnknownMagic([u8; 6]),
-    #[error("Hard links are not supported: {0}")]
-    HardLinksNotSupported(EscapedString<Vec<u8>>),
+    #[error("Hard links are not supported: {:?}", .0.as_bstr())]
+    HardLinksNotSupported(Vec<u8>),
     #[error("{0:?} field exceeds integer bounds")]
     IntegerTooLarge(&'static str),
     #[error("I/O error")]
@@ -60,7 +61,7 @@ fn read_int(mut reader: impl Read) -> io::Result<u32> {
         let digit = c.to_digit(16).ok_or_else(|| {
             io::Error::new(
                 io::ErrorKind::InvalidData,
-                format!("{0}: Invalid hex char: {1}", EscapedString::new(&buf), c),
+                format!("{:?}: Invalid hex char: {c}", buf.as_bstr()),
             )
         })?;
 
@@ -120,7 +121,7 @@ impl fmt::Debug for CpioEntryNew {
             .field("rdev_maj", &self.rdev_maj)
             .field("rdev_min", &self.rdev_min)
             .field("chksum", &self.chksum)
-            .field("name", &EscapedString::new(&self.name))
+            .field("name", &self.name.as_bstr())
             .field("content", &NumBytes(self.content.len()))
             .finish()
     }
@@ -140,7 +141,7 @@ impl fmt::Display for CpioEntryNew {
             m => Cow::Owned(format!("unknown ({m:o})")),
         };
 
-        writeln!(f, "Filename:  {}", EscapedString::new(&self.name))?;
+        writeln!(f, "Filename:  {:?}", self.name.as_bstr())?;
         writeln!(f, "Filetype:  {file_type_str}")?;
         writeln!(f, "Inode:     {}", self.ino)?;
         writeln!(f, "Mode:      {:o}", self.mode)?;
@@ -312,7 +313,7 @@ pub fn load(mut reader: impl Read, include_trailer: bool) -> Result<Vec<CpioEntr
     loop {
         let entry = CpioEntryNew::from_reader(&mut reader)?;
         if file_type(entry.mode) != S_IFDIR && entry.nlink > 1 {
-            return Err(Error::HardLinksNotSupported(EscapedString::new(entry.name)));
+            return Err(Error::HardLinksNotSupported(entry.name));
         }
 
         if entry.name == CPIO_TRAILER {

avbroot/src/lib.rs

@@ -16,6 +16,7 @@ extern crate alloc;
 pub mod boot;
 pub mod cli;
 pub mod crypto;
+pub mod escape;
 pub mod format;
 pub mod protobuf;
 pub mod stream;

avbroot/src/stream.rs

@@ -12,10 +12,11 @@ use std::{
     },
 };
 
+use bstr::ByteSlice;
 use num_traits::ToPrimitive;
 use ring::digest::Context;
 
-use crate::util::{self, EscapedString};
+use crate::util;
 
 /// A trait for seekable readers. This is only needed because `dyn Read + Seek`
 /// is not a valid construct in Rust yet.
@@ -118,7 +119,7 @@ impl<R: Read> ReadStringExt for R {
         String::from_utf8(buf).map_err(|e| {
             io::Error::new(
                 io::ErrorKind::InvalidData,
-                format!("Invalid UTF-8: {}: {e}", EscapedString::new(e.as_bytes())),
+                format!("Invalid UTF-8: {:?}: {e}", e.as_bytes().as_bstr()),
             )
         })
     }
@@ -138,7 +139,7 @@ impl<R: Read> ReadStringExt for R {
         String::from_utf8(buf).map_err(|e| {
             io::Error::new(
                 io::ErrorKind::InvalidData,
-                format!("Invalid UTF-8: {}: {e}", EscapedString::new(e.as_bytes())),
+                format!("Invalid UTF-8: {:?}: {e}", e.as_bytes().as_bstr()),
             )
         })
     }

avbroot/src/util.rs

@@ -23,72 +23,6 @@ impl fmt::Debug for NumBytes {
     }
 }
 
-/// A wrapper around a byte slice to format it as ASCII with invalid bytes
-/// escaped as `\x##`.
-#[derive(Clone)]
-pub struct EscapedString<T: AsRef<[u8]>> {
-    inner: T,
-    quoted: bool,
-}
-
-impl<T: AsRef<[u8]>> EscapedString<T> {
-    pub fn new(inner: T) -> Self {
-        Self {
-            inner,
-            quoted: true,
-        }
-    }
-
-    pub fn new_unquoted(inner: T) -> Self {
-        Self {
-            inner,
-            quoted: false,
-        }
-    }
-
-    pub fn into_inner(self) -> T {
-        self.inner
-    }
-
-    pub fn get_ref(&self) -> &T {
-        &self.inner
-    }
-
-    pub fn get_mut(&mut self) -> &mut T {
-        &mut self.inner
-    }
-}
-
-impl<T: AsRef<[u8]>> fmt::Debug for EscapedString<T> {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        let escaped: String = self
-            .inner
-            .as_ref()
-            .iter()
-            .flat_map(|b| b.escape_ascii())
-            .map(char::from)
-            .collect();
-
-        if self.quoted {
-            write!(f, "\"")?;
-        }
-
-        write!(f, "{escaped}")?;
-
-        if self.quoted {
-            write!(f, "\"")?;
-        }
-
-        Ok(())
-    }
-}
-
-impl<T: AsRef<[u8]>> fmt::Display for EscapedString<T> {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        fmt::Debug::fmt(&self, f)
-    }
-}
-
 /// Check if a byte slice is all zeros.
 pub fn is_zero(mut buf: &[u8]) -> bool {
     while !buf.is_empty() {