aya-log: Allow logging `core::net::Ipv4Addr` and `core::net::Ipv6Addr`

IP address types are available in `core`, so they can be used also in
eBPF programs. This change adds support of these types in aya-log.

* Add implementation of `WriteTuBuf` to these types.
* Support these types in `Ipv4Formatter` and `Ipv6Formatter`.
* Support them with `DisplayHint::Ip`.
* Add support for formatting `[u8; 4]`, to be able to handle
  `Ipv4Addr::octets`.

aya-log-common/src/lib.rs

@@ -1,6 +1,9 @@
 #![no_std]
 
-use core::num::{NonZeroUsize, TryFromIntError};
+use core::{
+    net::{IpAddr, Ipv4Addr, Ipv6Addr},
+    num::{NonZeroUsize, TryFromIntError},
+};
 
 use num_enum::IntoPrimitive;
 
@@ -52,7 +55,8 @@ impl_formatter_for_types!(
         f32, f64,
         char,
         str,
-        &str
+        &str,
+        IpAddr, Ipv4Addr, Ipv6Addr
     }
 );
 
@@ -75,7 +79,11 @@ impl_formatter_for_types!(
 );
 
 pub trait IpFormatter {}
+impl IpFormatter for IpAddr {}
+impl IpFormatter for Ipv4Addr {}
+impl IpFormatter for Ipv6Addr {}
 impl IpFormatter for u32 {}
+impl IpFormatter for [u8; 4] {}
 impl IpFormatter for [u8; 16] {}
 impl IpFormatter for [u16; 8] {}
 
@@ -118,6 +126,11 @@ pub enum Argument {
     F32,
     F64,
 
+    Ipv4Addr,
+    Ipv6Addr,
+
+    /// `[u8; 4]` array which represents an IPv4 address.
+    ArrU8Len4,
     /// `[u8; 6]` array which represents a MAC address.
     ArrU8Len6,
     /// `[u8; 16]` array which represents an IPv6 address.
@@ -203,6 +216,36 @@ impl_write_to_buf!(usize, Argument::Usize);
 impl_write_to_buf!(f32, Argument::F32);
 impl_write_to_buf!(f64, Argument::F64);
 
+impl WriteToBuf for IpAddr {
+    fn write(self, buf: &mut [u8]) -> Option<NonZeroUsize> {
+        match self {
+            IpAddr::V4(ipv4_addr) => write(Argument::Ipv4Addr.into(), &ipv4_addr.octets(), buf),
+            IpAddr::V6(ipv6_addr) => write(Argument::Ipv6Addr.into(), &ipv6_addr.octets(), buf),
+        }
+    }
+}
+
+impl WriteToBuf for Ipv4Addr {
+    fn write(self, buf: &mut [u8]) -> Option<NonZeroUsize> {
+        write(Argument::Ipv4Addr.into(), &self.octets(), buf)
+    }
+}
+
+impl WriteToBuf for [u8; 4] {
+    // This need not be inlined because the return value is Option<N> where N is 16, which is a
+    // compile-time constant.
+    #[inline(never)]
+    fn write(self, buf: &mut [u8]) -> Option<NonZeroUsize> {
+        write(Argument::ArrU8Len4.into(), &self, buf)
+    }
+}
+
+impl WriteToBuf for Ipv6Addr {
+    fn write(self, buf: &mut [u8]) -> Option<NonZeroUsize> {
+        write(Argument::Ipv6Addr.into(), &self.octets(), buf)
+    }
+}
+
 impl WriteToBuf for [u8; 16] {
     // This need not be inlined because the return value is Option<N> where N is 16, which is a
     // compile-time constant.

aya-log/src/lib.rs

@@ -310,6 +310,48 @@ impl Format for u32 {
     }
 }
 
+impl Format for Ipv4Addr {
+    fn format(&self, last_hint: Option<DisplayHintWrapper>) -> Result<String, ()> {
+        match last_hint.map(|DisplayHintWrapper(dh)| dh) {
+            Some(DisplayHint::Default) => Ok(Ipv4Formatter::format(*self)),
+            Some(DisplayHint::LowerHex) => Err(()),
+            Some(DisplayHint::UpperHex) => Err(()),
+            Some(DisplayHint::Ip) => Ok(Ipv4Formatter::format(*self)),
+            Some(DisplayHint::LowerMac) => Err(()),
+            Some(DisplayHint::UpperMac) => Err(()),
+            None => Ok(Ipv4Formatter::format(*self)),
+        }
+    }
+}
+
+impl Format for Ipv6Addr {
+    fn format(&self, last_hint: Option<DisplayHintWrapper>) -> Result<String, ()> {
+        match last_hint.map(|DisplayHintWrapper(dh)| dh) {
+            Some(DisplayHint::Default) => Ok(Ipv6Formatter::format(*self)),
+            Some(DisplayHint::LowerHex) => Err(()),
+            Some(DisplayHint::UpperHex) => Err(()),
+            Some(DisplayHint::Ip) => Ok(Ipv6Formatter::format(*self)),
+            Some(DisplayHint::LowerMac) => Err(()),
+            Some(DisplayHint::UpperMac) => Err(()),
+            None => Ok(Ipv6Formatter::format(*self)),
+        }
+    }
+}
+
+impl Format for [u8; 4] {
+    fn format(&self, last_hint: Option<DisplayHintWrapper>) -> Result<String, ()> {
+        match last_hint.map(|DisplayHintWrapper(dh)| dh) {
+            Some(DisplayHint::Default) => Ok(Ipv4Formatter::format(*self)),
+            Some(DisplayHint::LowerHex) => Err(()),
+            Some(DisplayHint::UpperHex) => Err(()),
+            Some(DisplayHint::Ip) => Ok(Ipv4Formatter::format(*self)),
+            Some(DisplayHint::LowerMac) => Err(()),
+            Some(DisplayHint::UpperMac) => Err(()),
+            None => Ok(Ipv4Formatter::format(*self)),
+        }
+    }
+}
+
 impl Format for [u8; 6] {
     fn format(&self, last_hint: Option<DisplayHintWrapper>) -> Result<String, ()> {
         match last_hint.map(|DisplayHintWrapper(dh)| dh) {
@@ -548,6 +590,20 @@ fn log_buf(mut buf: &[u8], logger: &dyn Log) -> Result<(), ()> {
                         .format(last_hint.take())?,
                 );
             }
+            Argument::Ipv4Addr => {
+                let value: [u8; 4] = value.try_into().map_err(|_| ())?;
+                let value = Ipv4Addr::from(value);
+                full_log_msg.push_str(&value.format(last_hint.take())?)
+            }
+            Argument::Ipv6Addr => {
+                let value: [u8; 16] = value.try_into().map_err(|_| ())?;
+                let value = Ipv6Addr::from(value);
+                full_log_msg.push_str(&value.format(last_hint.take())?)
+            }
+            Argument::ArrU8Len4 => {
+                let value: [u8; 4] = value.try_into().map_err(|_| ())?;
+                full_log_msg.push_str(&value.format(last_hint.take())?);
+            }
             Argument::ArrU8Len6 => {
                 let value: [u8; 6] = value.try_into().map_err(|_| ())?;
                 full_log_msg.push_str(&value.format(last_hint.take())?);
@@ -615,6 +671,8 @@ fn try_read<T: Pod>(mut buf: &[u8]) -> Result<(T, &[u8], &[u8]), ()> {
 
 #[cfg(test)]
 mod test {
+    use std::net::IpAddr;
+
     use aya_log_common::{write_record_header, WriteToBuf};
     use log::{logger, Level};
 
@@ -794,6 +852,52 @@ mod test {
         testing_logger::setup();
         let (mut len, mut input) = new_log(3).unwrap();
 
+        len += "ipv4: ".write(&mut input[len..]).unwrap().get();
+        len += DisplayHint::Ip.write(&mut input[len..]).unwrap().get();
+        len += Ipv4Addr::new(10, 0, 0, 1)
+            .write(&mut input[len..])
+            .unwrap()
+            .get();
+
+        _ = len;
+
+        let logger = logger();
+        let () = log_buf(&input, logger).unwrap();
+        testing_logger::validate(|captured_logs| {
+            assert_eq!(captured_logs.len(), 1);
+            assert_eq!(captured_logs[0].body, "ipv4: 10.0.0.1");
+            assert_eq!(captured_logs[0].level, Level::Info);
+        });
+    }
+
+    #[test]
+    fn test_display_hint_ip_ipv4() {
+        testing_logger::setup();
+        let (mut len, mut input) = new_log(3).unwrap();
+
+        len += "ipv4: ".write(&mut input[len..]).unwrap().get();
+        len += DisplayHint::Ip.write(&mut input[len..]).unwrap().get();
+        len += IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1))
+            .write(&mut input[len..])
+            .unwrap()
+            .get();
+
+        _ = len;
+
+        let logger = logger();
+        let () = log_buf(&input, logger).unwrap();
+        testing_logger::validate(|captured_logs| {
+            assert_eq!(captured_logs.len(), 1);
+            assert_eq!(captured_logs[0].body, "ipv4: 10.0.0.1");
+            assert_eq!(captured_logs[0].level, Level::Info);
+        });
+    }
+
+    #[test]
+    fn test_display_hint_ipv4_u32() {
+        testing_logger::setup();
+        let (mut len, mut input) = new_log(3).unwrap();
+
         len += "ipv4: ".write(&mut input[len..]).unwrap().get();
         len += DisplayHint::Ip.write(&mut input[len..]).unwrap().get();
         // 10.0.0.1 as u32
@@ -810,6 +914,56 @@ mod test {
         });
     }
 
+    #[test]
+    fn test_display_hint_ipv6() {
+        testing_logger::setup();
+        let (mut len, mut input) = new_log(3).unwrap();
+
+        len += "ipv6: ".write(&mut input[len..]).unwrap().get();
+        len += DisplayHint::Ip.write(&mut input[len..]).unwrap().get();
+        len += Ipv6Addr::new(
+            0x2001, 0x0db8, 0x0000, 0x0000, 0x0000, 0x0000, 0x0001, 0x0001,
+        )
+        .write(&mut input[len..])
+        .unwrap()
+        .get();
+
+        _ = len;
+
+        let logger = logger();
+        let () = log_buf(&input, logger).unwrap();
+        testing_logger::validate(|captured_logs| {
+            assert_eq!(captured_logs.len(), 1);
+            assert_eq!(captured_logs[0].body, "ipv6: 2001:db8::1:1");
+            assert_eq!(captured_logs[0].level, Level::Info);
+        });
+    }
+
+    #[test]
+    fn test_display_hint_ip_ipv6() {
+        testing_logger::setup();
+        let (mut len, mut input) = new_log(3).unwrap();
+
+        len += "ipv6: ".write(&mut input[len..]).unwrap().get();
+        len += DisplayHint::Ip.write(&mut input[len..]).unwrap().get();
+        len += IpAddr::V6(Ipv6Addr::new(
+            0x2001, 0x0db8, 0x0000, 0x0000, 0x0000, 0x0000, 0x0001, 0x0001,
+        ))
+        .write(&mut input[len..])
+        .unwrap()
+        .get();
+
+        _ = len;
+
+        let logger = logger();
+        let () = log_buf(&input, logger).unwrap();
+        testing_logger::validate(|captured_logs| {
+            assert_eq!(captured_logs.len(), 1);
+            assert_eq!(captured_logs[0].body, "ipv6: 2001:db8::1:1");
+            assert_eq!(captured_logs[0].level, Level::Info);
+        });
+    }
+
     #[test]
     fn test_display_hint_ipv6_arr_u8_len_16() {
         testing_logger::setup();

test/integration-ebpf/src/log.rs

@@ -1,6 +1,8 @@
 #![no_std]
 #![no_main]
 
+use core::net::{IpAddr, Ipv4Addr, Ipv6Addr};
+
 use aya_ebpf::{macros::uprobe, programs::ProbeContext};
 use aya_log_ebpf::{debug, error, info, trace, warn};
 
@@ -15,11 +17,43 @@ pub fn test_log(ctx: ProbeContext) {
         "wao",
         "wao".as_bytes()
     );
-    let ipv4 = 167772161u32; // 10.0.0.1
-    let ipv6 = [
-        32u8, 1u8, 13u8, 184u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 1u8,
-    ]; // 2001:db8::1
-    info!(&ctx, "ipv4: {:i}, ipv6: {:i}", ipv4, ipv6);
+
+    // 10.0.0.1
+    let ipv4 = Ipv4Addr::new(10, 0, 0, 1);
+    // 2001:db8::1
+    let ipv6 = Ipv6Addr::new(8193, 3512, 0, 0, 0, 0, 0, 1);
+    info!(
+        &ctx,
+        "ip structs, without format hint: ipv4: {}, ipv6: {}", ipv4, ipv6
+    );
+    info!(
+        &ctx,
+        "ip structs, with format hint: ipv4: {:i}, ipv6: {:i}", ipv4, ipv6
+    );
+
+    let ipv4_enum = IpAddr::V4(ipv4);
+    let ipv6_enum = IpAddr::V6(ipv6);
+    info!(
+        &ctx,
+        "ip enums, without format hint: ipv4: {}, ipv6: {}", ipv4_enum, ipv6_enum
+    );
+    info!(
+        &ctx,
+        "ip enums, with format hint: ipv4: {:i}, ipv6: {:i}", ipv4_enum, ipv6_enum
+    );
+
+    // We don't format `Ipv6Addr::to_bits`, because `u128` is not supported by
+    // eBPF. Even though Rust compiler does not complain, verifier would throw
+    // an error about returning values not fitting into 64-bit registers.
+    info!(&ctx, "ip as bits: ipv4: {:i}", ipv4.to_bits());
+
+    info!(
+        &ctx,
+        "ip as octets: ipv4: {:i}, ipv6: {:i}",
+        ipv4.octets(),
+        ipv6.octets()
+    );
+
     let mac = [4u8, 32u8, 6u8, 9u8, 0u8, 64u8];
     trace!(&ctx, "mac lc: {:mac}, mac uc: {:MAC}", mac, mac);
     let hex = 0x2f;

test/integration-test/src/tests/log.rs

@@ -106,7 +106,52 @@ async fn log() {
     assert_eq!(
         records.next(),
         Some(&CapturedLog {
-            body: "ipv4: 10.0.0.1, ipv6: 2001:db8::1".into(),
+            body: "ip structs, without format hint: ipv4: 10.0.0.1, ipv6: 2001:db8::1".into(),
+            level: Level::Info,
+            target: "log".into(),
+        })
+    );
+
+    assert_eq!(
+        records.next(),
+        Some(&CapturedLog {
+            body: "ip structs, with format hint: ipv4: 10.0.0.1, ipv6: 2001:db8::1".into(),
+            level: Level::Info,
+            target: "log".into(),
+        })
+    );
+
+    assert_eq!(
+        records.next(),
+        Some(&CapturedLog {
+            body: "ip enums, without format hint: ipv4: 10.0.0.1, ipv6: 2001:db8::1".into(),
+            level: Level::Info,
+            target: "log".into(),
+        })
+    );
+
+    assert_eq!(
+        records.next(),
+        Some(&CapturedLog {
+            body: "ip enums, with format hint: ipv4: 10.0.0.1, ipv6: 2001:db8::1".into(),
+            level: Level::Info,
+            target: "log".into(),
+        })
+    );
+
+    assert_eq!(
+        records.next(),
+        Some(&CapturedLog {
+            body: "ip as bits: ipv4: 10.0.0.1".into(),
+            level: Level::Info,
+            target: "log".into(),
+        })
+    );
+
+    assert_eq!(
+        records.next(),
+        Some(&CapturedLog {
+            body: "ip as octets: ipv4: 10.0.0.1, ipv6: 2001:db8::1".into(),
             level: Level::Info,
             target: "log".into(),
         })