improvements

src/main.c

@@ -1,5 +1,5 @@
 #include <linux/bpf.h>
-#define fuck_lsp
+#define fuck_formater
 #include <bpf/bpf_endian.h>
 #include <bpf/bpf_helpers.h>
 
@@ -30,29 +30,81 @@
 const __u16 open_ports[] = {4444};
 #define HIGH_PORT 5000
 
+#define MAX_ITERATIONS 1024
+
 char _license[] SEC("license") = "GPL";
 
+__u64 time_seed = 0xffffffffffffffff;
+
+#define DAY_LENGHT_IN_NS 86400000000000
+
+__u32 mulberry32_at(__u32 seed, __u32 index) {
+    __u32 z = seed + index * 0x6D2B79F5u;
+    z = (z ^ (z >> 15)) * (z | 1u);
+    z ^= z + (z ^ (z >> 7)) * (z | 61u);
+    return z ^ (z >> 14);
+}
+
+static __always_inline const bool
+binary_search(const __u16 *arr, const __u64 arr_len, const __u16 searched_for) {
+  if (arr_len == 0) {
+    return false;
+  }
+  if (searched_for < arr[0]) {
+    return false;
+  } else if (arr[0] == searched_for) {
+    return true;
+  } else {
+    if (searched_for > arr[arr_len - 1]) {
+      return false;
+    }
+    __u64 lover_bound_index = 0;
+    __u64 upper_bound_index = arr_len - 1;
+    for (__u64 i = 0; i < MAX_ITERATIONS; i++) {
+      __u64 half_index =
+          lover_bound_index + (upper_bound_index - lover_bound_index) / 2;
+      if (searched_for < arr[half_index]) {
+        upper_bound_index = half_index;
+      } else if (searched_for > arr[half_index]) {
+        lover_bound_index = half_index + 1;
+      } else if (searched_for == arr[half_index]) {
+        return true;
+      }
+    }
+  }
+  return false;
+}
+
 static __always_inline const __u8 *get_tcp_data(const __u16 EtherType,
                                                 const __u8 *ip_packet_data,
                                                 const __u8 *data_end) {
   if (EtherType == ETHER_TYPE_IPV4) {
+    if (ip_packet_data + 1 > data_end) {
+      return NULL;
+    }
 
     // check protocol
     const __u8 *protocol_id_ptr = ip_packet_data + IPV4_PROTOCOL_OFFSET;
-    if (protocol_id_ptr > data_end) {
+    if (protocol_id_ptr + 1 > data_end) {
       return NULL;
     }
     if (*protocol_id_ptr == TCP_PROTO_ID) {
+      if (ip_packet_data + ((*ip_packet_data & 0b00001111) * 4) > data_end) {
+        return NULL;
+      }
       return ip_packet_data + ((*ip_packet_data & 0b00001111) * 4);
     }
 
   } else if (EtherType == ETHER_TYPE_IPV6) {
     // check  Next Header
     const __u8 *protocol_id_ptr = ip_packet_data + IPV6_NEXT_HEADER_OFFSET;
-    if (protocol_id_ptr > data_end) {
+    if (protocol_id_ptr + 1 > data_end) {
       return NULL;
     }
     if (*protocol_id_ptr == TCP_PROTO_ID) {
+      if (ip_packet_data + IPV6_HEADER_SIZE > data_end) {
+        return NULL;
+      }
       return ip_packet_data + IPV6_HEADER_SIZE;
     }
   }
@@ -61,14 +113,17 @@ static __always_inline const __u8 *get_tcp_data(const __u16 EtherType,
 
 SEC("faker")
 int xdp_drop_prog(struct xdp_md *ctx) {
+  if (bpf_ktime_get_ns() - time_seed > DAY_LENGHT_IN_NS) {
+    time_seed = bpf_ktime_get_ns();
+  }
+
   const __u8 *data_end = (__u8 *)(long)ctx->data_end;
   const __u8 *data = (__u8 *)(long)ctx->data;
 
-  const __u16 *EtherType_ptr = (__u16 *)(data + ETHER_TYPE_FIELD_OFFSET);
-  if (EtherType_ptr > (__u16 *)data_end) {
+  if (data + ETHER_TYPE_FIELD_OFFSET + sizeof(__u16) > data_end) {
     return XDP_PASS;
   }
-  const __u16 EtherType = bpf_ntohs(*EtherType_ptr);
+  const __u16 EtherType = bpf_ntohs(*(__u16 *)(data + ETHER_TYPE_FIELD_OFFSET));
 
   const __u8 *ip_packet_data = data + ETH_HEADER_SIZE;
 
@@ -76,9 +131,14 @@ int xdp_drop_prog(struct xdp_md *ctx) {
       get_tcp_data(EtherType, ip_packet_data, data_end);
   if (!tcp_packet_data) {
     return XDP_PASS;
-  } // return if not tcp or IP
+  }
 
-  const __u16 dst_port = bpf_ntohs(*(tcp_packet_data + TCP_DST_PORT_OFFSET));
+  if (tcp_packet_data + TCP_FLAGS_OFFSET + 1 > data_end) {
+    return XDP_PASS;
+  }
+
+  const __u16 dst_port =
+      bpf_ntohs(*(__u16 *)(tcp_packet_data + TCP_DST_PORT_OFFSET));
 
   if (dst_port < LOW_PORT) {
     return XDP_PASS;
@@ -86,20 +146,23 @@ int xdp_drop_prog(struct xdp_md *ctx) {
   if (dst_port > HIGH_PORT) {
     return XDP_PASS;
   }
+  bpf_printk("packet to port: %d\n", dst_port);
 
-  for (__u16 i = 0; i < ARRAY_SIZE(open_ports); i++) {
-    if (dst_port == open_ports[i]) {
-      return XDP_PASS;
-    }
+  if (binary_search(open_ports, ARRAY_SIZE(open_ports), dst_port)) {
+    return XDP_PASS;
   }
 
   const __u8 tcp_flags = *(tcp_packet_data + TCP_FLAGS_OFFSET);
 
   // check SYN
   if (tcp_flags & SYN_MASK) {
-    if (bpf_get_prandom_u32() % 3 == 0) {
+    bpf_printk("packet is SYN\n");
+    if (mulberry32_at((__u32)time_seed, dst_port) % 3== 0) {
+      bpf_printk("DROPING!\n");
       return XDP_DROP;
-    };
+    }
+  } else {
+    bpf_printk("packet is  not SYN\n");
   }
 
   return XDP_PASS;