From cd4a40b584a5c0fbf0473c9c85db89de64d084e3 Mon Sep 17 00:00:00 2001
From: Kevin Krakauer <krakauer@google.com>
Date: Thu, 14 Dec 2023 22:23:00 -0800
Subject: [PATCH] xdp: use builtin constants and functions in xdp_loader eBPF

With cl/590826409's submission, we can use these.

PiperOrigin-RevId: 591148462
---
 tools/xdp/cmd/bpf/redirect_host.ebpf.c | 20 +++++---------------
 1 file changed, 5 insertions(+), 15 deletions(-)

diff --git a/tools/xdp/cmd/bpf/redirect_host.ebpf.c b/tools/xdp/cmd/bpf/redirect_host.ebpf.c
index cd99a4e21..484b92339 100644
--- a/tools/xdp/cmd/bpf/redirect_host.ebpf.c
+++ b/tools/xdp/cmd/bpf/redirect_host.ebpf.c
@@ -12,9 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <bpf/bpf_endian.h>
 #include <linux/bpf.h>
 #include <linux/if_ether.h>
 #include <linux/ip.h>
+#include <netinet/in.h>
 #include <netinet/tcp.h>
 
 #define section(secname) __attribute__((section(secname), used))
@@ -43,11 +45,6 @@ struct bpf_map_def section("maps") sock_map = {
     .max_entries = 1,
 };
 
-// TODO(b/240191988): It would be better to use bfp_htons() in
-// <bpf/bpf_endian.h>. However, we test on ARM64 and AMD64 with the same Docker
-// image, and each architecture requires different packages to get that header.
-const __u16 swapped_eth_p_ip = (__u16)((ETH_P_IP << 8) | (ETH_P_IP >> 8));
-
 // Redirect almost all incoming traffic to an AF_XDP socket. Certain packets are
 // allowed through to the Linux network stack:
 //
@@ -65,7 +62,7 @@ section("xdp") int xdp_prog(struct xdp_md *ctx) {
   cursor += sizeof(*eth);
 
   // Send all non-IPv4 traffic to the socket.
-  if (eth->h_proto != swapped_eth_p_ip) {
+  if (eth->h_proto != bpf_htons(ETH_P_IP)) {
     return bpf_redirect_map(&sock_map, ctx->rx_queue_index, XDP_PASS);
   }
 
@@ -76,10 +73,7 @@ section("xdp") int xdp_prog(struct xdp_md *ctx) {
   }
   cursor += sizeof(*ip);
 
-  // TODO(b/240191988): It would be better to use IPPROTO_TCP in <netinet/in.h>.
-  // However, we test on ARM64 and AMD64 with the same Docker image, and each
-  // architecture requires different packages to get that header.
-  if (ip->protocol != 6 /* IPPROTO_TCP */) {
+  if (ip->protocol != IPPROTO_TCP) {
     return bpf_redirect_map(&sock_map, ctx->rx_queue_index, XDP_PASS);
   }
   struct tcphdr *tcp = cursor;
@@ -88,11 +82,7 @@ section("xdp") int xdp_prog(struct xdp_md *ctx) {
   }
 
   // Allow port 22 traffic for SSH debugging.
-  // TODO(b/240191988): It would be better to use bfp_ntohs() in
-  // <bpf/bpf_endian.h>. However, we test on ARM64 and AMD64 with the same
-  // Docker image, and each architecture requires different packages to get that
-  // header.
-  if (__builtin_bswap16(tcp->th_dport) == 22) {
+  if (tcp->th_dport == bpf_htons(22)) {
     return XDP_PASS;
   }