Add basic Conntrack for UDP (#60)

Fixes #45 

Signed-off-by: Andrew Stoycos <[email protected]>

.github/workflows/test.yaml

@@ -67,3 +67,13 @@ jobs:
       env:
         BLIXT_CONTROLPLANE_IMAGE: "ghcr.io/kong/blixt-controlplane:integration-tests"
         BLIXT_DATAPLANE_IMAGE: "ghcr.io/kong/blixt-dataplane:integration-tests"
+        BLIXT_UDP_SERVER_IMAGE: "ghcr.io/kong/blixt-udp-test-server:integration-tests"
+
+    ## Upload diagnostics if integration test step failed.
+    - name: upload diagnostics
+      if: ${{ failure() }}
+      uses: actions/upload-artifact@v3
+      with:
+        name: blixt-integration-test-diag
+        path: /tmp/ktf-diag*
+        if-no-files-found: ignore

Makefile

@@ -138,6 +138,7 @@ build.image:
 .PHONY: build.all.images
 build.all.images: build.image
 	cd dataplane/ && make build.image TAG=$(TAG)
+	cd tools/udp-test-server && make build.image TAG=$(TAG)
 
 ##@ Deployment
 

config/samples/udproute/kustomization.yaml

@@ -2,4 +2,4 @@ resources:
 - gateway.yaml
 - server.yaml
 - udproute.yaml
-#+kubebuilder:scaffold:manifestskustomizesamples
+#+kubebuilder:scaffold:manifestskustomizesamples

config/tests/kustomization.yaml → config/tests/blixt/kustomization.yaml

@@ -1,5 +1,5 @@
 bases:
-- ../default
+- ../../default
 
 images:
 - name: ghcr.io/kong/blixt-dataplane

config/tests/udproute-noreach/kustomization.yaml

@@ -0,0 +1,23 @@
+resources: 
+- ../../samples/udproute
+
+patchesStrategicMerge:
+- |-
+  apiVersion: apps/v1
+  kind: Deployment
+  metadata:
+    name: blixt-udproute-sample
+  spec:
+    template:
+      spec:
+        containers:
+        - name: server
+          command:
+            - ./udp-test-server
+            # --dry-run disables UDP listeners in order to test failures to send
+            # data, and trigger ICMP port failure responses from the kernel
+            - --dry-run
+
+images:
+- name: ghcr.io/kong/blixt-udp-test-server
+  newTag: integration-tests

config/tests/udproute/kustomization.yaml

@@ -0,0 +1,6 @@
+bases:
+- ../../samples/udproute
+
+images:
+- name: ghcr.io/kong/blixt-udp-test-server
+  newTag: integration-tests

dataplane/ebpf/src/bindings.rs

@@ -1,4 +1,4 @@
-/* automatically generated by rust-bindgen 0.63.0 */
+/* automatically generated by rust-bindgen 0.64.0 */
 
 #[repr(C)]
 #[derive(Copy, Clone, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
@@ -334,3 +334,31 @@ pub struct udphdr {
     pub len: __be16,
     pub check: __sum16,
 }
+#[repr(C)]
+#[derive(Copy, Clone)]
+pub struct icmphdr {
+    pub type_: __u8,
+    pub code: __u8,
+    pub checksum: __sum16,
+    pub un: icmphdr__bindgen_ty_1,
+}
+#[repr(C)]
+#[derive(Copy, Clone)]
+pub union icmphdr__bindgen_ty_1 {
+    pub echo: icmphdr__bindgen_ty_1__bindgen_ty_1,
+    pub gateway: __be32,
+    pub frag: icmphdr__bindgen_ty_1__bindgen_ty_2,
+    pub reserved: [__u8; 4usize],
+}
+#[repr(C)]
+#[derive(Copy, Clone)]
+pub struct icmphdr__bindgen_ty_1__bindgen_ty_1 {
+    pub id: __be16,
+    pub sequence: __be16,
+}
+#[repr(C)]
+#[derive(Copy, Clone)]
+pub struct icmphdr__bindgen_ty_1__bindgen_ty_2 {
+    pub __unused: __be16,
+    pub mtu: __be16,
+}

dataplane/ebpf/src/egress/icmp.rs

@@ -0,0 +1,90 @@
+use core::mem;
+
+use aya_bpf::{
+    bindings::TC_ACT_PIPE,
+    helpers::bpf_csum_diff,
+    programs::TcContext,
+};
+use aya_log_ebpf::info;
+
+use crate::{
+    bindings::{iphdr, icmphdr},
+    utils::{csum_fold_helper, ip_from_int, ptr_at, ETH_HDR_LEN, IP_HDR_LEN},
+    BLIXT_CONNTRACK,
+};
+
+const ICMP_HDR_LEN: usize = mem::size_of::<icmphdr>();
+const ICMP_PROTO_TYPE_UNREACH: u8 = 3;
+
+pub fn handle_icmp_egress(ctx: TcContext) -> Result<i32, i64> {
+    let ip_hdr: *mut iphdr = unsafe { ptr_at(&ctx, ETH_HDR_LEN) }?;
+
+    let icmp_header_offset = ETH_HDR_LEN + IP_HDR_LEN;
+
+    let icmp_hdr: *mut icmphdr = unsafe { 
+        ptr_at(
+            &ctx,
+            icmp_header_offset
+        )?
+    };
+
+    // We only care about redirecting port unreachable messages currently so a 
+    // UDP client can tell when the server is shutdown
+    if unsafe { (*icmp_hdr).type_ } != ICMP_PROTO_TYPE_UNREACH { 
+        return Ok(TC_ACT_PIPE);
+    }   
+
+    let dest_addr = unsafe { (*ip_hdr).daddr };
+
+    let new_src = unsafe { BLIXT_CONNTRACK.get(&dest_addr) }.ok_or(TC_ACT_PIPE)?;
+
+    let daddr_dot_dec = ip_from_int(unsafe { (*ip_hdr).daddr });
+    info!(
+        &ctx,
+        "Received a ICMP Unreachable packet destined for svc ip: {}.{}.{}.{}",
+        daddr_dot_dec[0],
+        daddr_dot_dec[1],
+        daddr_dot_dec[2],
+        daddr_dot_dec[3],
+    );
+
+    // redirect icmp unreachable message back to client
+    unsafe { 
+        (*ip_hdr).saddr = *new_src;
+        (*ip_hdr).check = 0;
+    }
+
+    let full_cksum = unsafe { 
+        bpf_csum_diff(
+            mem::MaybeUninit::zeroed().assume_init(),
+            0,
+            ip_hdr as *mut u32,
+            mem::size_of::<iphdr>() as u32,
+            0
+        )
+    } as u64;
+    unsafe { (*ip_hdr).check = csum_fold_helper(full_cksum) };   
+
+    // Get inner ipheader since we need to update that as well
+    let icmp_inner_ip_hdr: *mut iphdr  = unsafe { ptr_at(&ctx, icmp_header_offset + ICMP_HDR_LEN)}?;
+
+    unsafe { 
+        (*icmp_inner_ip_hdr).daddr = *new_src;
+        (*icmp_inner_ip_hdr).check = 0;
+    }
+
+    let full_cksum = unsafe { 
+        bpf_csum_diff(
+            mem::MaybeUninit::zeroed().assume_init(),
+            0,
+            icmp_inner_ip_hdr as *mut u32,
+            mem::size_of::<iphdr>() as u32,
+            0
+        )
+    } as u64;
+    unsafe { (*icmp_inner_ip_hdr).check = csum_fold_helper(full_cksum) };
+
+    unsafe { BLIXT_CONNTRACK.remove(&dest_addr)? };
+
+    return Ok(TC_ACT_PIPE);
+}

dataplane/ebpf/src/egress/mod.rs

@@ -0,0 +1 @@
+pub mod icmp;

dataplane/ebpf/src/ingress/udp.rs

@@ -1,7 +1,7 @@
 use core::mem;
 
 use aya_bpf::{
-    bindings::TC_ACT_OK,
+    bindings::TC_ACT_PIPE,
     helpers::{bpf_csum_diff, bpf_redirect_neigh},
     programs::TcContext,
 };
@@ -11,6 +11,7 @@ use crate::{
     bindings::{iphdr, udphdr},
     utils::{csum_fold_helper, ip_from_int, ptr_at, ETH_HDR_LEN, IP_HDR_LEN},
     BACKENDS,
+    BLIXT_CONNTRACK,
 };
 use common::BackendKey;
 
@@ -21,12 +22,14 @@ pub fn handle_udp_ingress(ctx: TcContext) -> Result<i32, i64> {
 
     let udp_hdr: *mut udphdr = unsafe { ptr_at(&ctx, udp_header_offset)? };
 
+    let original_daddr = unsafe { (*ip_hdr).daddr };
+
     let key = BackendKey {
-        ip: u32::from_be(unsafe { (*ip_hdr).daddr }),
+        ip: u32::from_be(original_daddr),
         port: (u16::from_be(unsafe { (*udp_hdr).dest })) as u32,
     };
 
-    let backend = unsafe { BACKENDS.get(&key) }.ok_or(TC_ACT_OK)?;
+    let backend = unsafe { BACKENDS.get(&key) }.ok_or(TC_ACT_PIPE)?;
 
     let daddr_dot_dec = ip_from_int(unsafe { (*ip_hdr).daddr });
     info!(
@@ -40,12 +43,13 @@ pub fn handle_udp_ingress(ctx: TcContext) -> Result<i32, i64> {
     );
 
     unsafe {
+        BLIXT_CONNTRACK.insert(&(*ip_hdr).saddr, &original_daddr, 0 as u64)?;
         (*ip_hdr).daddr = backend.daddr.to_be();
-    }
-
+    };
+    
     if (ctx.data() + ETH_HDR_LEN + mem::size_of::<iphdr>()) > ctx.data_end() {
         info!(&ctx, "Iphdr is out of bounds");
-        return Ok(TC_ACT_OK);
+        return Ok(TC_ACT_PIPE);
     }
 
     // Calculate l3 cksum

dataplane/ebpf/src/main.rs

@@ -8,11 +8,12 @@
 mod bindings;
 mod ingress;
 mod utils;
+mod egress;
 
 use memoffset::offset_of;
 
 use aya_bpf::{
-    bindings::{TC_ACT_OK, TC_ACT_PIPE, TC_ACT_SHOT},
+    bindings::{TC_ACT_PIPE, TC_ACT_SHOT, TC_ACT_OK},
     macros::{classifier, map},
     maps::HashMap,
     programs::TcContext,
@@ -21,7 +22,8 @@ use aya_bpf::{
 use bindings::{ethhdr, iphdr};
 use common::{Backend, BackendKey};
 use ingress::{tcp::handle_tcp_ingress, udp::handle_udp_ingress};
-use utils::{ETH_HDR_LEN, ETH_P_IP, IPPROTO_TCP, IPPROTO_UDP};
+use egress::{icmp::handle_icmp_egress};
+use utils::{ETH_HDR_LEN, ETH_P_IP, IPPROTO_TCP, IPPROTO_UDP, IPPROTO_ICMP};
 
 // -----------------------------------------------------------------------------
 // Maps
@@ -31,6 +33,10 @@ use utils::{ETH_HDR_LEN, ETH_P_IP, IPPROTO_TCP, IPPROTO_UDP};
 static mut BACKENDS: HashMap<BackendKey, Backend> =
     HashMap::<BackendKey, Backend>::with_max_entries(128, 0);
 
+#[map(name = "BLIXT_CONNTRACK")] 
+static mut BLIXT_CONNTRACK: HashMap<u32, u32> =
+    HashMap::<u32, u32>::with_max_entries(128, 0);
+
 // -----------------------------------------------------------------------------
 // Ingress
 // -----------------------------------------------------------------------------
@@ -42,6 +48,7 @@ pub fn tc_ingress(ctx: TcContext) -> i32 {
         Err(_) => TC_ACT_SHOT,
     };
 
+    // TODO(https://github.com/Kong/blixt/issues/69) better Error reporting framework
     return TC_ACT_OK;
 }
 
@@ -78,12 +85,28 @@ pub fn tc_egress(ctx: TcContext) -> i32 {
         Err(_) => TC_ACT_SHOT,
     };
 
+    // TODO(https://github.com/Kong/blixt/issues/69) better Error reporting framework
     return TC_ACT_OK;
 }
 
-fn try_tc_egress(_ctx: TcContext) -> Result<i32, i64> {
-    // TODO: not implemented yet
-    Ok(TC_ACT_PIPE)
+fn try_tc_egress(ctx: TcContext) -> Result<i32, i64> {
+    let h_proto = u16::from_be(
+        ctx.load(offset_of!(ethhdr, h_proto))
+            .map_err(|_| TC_ACT_PIPE)?,
+    );
+
+    if h_proto != ETH_P_IP {
+        return Ok(TC_ACT_PIPE);
+    }
+
+    let protocol = ctx
+        .load::<u8>(ETH_HDR_LEN + offset_of!(iphdr, protocol))
+        .map_err(|_| TC_ACT_PIPE)?;
+
+    match protocol {
+        IPPROTO_ICMP => handle_icmp_egress(ctx),
+        _ => Ok(TC_ACT_PIPE),
+    }
 }
 
 // -----------------------------------------------------------------------------

dataplane/ebpf/src/utils.rs

@@ -12,6 +12,7 @@ pub const ETH_P_IP: u16 = 0x0800;
 
 pub const IPPROTO_TCP: u8 = 6;
 pub const IPPROTO_UDP: u8 = 17;
+pub const IPPROTO_ICMP: u8 = 1;
 
 pub const ETH_HDR_LEN: usize = mem::size_of::<ethhdr>();
 pub const IP_HDR_LEN: usize = mem::size_of::<iphdr>();

dataplane/loader/Cargo.toml

@@ -11,7 +11,7 @@ common = { path = "../common", features=["user"] }
 clap = { version = "3.1", features = ["derive"] }
 env_logger = "0.9"
 log = "0.4"
-tokio = { version = "1.18", features = ["macros", "rt", "rt-multi-thread", "net", "signal"] }
+tokio = { version = "1.25.0", features = ["macros", "rt", "rt-multi-thread", "net", "signal"] }
 api-server = { path = "../api-server" }
 anyhow = "1"
 

dataplane/xtask/src/codegen.rs

@@ -3,7 +3,7 @@ use std::{fs::File, io::Write, path::PathBuf};
 
 pub fn generate() -> Result<(), anyhow::Error> {
     let dir = PathBuf::from("ebpf/src");
-    let names: Vec<&str> = vec!["ethhdr", "iphdr", "udphdr", "tcphdr"];
+    let names: Vec<&str> = vec!["ethhdr", "iphdr", "udphdr", "tcphdr","icmphdr"];
     let bindings = aya_tool::generate(
         InputFile::Btf(PathBuf::from("/sys/kernel/btf/vmlinux")),
         &names,