Add ICMP type/code and udplite/sctp in Nettrace.

This adds support for additional protocols and also provides more clear
semantics of protocol specific fields, such as only including tcpFlags
when the protocol is IPPROTO_TCP.

Test: atest libnetworkstats_test & flash and trace
Change-Id: Ic69fc75e85ebf8734027c942a253af4972ad14d8
diff --git a/service-t/native/libs/libnetworkstats/NetworkTraceHandler.cpp b/service-t/native/libs/libnetworkstats/NetworkTraceHandler.cpp
index 3520254..9b1b72d 100644
--- a/service-t/native/libs/libnetworkstats/NetworkTraceHandler.cpp
+++ b/service-t/native/libs/libnetworkstats/NetworkTraceHandler.cpp
@@ -18,6 +18,7 @@
 
 #include "netdbpf/NetworkTraceHandler.h"
 
+#include <android-base/macros.h>
 #include <arpa/inet.h>
 #include <bpf/BpfUtils.h>
 #include <log/log.h>
@@ -81,14 +82,29 @@
       tag(pkt.tag),
       egress(pkt.egress),
       ipProto(pkt.ipProto),
-      ipVersion(pkt.ipVersion),
-      tcpFlags(pkt.tcpFlags),
-      localPort(ntohs(pkt.egress ? pkt.sport : pkt.dport)),
-      remotePort(ntohs(pkt.egress ? pkt.dport : pkt.sport)) {}
+      ipVersion(pkt.ipVersion) {
+  switch (ipProto) {
+    case IPPROTO_TCP:
+      tcpFlags = pkt.tcpFlags;
+      FALLTHROUGH_INTENDED;
+    case IPPROTO_DCCP:
+    case IPPROTO_UDP:
+    case IPPROTO_UDPLITE:
+    case IPPROTO_SCTP:
+      localPort = ntohs(pkt.egress ? pkt.sport : pkt.dport);
+      remotePort = ntohs(pkt.egress ? pkt.dport : pkt.sport);
+      break;
+    case IPPROTO_ICMP:
+    case IPPROTO_ICMPV6:
+      icmpType = ntohs(pkt.sport);
+      icmpCode = ntohs(pkt.dport);
+      break;
+  }
+}
 
 #define AGG_FIELDS(x)                                                    \
   (x).ifindex, (x).uid, (x).tag, (x).egress, (x).ipProto, (x).ipVersion, \
-      (x).tcpFlags, (x).localPort, (x).remotePort
+      (x).tcpFlags, (x).localPort, (x).remotePort, (x).icmpType, (x).icmpCode
 
 std::size_t BundleHash::operator()(const BundleKey& a) const {
   std::size_t seed = 0;
@@ -265,6 +281,8 @@
   if (src.tcpFlags.has_value()) event->set_tcp_flags(*src.tcpFlags);
   if (src.localPort.has_value()) event->set_local_port(*src.localPort);
   if (src.remotePort.has_value()) event->set_remote_port(*src.remotePort);
+  if (src.icmpType.has_value()) event->set_icmp_type(*src.icmpType);
+  if (src.icmpCode.has_value()) event->set_icmp_code(*src.icmpCode);
 
   event->set_ip_proto(src.ipProto);
 
diff --git a/service-t/native/libs/libnetworkstats/NetworkTraceHandlerTest.cpp b/service-t/native/libs/libnetworkstats/NetworkTraceHandlerTest.cpp
index 2eb2e4f..0c4f049 100644
--- a/service-t/native/libs/libnetworkstats/NetworkTraceHandlerTest.cpp
+++ b/service-t/native/libs/libnetworkstats/NetworkTraceHandlerTest.cpp
@@ -113,7 +113,7 @@
           .length = 100,
           .uid = 10,
           .tag = 123,
-          .ipProto = 6,
+          .ipProto = IPPROTO_TCP,
           .tcpFlags = 1,
       },
   };
@@ -138,12 +138,14 @@
           .sport = htons(8080),
           .dport = htons(443),
           .egress = true,
+          .ipProto = IPPROTO_TCP,
       },
       PacketTrace{
           .timestampNs = 2,
           .sport = htons(443),
           .dport = htons(8080),
           .egress = false,
+          .ipProto = IPPROTO_TCP,
       },
   };
 
@@ -161,6 +163,42 @@
               TrafficDirection::DIR_INGRESS);
 }
 
+TEST_F(NetworkTraceHandlerTest, WriteIcmpTypeAndCode) {
+  std::vector<PacketTrace> input = {
+      PacketTrace{
+          .timestampNs = 1,
+          .sport = htons(11),  // type
+          .dport = htons(22),  // code
+          .egress = true,
+          .ipProto = IPPROTO_ICMP,
+      },
+      PacketTrace{
+          .timestampNs = 2,
+          .sport = htons(33),  // type
+          .dport = htons(44),  // code
+          .egress = false,
+          .ipProto = IPPROTO_ICMPV6,
+      },
+  };
+
+  std::vector<TracePacket> events;
+  ASSERT_TRUE(TraceAndSortPackets(input, &events));
+
+  ASSERT_EQ(events.size(), 2);
+  EXPECT_FALSE(events[0].network_packet().has_local_port());
+  EXPECT_FALSE(events[0].network_packet().has_remote_port());
+  EXPECT_THAT(events[0].network_packet().icmp_type(), 11);
+  EXPECT_THAT(events[0].network_packet().icmp_code(), 22);
+  EXPECT_THAT(events[0].network_packet().direction(),
+              TrafficDirection::DIR_EGRESS);
+  EXPECT_FALSE(events[1].network_packet().local_port());
+  EXPECT_FALSE(events[1].network_packet().remote_port());
+  EXPECT_THAT(events[1].network_packet().icmp_type(), 33);
+  EXPECT_THAT(events[1].network_packet().icmp_code(), 44);
+  EXPECT_THAT(events[1].network_packet().direction(),
+              TrafficDirection::DIR_INGRESS);
+}
+
 TEST_F(NetworkTraceHandlerTest, BasicBundling) {
   // TODO: remove this once bundling becomes default. Until then, set arbitrary
   // aggregation threshold to enable bundling.
@@ -245,6 +283,11 @@
       PacketTrace{.timestampNs = 4, .length = 8, .dport = b, .egress = false},
   };
 
+  // Set common fields.
+  for (PacketTrace& pkt : input) {
+    pkt.ipProto = IPPROTO_TCP;
+  }
+
   std::vector<TracePacket> events;
   ASSERT_TRUE(TraceAndSortPackets(input, &events, config));
   ASSERT_EQ(events.size(), 2);
@@ -280,6 +323,11 @@
       PacketTrace{.timestampNs = 4, .length = 8, .sport = b, .egress = false},
   };
 
+  // Set common fields.
+  for (PacketTrace& pkt : input) {
+    pkt.ipProto = IPPROTO_TCP;
+  }
+
   std::vector<TracePacket> events;
   ASSERT_TRUE(TraceAndSortPackets(input, &events, config));
   ASSERT_EQ(events.size(), 2);
@@ -312,6 +360,11 @@
       PacketTrace{.timestampNs = 4, .length = 4, .uid = 456, .tcpFlags = 2},
   };
 
+  // Set common fields.
+  for (PacketTrace& pkt : input) {
+    pkt.ipProto = IPPROTO_TCP;
+  }
+
   std::vector<TracePacket> events;
   ASSERT_TRUE(TraceAndSortPackets(input, &events, config));
 
diff --git a/service-t/native/libs/libnetworkstats/include/netdbpf/NetworkTraceHandler.h b/service-t/native/libs/libnetworkstats/include/netdbpf/NetworkTraceHandler.h
index 853e26e..6bf186a 100644
--- a/service-t/native/libs/libnetworkstats/include/netdbpf/NetworkTraceHandler.h
+++ b/service-t/native/libs/libnetworkstats/include/netdbpf/NetworkTraceHandler.h
@@ -32,7 +32,7 @@
 
 // BundleKey encodes a PacketTrace minus timestamp and length. The key should
 // match many packets over time for interning. For convenience, sport/dport
-// are parsed here as either local/remote port.
+// are parsed here as either local/remote port or icmp type/code.
 struct BundleKey {
   explicit BundleKey(const PacketTrace& pkt);
 
@@ -47,6 +47,8 @@
   std::optional<uint8_t> tcpFlags;
   std::optional<uint16_t> localPort;
   std::optional<uint16_t> remotePort;
+  std::optional<uint8_t> icmpType;
+  std::optional<uint8_t> icmpCode;
 };
 
 // BundleKeys are hashed using a simple hash combine.