Support KDE zone checking in TPR

Summary:
Today when shivless is enabled on server, we have a case where if client enables shivless but if server rejects the zone, tpr still gets disabled from server.
This diffs adds zone support in TPR.
Next diff will add bpf prog to check the zone before skipping when shivless hdr_opt is present

Reviewed By: avasylev

Differential Revision: D79602910

fbshipit-source-id: c13023721b672dac10b2b7c35e2263e71d43f7fe

Author: nikhildl12

katran/tpr/TPRTypes.h

@@ -30,6 +30,7 @@ struct tcp_opt {
 struct server_info {
   RunningMode running_mode;
   uint8_t kde_enabled;
+  uint8_t kde_zones;
   // only ipv6 is supported
   uint32_t server_id;
 };

katran/tpr/TcpPktRouter.cpp

@@ -128,12 +128,29 @@ bool TcpPktRouter::setServerIdV6(uint32_t id) {
   return true;
 }
 
+folly::Expected<folly::Unit, std::system_error> TcpPktRouter::setServerKDEZone(
+    uint8_t kdeZones) {
+  CHECK_EQ(mode_, RunningMode::SERVER);
+  LOG(INFO) << "Setting kde zone=" << kdeZones;
+
+  kdeZones_ = kdeZones;
+  if (isInitialized_) {
+    auto updateRes = updateServerInfo();
+    if (updateRes.hasError()) {
+      LOG(ERROR) << "Failed to update KDE zone: " << updateRes.error().what();
+      return updateRes;
+    }
+  }
+  return folly::Unit();
+}
+
 folly::Expected<folly::Unit, std::system_error>
 TcpPktRouter::updateServerInfo() noexcept {
   struct server_info info = {};
   if (mode_ == RunningMode::SERVER) {
     info.running_mode = RunningMode::SERVER;
     info.kde_enabled = kdeEnabled_;
+    info.kde_zones = kdeZones_;
     info.server_id = v6Id_;
     if (info.server_id == 0) {
       LOG(WARNING) << "TCP Pkt router is set but server_id is 0. Please check "

katran/tpr/TcpPktRouter.h

@@ -71,6 +71,9 @@ class TcpPktRouter {
     return v6Id_;
   }
 
+  folly::Expected<folly::Unit, std::system_error> setServerKDEZone(
+      uint8_t kdeZones);
+
   RunningMode getMode() {
     return mode_;
   }
@@ -102,6 +105,7 @@ class TcpPktRouter {
   bool isInitialized_{false};
   uint32_t v6Id_;
   bool kdeEnabled_;
+  uint8_t kdeZones_{0};
   std::optional<uint32_t> serverPort_;
   /**
    * Polls stats for packet level events periodically, and

katran/tpr/bpf/tcp_pkt_router_passive_hdlr.h

@@ -15,23 +15,27 @@
 #include "tcp_pkt_router_maps.h"
 #include "tcp_pkt_router_structs.h"
 
-static inline bool kde_enabled() {
-  __u32 sinfo_key = SERVER_INFO_INDEX;
-  struct server_info* s_info = bpf_map_lookup_elem(&server_infos, &sinfo_key);
-  return (s_info && s_info->kde_enabled != 0);
-}
-
-static inline bool has_syn_with_kde_opt(struct bpf_sock_ops* skops) {
-  __u64 load_flags = BPF_LOAD_HDR_OPT_TCP_SYN;
-  struct kde_clt_tcp_opt_v2 kde_opt = {};
-  kde_opt.kind = KDE_CLT_TCP_HDR_OPT_KIND;
-  int ret = bpf_load_hdr_opt(skops, &kde_opt, sizeof(kde_opt), load_flags);
-  return (ret == KDE_CLT_TCP_HDR_OPT_LEN) ||
-      (ret == KDE_CLT_TCP_HDR_OPT_V2_LEN);
-}
+const volatile __u8 KDE_ZONE_ALL = 0xFF;
 
 static inline bool should_ignore_due_to_kde(struct bpf_sock_ops* skops) {
-  return kde_enabled() && has_syn_with_kde_opt(skops);
+  __u32 sinfo_key = SERVER_INFO_INDEX;
+  struct server_info* s_info = bpf_map_lookup_elem(&server_infos, &sinfo_key);
+  // if kde is enabled, check if zone in hdr opts is matching server zone
+  if (s_info && s_info->kde_enabled) {
+    __u64 load_flags = BPF_LOAD_HDR_OPT_TCP_SYN;
+    struct kde_clt_tcp_opt_v2 kde_opt = {};
+    kde_opt.kind = KDE_CLT_TCP_HDR_OPT_KIND;
+    int ret = bpf_load_hdr_opt(skops, &kde_opt, sizeof(kde_opt), load_flags);
+    if (ret == KDE_CLT_TCP_HDR_OPT_V2_LEN) {
+      if (s_info->kde_zones)
+        if ((s_info->kde_zones & kde_opt.zone) == kde_opt.zone) {
+          return true;
+        } else if (s_info->kde_zones == KDE_ZONE_ALL) {
+          return true;
+        }
+    }
+  }
+  return false;
 }
 
 static inline int handle_passive_parse_hdr(

katran/tpr/bpf/tcp_pkt_router_structs.h

@@ -7,6 +7,8 @@ struct server_info {
   __u8 running_mode;
   // 0 = no, otherwise yes
   __u8 kde_enabled;
+  // zones supported by kde
+  __u8 kde_zones;
   __u32 server_id;
 };