Merge pull request #585 from probulate/tag-len-value

aya-log-common: support logging byte slices

.github/workflows/build-aya.yml

@@ -32,7 +32,7 @@ jobs:
 
       - uses: taiki-e/install-action@cargo-hack
       - name: Check
-        run: cargo hack check --feature-powerset --ignore-private
+        run: cargo hack check --all-targets --feature-powerset --ignore-private
 
       - uses: Swatinem/rust-cache@v1
       - name: Prereqs

aya-log-common/src/lib.rs

@@ -1,15 +1,17 @@
 #![no_std]
 
-use core::{cmp, mem, ptr, slice};
+use core::{mem, num, ptr};
 
 use num_enum::IntoPrimitive;
 
 pub const LOG_BUF_CAPACITY: usize = 8192;
 
 pub const LOG_FIELDS: usize = 6;
 
-#[repr(usize)]
-#[derive(Copy, Clone, Eq, PartialEq, Debug, Hash)]
+pub type LogValueLength = u16;
+
+#[repr(u8)]
+#[derive(Copy, Clone, Eq, PartialEq, Debug, Hash, IntoPrimitive)]
 pub enum Level {
     /// The "error" level.
     ///
@@ -33,7 +35,7 @@ pub enum Level {
     Trace,
 }
 
-#[repr(usize)]
+#[repr(u8)]
 #[derive(Copy, Clone, Debug)]
 pub enum RecordField {
     Target = 1,
@@ -46,7 +48,7 @@ pub enum RecordField {
 
 /// Types which are supported by aya-log and can be safely sent from eBPF
 /// programs to userspace.
-#[repr(usize)]
+#[repr(u8)]
 #[derive(Copy, Clone, Debug)]
 pub enum Argument {
     DisplayHint,
@@ -73,6 +75,7 @@ pub enum Argument {
     /// `[u16; 8]` array which represents an IPv6 address.
     ArrU16Len8,
 
+    Bytes,
     Str,
 }
 
@@ -96,55 +99,65 @@ pub enum DisplayHint {
     UpperMac,
 }
 
-struct TagLenValue<'a, T> {
-    tag: T,
-    value: &'a [u8],
+struct TagLenValue<T, V> {
+    pub tag: T,
+    pub value: V,
 }
 
-impl<'a, T> TagLenValue<'a, T>
+impl<T, V> TagLenValue<T, V>
 where
-    T: Copy,
+    V: IntoIterator<Item = u8>,
+    <V as IntoIterator>::IntoIter: ExactSizeIterator,
 {
-    #[inline(always)]
-    pub(crate) fn new(tag: T, value: &'a [u8]) -> TagLenValue<'a, T> {
-        TagLenValue { tag, value }
-    }
-
-    pub(crate) fn write(&self, mut buf: &mut [u8]) -> Result<usize, ()> {
-        let size = mem::size_of::<T>() + mem::size_of::<usize>() + self.value.len();
-        let remaining = cmp::min(buf.len(), LOG_BUF_CAPACITY);
-        // Check if the size doesn't exceed the buffer bounds.
-        if size > remaining {
+    pub(crate) fn write(self, mut buf: &mut [u8]) -> Result<usize, ()> {
+        // Break the abstraction to please the verifier.
+        if buf.len() > LOG_BUF_CAPACITY {
+            buf = &mut buf[..LOG_BUF_CAPACITY];
+        }
+        let Self { tag, value } = self;
+        let value = value.into_iter();
+        let len = value.len();
+        let wire_len: LogValueLength = value
+            .len()
+            .try_into()
+            .map_err(|num::TryFromIntError { .. }| ())?;
+        let size = mem::size_of_val(&tag) + mem::size_of_val(&wire_len) + len;
+        if size > buf.len() {
             return Err(());
         }
 
-        unsafe { ptr::write_unaligned(buf.as_mut_ptr() as *mut _, self.tag) };
-        buf = &mut buf[mem::size_of::<T>()..];
+        let tag_size = mem::size_of_val(&tag);
+        unsafe { ptr::write_unaligned(buf.as_mut_ptr() as *mut _, tag) };
+        buf = &mut buf[tag_size..];
 
-        unsafe { ptr::write_unaligned(buf.as_mut_ptr() as *mut _, self.value.len()) };
-        buf = &mut buf[mem::size_of::<usize>()..];
+        unsafe { ptr::write_unaligned(buf.as_mut_ptr() as *mut _, wire_len) };
+        buf = &mut buf[mem::size_of_val(&wire_len)..];
+
+        buf.iter_mut().zip(value).for_each(|(dst, src)| {
+            *dst = src;
+        });
 
-        let len = cmp::min(buf.len(), self.value.len());
-        // The verifier isn't happy with `len` being unbounded, so compare it
-        // with `LOG_BUF_CAPACITY`.
-        if len > LOG_BUF_CAPACITY {
-            return Err(());
-        }
-        buf[..len].copy_from_slice(&self.value[..len]);
         Ok(size)
     }
 }
 
+impl<T, V> TagLenValue<T, V> {
+    #[inline(always)]
+    pub(crate) fn new(tag: T, value: V) -> TagLenValue<T, V> {
+        TagLenValue { tag, value }
+    }
+}
+
 pub trait WriteToBuf {
     #[allow(clippy::result_unit_err)]
-    fn write(&self, buf: &mut [u8]) -> Result<usize, ()>;
+    fn write(self, buf: &mut [u8]) -> Result<usize, ()>;
 }
 
 macro_rules! impl_write_to_buf {
     ($type:ident, $arg_type:expr) => {
         impl WriteToBuf for $type {
-            fn write(&self, buf: &mut [u8]) -> Result<usize, ()> {
-                TagLenValue::<Argument>::new($arg_type, &self.to_ne_bytes()).write(buf)
+            fn write(self, buf: &mut [u8]) -> Result<usize, ()> {
+                TagLenValue::new($arg_type, self.to_ne_bytes()).write(buf)
             }
         }
     };
@@ -166,35 +179,40 @@ impl_write_to_buf!(f32, Argument::F32);
 impl_write_to_buf!(f64, Argument::F64);
 
 impl WriteToBuf for [u8; 16] {
-    fn write(&self, buf: &mut [u8]) -> Result<usize, ()> {
-        TagLenValue::<Argument>::new(Argument::ArrU8Len16, self).write(buf)
+    fn write(self, buf: &mut [u8]) -> Result<usize, ()> {
+        TagLenValue::new(Argument::ArrU8Len16, self).write(buf)
     }
 }
 
 impl WriteToBuf for [u16; 8] {
-    fn write(&self, buf: &mut [u8]) -> Result<usize, ()> {
-        let ptr = self.as_ptr().cast::<u8>();
-        let bytes = unsafe { slice::from_raw_parts(ptr, 16) };
-        TagLenValue::<Argument>::new(Argument::ArrU16Len8, bytes).write(buf)
+    fn write(self, buf: &mut [u8]) -> Result<usize, ()> {
+        let bytes = unsafe { core::mem::transmute::<_, [u8; 16]>(self) };
+        TagLenValue::new(Argument::ArrU16Len8, bytes).write(buf)
     }
 }
 
 impl WriteToBuf for [u8; 6] {
-    fn write(&self, buf: &mut [u8]) -> Result<usize, ()> {
-        TagLenValue::<Argument>::new(Argument::ArrU8Len6, self).write(buf)
+    fn write(self, buf: &mut [u8]) -> Result<usize, ()> {
+        TagLenValue::new(Argument::ArrU8Len6, self).write(buf)
+    }
+}
+
+impl WriteToBuf for &[u8] {
+    fn write(self, buf: &mut [u8]) -> Result<usize, ()> {
+        TagLenValue::new(Argument::Bytes, self.iter().copied()).write(buf)
     }
 }
 
-impl WriteToBuf for str {
-    fn write(&self, buf: &mut [u8]) -> Result<usize, ()> {
-        TagLenValue::<Argument>::new(Argument::Str, self.as_bytes()).write(buf)
+impl WriteToBuf for &str {
+    fn write(self, buf: &mut [u8]) -> Result<usize, ()> {
+        TagLenValue::new(Argument::Str, self.as_bytes().iter().copied()).write(buf)
     }
 }
 
 impl WriteToBuf for DisplayHint {
-    fn write(&self, buf: &mut [u8]) -> Result<usize, ()> {
-        let v: u8 = (*self).into();
-        TagLenValue::<Argument>::new(Argument::DisplayHint, &v.to_ne_bytes()).write(buf)
+    fn write(self, buf: &mut [u8]) -> Result<usize, ()> {
+        let v: u8 = self.into();
+        TagLenValue::new(Argument::DisplayHint, v.to_ne_bytes()).write(buf)
     }
 }
 
@@ -210,17 +228,31 @@ pub fn write_record_header(
     line: u32,
     num_args: usize,
 ) -> Result<usize, ()> {
+    let level: u8 = level.into();
     let mut size = 0;
-    for attr in [
-        TagLenValue::<RecordField>::new(RecordField::Target, target.as_bytes()),
-        TagLenValue::<RecordField>::new(RecordField::Level, &(level as usize).to_ne_bytes()),
-        TagLenValue::<RecordField>::new(RecordField::Module, module.as_bytes()),
-        TagLenValue::<RecordField>::new(RecordField::File, file.as_bytes()),
-        TagLenValue::<RecordField>::new(RecordField::Line, &line.to_ne_bytes()),
-        TagLenValue::<RecordField>::new(RecordField::NumArgs, &num_args.to_ne_bytes()),
-    ] {
-        size += attr.write(&mut buf[size..])?;
-    }
-
+    size += TagLenValue::new(RecordField::Target, target.as_bytes().iter().copied())
+        .write(&mut buf[size..])?;
+    size += TagLenValue::new(RecordField::Level, level.to_ne_bytes()).write(&mut buf[size..])?;
+    size += TagLenValue::new(RecordField::Module, module.as_bytes().iter().copied())
+        .write(&mut buf[size..])?;
+    size += TagLenValue::new(RecordField::File, file.as_bytes().iter().copied())
+        .write(&mut buf[size..])?;
+    size += TagLenValue::new(RecordField::Line, line.to_ne_bytes()).write(&mut buf[size..])?;
+    size +=
+        TagLenValue::new(RecordField::NumArgs, num_args.to_ne_bytes()).write(&mut buf[size..])?;
     Ok(size)
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    fn log_value_length_sufficient() {
+        assert!(
+            LOG_BUF_CAPACITY >= LogValueLength::MAX.into(),
+            "{} < {}",
+            LOG_BUF_CAPACITY,
+            LogValueLength::MAX
+        );
+    }
+}

aya-log-ebpf-macros/src/expand.rs

@@ -151,12 +151,11 @@ pub(crate) fn log(args: LogArgs, level: Option<TokenStream>) -> Result<TokenStre
                     let record_len = header_len;
 
                     if let Ok(record_len) = {
-                        use ::aya_log_ebpf::WriteToBuf;
                         Ok::<_, ()>(record_len) #( .and_then(|record_len| {
                             if record_len >= buf.buf.len() {
                                 return Err(());
                             }
-                            { #values_iter }.write(&mut buf.buf[record_len..]).map(|len| record_len + len)
+                            aya_log_ebpf::WriteToBuf::write({ #values_iter }, &mut buf.buf[record_len..]).map(|len| record_len + len)
                         }) )*
                     } {
                         unsafe { ::aya_log_ebpf::AYA_LOGS.output(