diff --git a/service-t/native/libs/libnetworkstats/NetworkTraceHandler.cpp b/service-t/native/libs/libnetworkstats/NetworkTraceHandler.cpp index 3520254acb..9b1b72dac6 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 #include #include #include @@ -81,14 +82,29 @@ BundleKey::BundleKey(const PacketTrace& pkt) 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 @@ void NetworkTraceHandler::Fill(const BundleKey& src, 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 2eb2e4f313..0c4f049ce1 100644 --- a/service-t/native/libs/libnetworkstats/NetworkTraceHandlerTest.cpp +++ b/service-t/native/libs/libnetworkstats/NetworkTraceHandlerTest.cpp @@ -113,7 +113,7 @@ TEST_F(NetworkTraceHandlerTest, WriteBasicFields) { .length = 100, .uid = 10, .tag = 123, - .ipProto = 6, + .ipProto = IPPROTO_TCP, .tcpFlags = 1, }, }; @@ -138,12 +138,14 @@ TEST_F(NetworkTraceHandlerTest, WriteDirectionAndPorts) { .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 @@ TEST_F(NetworkTraceHandlerTest, WriteDirectionAndPorts) { TrafficDirection::DIR_INGRESS); } +TEST_F(NetworkTraceHandlerTest, WriteIcmpTypeAndCode) { + std::vector 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 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 @@ TEST_F(NetworkTraceHandlerTest, DropLocalPort) { PacketTrace{.timestampNs = 4, .length = 8, .dport = b, .egress = false}, }; + // Set common fields. + for (PacketTrace& pkt : input) { + pkt.ipProto = IPPROTO_TCP; + } + std::vector events; ASSERT_TRUE(TraceAndSortPackets(input, &events, config)); ASSERT_EQ(events.size(), 2); @@ -280,6 +323,11 @@ TEST_F(NetworkTraceHandlerTest, DropRemotePort) { PacketTrace{.timestampNs = 4, .length = 8, .sport = b, .egress = false}, }; + // Set common fields. + for (PacketTrace& pkt : input) { + pkt.ipProto = IPPROTO_TCP; + } + std::vector events; ASSERT_TRUE(TraceAndSortPackets(input, &events, config)); ASSERT_EQ(events.size(), 2); @@ -312,6 +360,11 @@ TEST_F(NetworkTraceHandlerTest, DropTcpFlags) { PacketTrace{.timestampNs = 4, .length = 4, .uid = 456, .tcpFlags = 2}, }; + // Set common fields. + for (PacketTrace& pkt : input) { + pkt.ipProto = IPPROTO_TCP; + } + std::vector 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 853e26e09e..6bf186aa63 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 @@ namespace bpf { // 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 @@ struct BundleKey { std::optional tcpFlags; std::optional localPort; std::optional remotePort; + std::optional icmpType; + std::optional icmpCode; }; // BundleKeys are hashed using a simple hash combine.