From 96b49fdea1a1c3e13dd9186375e1a2c5d54abd79 Mon Sep 17 00:00:00 2001
From: Sri Goli <sri+github@goli.dev>
Date: Wed, 6 Mar 2024 12:06:36 -0800
Subject: [PATCH] NAT  handling for IPv6 formatted IPv4 addresses

---
 collector/kernel/buffered_poller.cc |  1 +
 collector/kernel/nat_handler.cc     | 57 +++++++++++++++++++++++++++++
 collector/kernel/nat_handler.h      |  1 +
 3 files changed, 59 insertions(+)

diff --git a/collector/kernel/buffered_poller.cc b/collector/kernel/buffered_poller.cc
index 52a958c3..e5482491 100644
--- a/collector/kernel/buffered_poller.cc
+++ b/collector/kernel/buffered_poller.cc
@@ -643,6 +643,7 @@ void BufferedPoller::handle_set_state_ipv6(message_metadata const &metadata, jb_
   }
 
   writer_.set_state_ipv6_tstamp(metadata.timestamp, msg.dest, msg.src, msg.dport, msg.sport, msg.sk, msg.tx_rx);
+  nat_handler_.handle_set_state_ipv6(metadata.timestamp, &msg);
 }
 
 void BufferedPoller::handle_close_socket(message_metadata const &metadata, jb_agent_internal__close_sock_info &msg)
diff --git a/collector/kernel/nat_handler.cc b/collector/kernel/nat_handler.cc
index a427dc66..e4a630b0 100644
--- a/collector/kernel/nat_handler.cc
+++ b/collector/kernel/nat_handler.cc
@@ -208,6 +208,63 @@ void NatHandler::handle_set_state_ipv4(u64 timestamp, jb_agent_internal__set_sta
   }
 }
 
+void NatHandler::handle_set_state_ipv6(u64 timestamp, jb_agent_internal__set_state_ipv6 *msg)
+{
+  // Check if it is an ipv4 address
+  IPv6Address src = IPv6Address::from(msg->src);
+  IPv6Address dst = IPv6Address::from(msg->dest);
+
+  if (!src.is_ipv4() || !dst.is_ipv4()) {
+    if (is_log_whitelisted(AgentLogKind::NAT)) {
+      LOG::trace_in(
+          AgentLogKind::NAT,
+          "NatHandler::handle_set_state_ipv6: "
+          "not a v4 address - skipping nat handling sk={}, src={}:{}, dest={}:{}, tx_rx={}",
+          msg->sk,
+          src,
+          msg->sport,
+          dst,
+          msg->dport,
+          msg->tx_rx);
+    }
+    return;
+  }
+
+  if (is_log_whitelisted(AgentLogKind::NAT)) {
+    LOG::trace_in(
+        AgentLogKind::NAT,
+        "NatHandler::handle_set_state_ipv6: "
+        "sk={}, src={}:{}, dest={}:{}, tx_rx={}",
+        msg->sk,
+        src.to_ipv4().value(),
+        msg->sport,
+        dst.to_ipv4().value(),
+        msg->dport,
+        msg->tx_rx);
+  }
+
+  const u64 sk = msg->sk;
+
+  const hostport_tuple ft = {
+      .src_ip = src.to_ipv4().value().as_int(),
+      .dst_ip = dst.to_ipv4().value().as_int(),
+      .src_port = htons(msg->sport),
+      .dst_port = htons(msg->dport),
+      .proto = IPPROTO_TCP,
+  };
+
+  record_sk(sk, ft);
+
+  // We had a NAT table entry before getting the socket info.
+  if (auto mapping = nat_table_.find(ft); mapping != nat_table_.end()) {
+    send_nat_remapping(timestamp, sk, mapping->second);
+  }
+
+  if (auto rev_mapping = nat_table_rev_.find(ft.reversed()); rev_mapping != nat_table_rev_.end()) {
+    send_nat_remapping(timestamp, sk, rev_mapping->second.reversed());
+  }
+}
+
 void NatHandler::handle_close_socket(u64 timestamp, jb_agent_internal__close_sock_info *msg)
 {
   remove_sk(msg->sk);
diff --git a/collector/kernel/nat_handler.h b/collector/kernel/nat_handler.h
index af7d2ff8..a341a0a3 100644
--- a/collector/kernel/nat_handler.h
+++ b/collector/kernel/nat_handler.h
@@ -29,6 +29,7 @@ class NatHandler {
   void handle_existing_conntrack_tuple(u64 timestamp, struct jb_agent_internal__existing_conntrack_tuple *msg);
 
   void handle_set_state_ipv4(u64 timestamp, jb_agent_internal__set_state_ipv4 *msg);
+  void handle_set_state_ipv6(u64 timestamp, jb_agent_internal__set_state_ipv6 *msg);
 
   void handle_close_socket(u64 timestamp, jb_agent_internal__close_sock_info *msg);