feat(s2n-quic-platform): add ancillary instructions for bpf assembler

quic/s2n-quic-platform/Cargo.toml

@@ -42,6 +42,7 @@ bolero-generator = "0.13"
 futures = { version = "0.3", features = ["std"] }
 insta = { version = "1", features = ["json"] }
 s2n-quic-core = { path = "../s2n-quic-core", features = ["testing"] }
+tempfile = "3"
 tokio = { version = "1", features = ["full"] }
 tracing = { version = "0.1" }
 

quic/s2n-quic-platform/src/bpf.rs

@@ -5,6 +5,8 @@
 pub mod instruction;
 #[macro_use]
 mod common;
+#[macro_use]
+pub mod ancillary;
 
 pub mod cbpf;
 pub mod ebpf;

quic/s2n-quic-platform/src/bpf/ancillary.rs

@@ -0,0 +1,194 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+//! Ancillary data extensions are used to look up additional information about a packet
+//!
+//! See <https://docs.kernel.org/networking/filter.html#bpf-engine-and-instruction-set>
+
+macro_rules! skf_value {
+    ($name:ident) => {
+        (libc::SKF_AD_OFF + libc::$name) as _
+    };
+}
+
+#[derive(Clone, Copy, Debug)]
+pub struct Info {
+    /// The special extension name used in `bpf_asm`
+    pub extension: &'static str,
+    /// The `C` interface
+    pub capi: &'static str,
+}
+
+/// Returns the [`Info`] for a given offset
+pub const fn lookup(offset: u32) -> Option<Info> {
+    macro_rules! map {
+        ($(($value:expr, $extension:expr, $capi:expr),)*) => {
+            match offset {
+                $(_ if offset == $value => Some(Info { extension: $extension, capi: $capi }),)*
+                _ => None,
+            }
+        };
+    }
+
+    map!(
+        (skb::protocol(), "proto", "skb->protocol"),
+        (skb::pkt_type(), "type", "skb->pkt_type"),
+        (skb::ifindex(), "ifidx", "skb->dev->ifindex"),
+        (skb::mark(), "mark", "skb->mark"),
+        (skb::queue_mapping(), "queue", "skb->queue_mapping"),
+        (skb::dev_type(), "hatype", "skb->dev->type"),
+        (skb::hash(), "rxhash", "skb->hash"),
+        (skb::vlan_tci(), "vlan_tci", "skb_vlan_tag_get(skb)"),
+        (skb::vlan_avail(), "vlan_avail", "skb_vlan_tag_present(skb)"),
+        (skb::vlan_proto(), "vlan_tpid", "skb->vlan_tproto"),
+        (payload_offset(), "poff", "payload_offset()"),
+        (raw_smp_processor_id(), "cpu", "raw_smp_processor_id()"),
+        (get_random_u32(), "rand", "get_random_u32()"),
+    )
+}
+
+pub mod skb {
+    macro_rules! skb_value {
+        ($name:ident, $value:ident) => {
+            #[inline]
+            pub const fn $name() -> u32 {
+                skf_value!($value)
+            }
+        };
+    }
+
+    skb_value!(protocol, SKF_AD_PROTOCOL);
+    skb_value!(pkt_type, SKF_AD_PKTTYPE);
+    skb_value!(ifindex, SKF_AD_IFINDEX);
+    skb_value!(mark, SKF_AD_MARK);
+    skb_value!(queue_mapping, SKF_AD_QUEUE);
+    skb_value!(dev_type, SKF_AD_HATYPE);
+    skb_value!(hash, SKF_AD_RXHASH);
+    skb_value!(vlan_tci, SKF_AD_VLAN_TAG);
+    skb_value!(vlan_avail, SKF_AD_VLAN_TAG_PRESENT);
+    skb_value!(vlan_proto, SKF_AD_VLAN_TPID);
+}
+
+#[inline]
+pub const fn payload_offset() -> u32 {
+    skf_value!(SKF_AD_PAY_OFFSET)
+}
+
+#[inline]
+pub const fn raw_smp_processor_id() -> u32 {
+    skf_value!(SKF_AD_CPU)
+}
+
+#[inline]
+pub const fn get_random_u32() -> u32 {
+    skf_value!(SKF_AD_RANDOM)
+}
+
+macro_rules! impl_ancillary {
+    () => {
+        /// Ancillary data extensions are used to look up additional information about a packet
+        ///
+        /// See <https://docs.kernel.org/networking/filter.html#bpf-engine-and-instruction-set>
+        pub mod ancillary {
+            use super::{super::ancillary, *};
+
+            #[inline]
+            const fn ld_ext(k: u32) -> Instruction {
+                Instruction::raw(
+                    Mode::ABS as u16 | Class::LD as u16 | Size::W as u16,
+                    0,
+                    0,
+                    k,
+                )
+            }
+
+            /// Data associated with the socket buffer (skb)
+            pub mod skb {
+                use super::{ancillary::skb, *};
+
+                pub const fn len() -> Instruction {
+                    super::super::len()
+                }
+
+                /// Loads the `skb->protocol` into the `A` register
+                pub const fn protocol() -> Instruction {
+                    ld_ext(skb::protocol())
+                }
+
+                /// Loads the `skb->pkt_type` into the `A` register
+                pub const fn pkt_type() -> Instruction {
+                    ld_ext(skb::pkt_type())
+                }
+
+                /// Loads the `skb->ifindex` into the `A` register
+                pub const fn ifindex() -> Instruction {
+                    ld_ext(skb::ifindex())
+                }
+
+                /// Loads the `skb->mark` into the `A` register
+                pub const fn mark() -> Instruction {
+                    ld_ext(skb::mark())
+                }
+
+                /// Loads the `skb->queue_mapping` into the `A` register
+                pub const fn queue_mapping() -> Instruction {
+                    ld_ext(skb::queue_mapping())
+                }
+
+                /// Loads the `skb->dev->type` into the `A` register
+                pub const fn dev_type() -> Instruction {
+                    ld_ext(skb::dev_type())
+                }
+
+                /// Loads the `skb->hash` into the `A` register
+                pub const fn hash() -> Instruction {
+                    ld_ext(skb::hash())
+                }
+
+                /// Loads the VLAN Tag value into the `A` register
+                pub const fn vlan_tci() -> Instruction {
+                    ld_ext(skb::vlan_tci())
+                }
+
+                /// Loads the VLAN Tag value into the `A` register
+                ///
+                /// This is used for compatibility with the C API
+                pub const fn vlan_tag_get() -> Instruction {
+                    vlan_tci()
+                }
+
+                /// Loads if the VLAN Tag is present into the `A` register
+                pub const fn vlan_avail() -> Instruction {
+                    ld_ext(skb::vlan_avail())
+                }
+
+                /// Loads if the VLAN Tag is present into the `A` register
+                ///
+                /// This is used for compatibility with the C API
+                pub const fn vlan_tag_present() -> Instruction {
+                    vlan_avail()
+                }
+
+                /// Loads the `skb->vlan_proto` (VLAN Protocol) into the `A` register
+                pub const fn vlan_proto() -> Instruction {
+                    ld_ext(skb::vlan_proto())
+                }
+            }
+
+            /// Loads the payload offset into the `A` register
+            pub const fn payload_offset() -> Instruction {
+                ld_ext(ancillary::payload_offset())
+            }
+
+            /// Loads the CPU ID into the `A` register
+            pub const fn raw_smp_processor_id() -> Instruction {
+                ld_ext(ancillary::raw_smp_processor_id())
+            }
+
+            /// Loads a random `u32` into the `A` register
+            pub const fn get_random_u32() -> Instruction {
+                ld_ext(ancillary::get_random_u32())
+            }
+        }
+    };
+}

quic/s2n-quic-platform/src/bpf/cbpf.rs

@@ -3,6 +3,9 @@
 
 use core::fmt;
 
+#[cfg(all(test, not(miri)))]
+mod tests;
+
 pub use super::common::*;
 pub type Instruction = super::Instruction<Cbpf>;
 pub type Program<'a> = super::Program<'a, Cbpf>;
@@ -53,10 +56,11 @@ impl super::instruction::Dialect for Cbpf {
             f.field("jf", &jf);
         }
 
-        f.field("k", &k).finish()
+        let prefix = if k == 0 { "" } else { "0x" };
+        f.field("k", &format_args!("{prefix}{k:x}")).finish()
     }
 
-    fn display(i: &Instruction, f: &mut fmt::Formatter) -> fmt::Result {
+    fn display(i: &Instruction, f: &mut fmt::Formatter, line: Option<usize>) -> fmt::Result {
         let code = i.code;
         let k = i.k;
         let jt = i.jt;
@@ -71,7 +75,19 @@ impl super::instruction::Dialect for Cbpf {
 
                 match mode {
                     Mode::IMM => return write!(f, "{class}{size} #{k}"),
-                    Mode::ABS => return write!(f, "{class}{size} [{k}]"),
+                    Mode::ABS => {
+                        let prefix = if k == 0 { "" } else { "0x" };
+                        return if let Some(info) = super::ancillary::lookup(k) {
+                            write!(
+                                f,
+                                "{class}{size} {} ; [{prefix}{k:x}] // {}",
+                                info.extension, info.capi
+                            )
+                        } else {
+                            write!(f, "{class}{size} [{prefix}{k:x}]")
+                        };
+                    }
+                    Mode::LEN => return write!(f, "{class}{size} len"),
                     _ => {}
                 }
             }
@@ -88,17 +104,30 @@ impl super::instruction::Dialect for Cbpf {
                 let op = Jump::decode(code);
                 let source = Source::decode(code);
 
-                return match source {
-                    Source::K => write!(f, "{op} #{k},{jt},{jf}"),
-                    Source::X => write!(f, "{op} x,{jt},{jf}"),
-                };
+                match source {
+                    Source::K => write!(f, "{op} #{k}")?,
+                    Source::X => write!(f, "{op} x")?,
+                }
+
+                if let Some(line) = line {
+                    let line = line + 1;
+                    let jt = line + jt as usize;
+                    let jf = line + jf as usize;
+                    write!(f, ",l{jt},l{jf}")?
+                } else {
+                    write!(f, ",{jt},{jf}")?
+                }
+
+                return Ok(());
             }
             Class::RET => {
+                let size = Size::decode(code);
                 let source = Source::decode(code);
 
-                return match source {
-                    Source::K => write!(f, "{class} #{k}"),
-                    Source::X => write!(f, "{class} x"),
+                return match (source, size) {
+                    (Source::K, Size::B) if k == 0 => write!(f, "{class} %a"),
+                    (Source::K, _) => write!(f, "{class} #{k}"),
+                    (Source::X, _) => write!(f, "{class} %x"),
                 };
             }
             _ => {}
@@ -192,39 +221,6 @@ define!(
 impl_ops!();
 impl_ret!();
 
-#[cfg(test)]
-mod tests {
-    use crate::bpf::cbpf::*;
-
-    static PROGRAM: Program = {
-        const MAX: u8 = 0b0011_1111;
-
-        Program::new(&[
-            // load the first byte of the packet
-            ldb(0),
-            // mask off the LSBs
-            and(!MAX as _),
-            // IF:
-            // the control bit is set
-            jneq(MAX as u32 + 1, 1, 0),
-            // THEN:
-            // return a 0 indicating we want to route to the writer socket
-            ret(0),
-            // ELSE:
-            // return a 1 indicating we want to route to the reader socket
-            ret(1),
-        ])
-    };
-
-    #[test]
-    #[cfg_attr(miri, ignore)]
-    fn static_program_debug() {
-        insta::assert_debug_snapshot!(PROGRAM);
-    }
-
-    #[test]
-    #[cfg_attr(miri, ignore)]
-    fn static_program_display() {
-        insta::assert_snapshot!(PROGRAM);
-    }
+pub const fn len() -> Instruction {
+    Instruction::raw(Mode::LEN as _, 0, 0, 0)
 }

quic/s2n-quic-platform/src/bpf/cbpf/snapshots/s2n_quic_platform__bpf__cbpf__tests__comparisons_c_literal.snap

@@ -0,0 +1,10 @@
+---
+source: quic/s2n-quic-platform/src/bpf/cbpf/tests.rs
+expression: actual_c_literal
+---
+{ 0x20,  0,  0, 0000000000 },
+{ 0x25,  3,  0, 0x00000001 },
+{ 0x25,  0,  2, 0x00000002 },
+{ 0x15,  1,  0, 0x00000003 },
+{ 0x15,  0,  0, 0x00000004 },
+{ 0x06,  0,  0, 0000000000 },