Merge pull request rex-rs#33 from rex-rs/electrode

 - samples/electrode: replace reads of `iphdr.protocol` and `udphdr.dest` with `read_field!`
 - rex/utils: add `read_field!`
 - samples/electrode: fix the access of wrong index of magic bits in skb payload
 - samples/electrode: introduce PAXOS_PORT and MAGIC_BITS constants

Signed-off-by: Jinghao Jia <[email protected]>

Author: jinghao-jia

rex/src/utils.rs

@@ -258,3 +258,31 @@ where
         unsafe { AlignedMut::from_val(ptr.read_unaligned(), slice) }
     }
 }
+
+/// Read a numeric field that is stored **big-endian** inside a header already
+/// sitting in `payload_slice` at offset `hdr_base`.
+///
+/// Example:
+/// ```
+/// let iphdr_base = size_of::<ethhdr>();
+/// let proto_be = read_field!(skb.data_slice, iphdr_base, iphdr, protocol, u8);
+/// match u8::from_be(proto_be) as u32 {
+///     IPPROTO_TCP => handle_tcp(),
+///     IPPROTO_UDP => handle_udp(),
+///     _  => {}
+/// }
+/// ```
+#[macro_export]
+macro_rules! read_field {
+    ($slice:expr,     // payload slice
+     $hdr_base:expr,  // where this header starts inside the buffer
+     $hdr:path,       // concrete header type, for `offset_of!`
+     $field:ident,    // field we want
+     $ty:ty           // Rust type of that field (u8, u16, …)
+    ) => {{
+        let start = $hdr_base + core::mem::offset_of!($hdr, $field);
+        *$crate::utils::convert_slice_to_struct::<$ty>(
+            &$slice[start..start + size_of::<$ty>()],
+        )
+    }};
+}

samples/electrode/src/common.rs

@@ -31,6 +31,9 @@ pub(crate) const BROADCAST_SIGN_BIT: u32 = 1 << 31;
 pub(crate) const QUORUM_SIZE: u32 = (CLUSTER_SIZE as u32 + 1) >> 1;
 pub(crate) const QUORUM_BITSET_ENTRY: u32 = 1024; // must be 2^t
 
+pub(crate) const PAXOS_PORT: u16 = 12345;
+pub(crate) const MAGIC_BITS: [u8; MAGIC_LEN] = [0x18, 0x03, 0x05, 0x20];
+
 #[derive(Debug, PartialEq, Eq, Clone, Copy)]
 #[repr(C)]
 pub(crate) enum ReplicaStatus {

samples/electrode/src/main.rs

@@ -1,15 +1,16 @@
 #![no_std]
 #![no_main]
 #![allow(non_camel_case_types)]
+#![allow(non_snake_case)]
 
 extern crate rex;
 
-use core::mem::{offset_of, size_of, swap};
+use core::mem::{size_of, swap};
 
 use rex::sched_cls::*;
 use rex::utils::*;
 use rex::xdp::*;
-use rex::{rex_printk, rex_tc, rex_xdp};
+use rex::{read_field, rex_printk, rex_tc, rex_xdp};
 
 pub mod common;
 pub mod maps;
@@ -33,37 +34,36 @@ macro_rules! swap_field {
 fn fast_paxos_main(obj: &xdp, ctx: &mut xdp_md) -> Result {
     let header_len =
         size_of::<ethhdr>() + size_of::<iphdr>() + size_of::<udphdr>();
-    let iphdr_start = size_of::<ethhdr>();
-    let iphdr_end = iphdr_start + size_of::<iphdr>();
-    let udphdr_end = iphdr_end + size_of::<udphdr>();
-
-    let protocol_start = iphdr_start + offset_of!(iphdr, protocol);
-    let protocol = convert_slice_to_struct::<u8>(
-        &ctx.data_slice[protocol_start..protocol_start + size_of::<u8>()],
-    );
-
-    match u8::from_be(*protocol) as u32 {
+    let iphdr_base = size_of::<ethhdr>();
+    let iphdr_end = iphdr_base + size_of::<iphdr>();
+
+    match u8::from_be(read_field!(
+        ctx.data_slice,
+        iphdr_base,
+        iphdr,
+        protocol,
+        u8
+    )) as u32
+    {
         IPPROTO_TCP => {
             // NOTE: currently we only take care of UDP memcached
         }
         IPPROTO_UDP => {
-            let port_start = iphdr_end + offset_of!(udphdr, dest);
-            let port = u16::from_be(*convert_slice_to_struct::<u16>(
-                &ctx.data_slice[port_start..port_start + size_of::<u16>()],
+            let port = u16::from_be(read_field!(
+                ctx.data_slice,
+                iphdr_end,
+                udphdr,
+                dest,
+                u16
             ));
+
             let payload = &mut ctx.data_slice[header_len..];
 
             // port check, our process bound to 12345.
             // don't have magic bits...
-            if port != 12345 || payload.len() < MAGIC_LEN + size_of::<u64>() {
-                return Ok(XDP_PASS as i32);
-            }
-
-            // NOTE: currently, we don't support reassembly.
-            if payload[0] != 0x18 ||
-                payload[1] != 0x03 ||
-                payload[2] != 0x05 ||
-                payload[3] != 0x20
+            if port != PAXOS_PORT ||
+                payload.len() < MAGIC_LEN + size_of::<u64>() ||
+                !payload.starts_with(&MAGIC_BITS)
             {
                 return Ok(XDP_PASS as i32);
             }
@@ -79,44 +79,39 @@ fn fast_paxos_main(obj: &xdp, ctx: &mut xdp_md) -> Result {
 
 #[rex_tc]
 fn fast_broad_cast_main(obj: &sched_cls, skb: &mut __sk_buff) -> Result {
-    let mut header_len =
+    let header_len =
         size_of::<iphdr>() + size_of::<ethhdr>() + size_of::<udphdr>();
-    let iphdr_start = size_of::<ethhdr>();
-    let iphdr_end = iphdr_start + size_of::<iphdr>();
+    let iphdr_base = size_of::<ethhdr>();
+    let iphdr_end = iphdr_base + size_of::<iphdr>();
 
     // check if the packet is long enough
     if skb.data_slice.len() <= header_len {
         return Ok(TC_ACT_OK as i32);
     }
 
-    let protocol_start = iphdr_start + offset_of!(iphdr, protocol);
-    let protocol = convert_slice_to_struct::<u8>(
-        &skb.data_slice[protocol_start..protocol_start + size_of::<u8>()],
-    );
-    match u8::from_be(*protocol) as u32 {
+    match u8::from_be(read_field!(
+        skb.data_slice,
+        iphdr_base,
+        iphdr,
+        protocol,
+        u8
+    )) as u32
+    {
         IPPROTO_UDP => {
-            // only port 12345 is allowed
-            let port_start = iphdr_end + offset_of!(udphdr, dest);
-            let port = u16::from_be(*convert_slice_to_struct::<u16>(
-                &skb.data_slice[port_start..port_start + size_of::<u16>()],
+            let port = u16::from_be(read_field!(
+                skb.data_slice,
+                iphdr_end,
+                udphdr,
+                dest,
+                u16
             ));
-            if port != 12345 {
-                return Ok(TC_ACT_OK as i32);
-            }
-
-            // check for the magic bits
-            header_len += MAGIC_LEN + size_of::<u64>();
 
-            if skb.data_slice.len() < header_len {
-                rex_printk!("data_slice.len() < header_len\n").ok();
-                return Ok(TC_ACT_OK as i32);
-            }
-            let payload = &skb.data_slice;
-
-            if payload[0] != 0x18 ||
-                payload[1] != 0x03 ||
-                payload[2] != 0x05 ||
-                payload[3] != 0x20
+            let payload = &skb.data_slice[header_len..];
+            // check for the magic bits and Paxos port
+            // only port 12345 is allowed
+            if port != PAXOS_PORT ||
+                payload.len() < MAGIC_LEN + size_of::<u64>() ||
+                !payload.starts_with(&MAGIC_BITS)
             {
                 return Ok(TC_ACT_OK as i32);
             }