1 /*
2  * Copyright (C) 2023 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #include <gmock/gmock.h>
18 #include <gtest/gtest.h>
19 
20 #include <vector>
21 
22 #include "netdbpf/NetworkTraceHandler.h"
23 #include "protos/perfetto/config/android/network_trace_config.gen.h"
24 #include "protos/perfetto/trace/android/network_trace.pb.h"
25 #include "protos/perfetto/trace/trace.pb.h"
26 #include "protos/perfetto/trace/trace_packet.pb.h"
27 
28 namespace android {
29 namespace bpf {
30 using ::perfetto::protos::NetworkPacketEvent;
31 using ::perfetto::protos::NetworkPacketTraceConfig;
32 using ::perfetto::protos::Trace;
33 using ::perfetto::protos::TracePacket;
34 using ::perfetto::protos::TrafficDirection;
35 
36 class NetworkTraceHandlerTest : public testing::Test {
37  protected:
38   // Starts a tracing session with the handler under test.
StartTracing(NetworkPacketTraceConfig settings)39   std::unique_ptr<perfetto::TracingSession> StartTracing(
40       NetworkPacketTraceConfig settings) {
41     perfetto::TracingInitArgs args;
42     args.backends = perfetto::kInProcessBackend;
43     perfetto::Tracing::Initialize(args);
44 
45     perfetto::DataSourceDescriptor dsd;
46     dsd.set_name("test.network_packets");
47     NetworkTraceHandler::Register(dsd, /*isTest=*/true);
48 
49     perfetto::TraceConfig cfg;
50     cfg.add_buffers()->set_size_kb(1024);
51     auto* config = cfg.add_data_sources()->mutable_config();
52     config->set_name("test.network_packets");
53     config->set_network_packet_trace_config_raw(settings.SerializeAsString());
54 
55     auto session = perfetto::Tracing::NewTrace(perfetto::kInProcessBackend);
56     session->Setup(cfg);
57     session->StartBlocking();
58     return session;
59   }
60 
61   // Stops the trace session and reports all relevant trace packets.
StopTracing(perfetto::TracingSession * session,std::vector<TracePacket> * output)62   bool StopTracing(perfetto::TracingSession* session,
63                    std::vector<TracePacket>* output) {
64     session->StopBlocking();
65 
66     Trace trace;
67     std::vector<char> raw_trace = session->ReadTraceBlocking();
68     if (!trace.ParseFromArray(raw_trace.data(), raw_trace.size())) {
69       ADD_FAILURE() << "trace.ParseFromArray failed";
70       return false;
71     }
72 
73     // This is a real trace and includes irrelevant trace packets such as trace
74     // metadata. The following strips the results to just the packets we want.
75     for (const auto& pkt : trace.packet()) {
76       if (pkt.has_network_packet() || pkt.has_network_packet_bundle()) {
77         output->emplace_back(pkt);
78       }
79     }
80 
81     return true;
82   }
83 
84   // This runs a trace with a single call to Write.
TraceAndSortPackets(const std::vector<PacketTrace> & input,std::vector<TracePacket> * output,NetworkPacketTraceConfig config={})85   bool TraceAndSortPackets(const std::vector<PacketTrace>& input,
86                            std::vector<TracePacket>* output,
87                            NetworkPacketTraceConfig config = {}) {
88     auto session = StartTracing(config);
__anon461f4bb00102(NetworkTraceHandler::TraceContext ctx) 89     NetworkTraceHandler::Trace([&](NetworkTraceHandler::TraceContext ctx) {
90       ctx.GetDataSourceLocked()->Write(input, ctx);
91       ctx.Flush();
92     });
93 
94     if (!StopTracing(session.get(), output)) {
95       return false;
96     }
97 
98     // Sort to provide deterministic ordering regardless of Perfetto internals
99     // or implementation-defined (e.g. hash map) reshuffling.
100     std::sort(output->begin(), output->end(),
__anon461f4bb00202(const TracePacket& a, const TracePacket& b) 101               [](const TracePacket& a, const TracePacket& b) {
102                 return a.timestamp() < b.timestamp();
103               });
104 
105     return true;
106   }
107 };
108 
TEST_F(NetworkTraceHandlerTest,WriteBasicFields)109 TEST_F(NetworkTraceHandlerTest, WriteBasicFields) {
110   std::vector<PacketTrace> input = {
111       PacketTrace{
112           .timestampNs = 1000,
113           .length = 100,
114           .uid = 10,
115           .tag = 123,
116           .ipProto = IPPROTO_TCP,
117           .tcpFlags = 1,
118       },
119   };
120 
121   std::vector<TracePacket> events;
122   ASSERT_TRUE(TraceAndSortPackets(input, &events));
123 
124   ASSERT_EQ(events.size(), 1);
125   EXPECT_THAT(events[0].timestamp(), 1000);
126   EXPECT_THAT(events[0].network_packet().uid(), 10);
127   EXPECT_THAT(events[0].network_packet().tag(), 123);
128   EXPECT_THAT(events[0].network_packet().ip_proto(), 6);
129   EXPECT_THAT(events[0].network_packet().tcp_flags(), 1);
130   EXPECT_THAT(events[0].network_packet().length(), 100);
131   EXPECT_THAT(events[0].has_sequence_flags(), false);
132 }
133 
TEST_F(NetworkTraceHandlerTest,WriteDirectionAndPorts)134 TEST_F(NetworkTraceHandlerTest, WriteDirectionAndPorts) {
135   std::vector<PacketTrace> input = {
136       PacketTrace{
137           .timestampNs = 1,
138           .sport = htons(8080),
139           .dport = htons(443),
140           .egress = true,
141           .ipProto = IPPROTO_TCP,
142       },
143       PacketTrace{
144           .timestampNs = 2,
145           .sport = htons(443),
146           .dport = htons(8080),
147           .egress = false,
148           .ipProto = IPPROTO_TCP,
149       },
150   };
151 
152   std::vector<TracePacket> events;
153   ASSERT_TRUE(TraceAndSortPackets(input, &events));
154 
155   ASSERT_EQ(events.size(), 2);
156   EXPECT_THAT(events[0].network_packet().local_port(), 8080);
157   EXPECT_THAT(events[0].network_packet().remote_port(), 443);
158   EXPECT_THAT(events[0].network_packet().direction(),
159               TrafficDirection::DIR_EGRESS);
160   EXPECT_THAT(events[1].network_packet().local_port(), 8080);
161   EXPECT_THAT(events[1].network_packet().remote_port(), 443);
162   EXPECT_THAT(events[1].network_packet().direction(),
163               TrafficDirection::DIR_INGRESS);
164 }
165 
TEST_F(NetworkTraceHandlerTest,WriteIcmpTypeAndCode)166 TEST_F(NetworkTraceHandlerTest, WriteIcmpTypeAndCode) {
167   std::vector<PacketTrace> input = {
168       PacketTrace{
169           .timestampNs = 1,
170           .sport = htons(11),  // type
171           .dport = htons(22),  // code
172           .egress = true,
173           .ipProto = IPPROTO_ICMP,
174       },
175       PacketTrace{
176           .timestampNs = 2,
177           .sport = htons(33),  // type
178           .dport = htons(44),  // code
179           .egress = false,
180           .ipProto = IPPROTO_ICMPV6,
181       },
182   };
183 
184   std::vector<TracePacket> events;
185   ASSERT_TRUE(TraceAndSortPackets(input, &events));
186 
187   ASSERT_EQ(events.size(), 2);
188   EXPECT_FALSE(events[0].network_packet().has_local_port());
189   EXPECT_FALSE(events[0].network_packet().has_remote_port());
190   EXPECT_THAT(events[0].network_packet().icmp_type(), 11);
191   EXPECT_THAT(events[0].network_packet().icmp_code(), 22);
192   EXPECT_THAT(events[0].network_packet().direction(),
193               TrafficDirection::DIR_EGRESS);
194   EXPECT_FALSE(events[1].network_packet().local_port());
195   EXPECT_FALSE(events[1].network_packet().remote_port());
196   EXPECT_THAT(events[1].network_packet().icmp_type(), 33);
197   EXPECT_THAT(events[1].network_packet().icmp_code(), 44);
198   EXPECT_THAT(events[1].network_packet().direction(),
199               TrafficDirection::DIR_INGRESS);
200 }
201 
TEST_F(NetworkTraceHandlerTest,BasicBundling)202 TEST_F(NetworkTraceHandlerTest, BasicBundling) {
203   // TODO: remove this once bundling becomes default. Until then, set arbitrary
204   // aggregation threshold to enable bundling.
205   NetworkPacketTraceConfig config;
206   config.set_aggregation_threshold(10);
207 
208   std::vector<PacketTrace> input = {
209       PacketTrace{.timestampNs = 2, .length = 200, .uid = 123},
210       PacketTrace{.timestampNs = 1, .length = 100, .uid = 123},
211       PacketTrace{.timestampNs = 4, .length = 300, .uid = 123},
212 
213       PacketTrace{.timestampNs = 2, .length = 400, .uid = 456},
214       PacketTrace{.timestampNs = 4, .length = 100, .uid = 456},
215   };
216 
217   std::vector<TracePacket> events;
218   ASSERT_TRUE(TraceAndSortPackets(input, &events, config));
219 
220   ASSERT_EQ(events.size(), 2);
221 
222   EXPECT_THAT(events[0].timestamp(), 1);
223   EXPECT_THAT(events[0].network_packet_bundle().ctx().uid(), 123);
224   EXPECT_THAT(events[0].network_packet_bundle().packet_lengths(),
225               testing::ElementsAre(200, 100, 300));
226   EXPECT_THAT(events[0].network_packet_bundle().packet_timestamps(),
227               testing::ElementsAre(1, 0, 3));
228 
229   EXPECT_THAT(events[1].timestamp(), 2);
230   EXPECT_THAT(events[1].network_packet_bundle().ctx().uid(), 456);
231   EXPECT_THAT(events[1].network_packet_bundle().packet_lengths(),
232               testing::ElementsAre(400, 100));
233   EXPECT_THAT(events[1].network_packet_bundle().packet_timestamps(),
234               testing::ElementsAre(0, 2));
235 }
236 
TEST_F(NetworkTraceHandlerTest,AggregationThreshold)237 TEST_F(NetworkTraceHandlerTest, AggregationThreshold) {
238   // With an aggregation threshold of 3, the set of packets with uid=123 will
239   // be aggregated (3>=3) whereas packets with uid=456 get per-packet info.
240   NetworkPacketTraceConfig config;
241   config.set_aggregation_threshold(3);
242 
243   std::vector<PacketTrace> input = {
244       PacketTrace{.timestampNs = 2, .length = 200, .uid = 123},
245       PacketTrace{.timestampNs = 1, .length = 100, .uid = 123},
246       PacketTrace{.timestampNs = 4, .length = 300, .uid = 123},
247 
248       PacketTrace{.timestampNs = 2, .length = 400, .uid = 456},
249       PacketTrace{.timestampNs = 4, .length = 100, .uid = 456},
250   };
251 
252   std::vector<TracePacket> events;
253   ASSERT_TRUE(TraceAndSortPackets(input, &events, config));
254 
255   ASSERT_EQ(events.size(), 2);
256 
257   EXPECT_EQ(events[0].timestamp(), 1);
258   EXPECT_EQ(events[0].network_packet_bundle().ctx().uid(), 123);
259   EXPECT_EQ(events[0].network_packet_bundle().total_duration(), 3);
260   EXPECT_EQ(events[0].network_packet_bundle().total_packets(), 3);
261   EXPECT_EQ(events[0].network_packet_bundle().total_length(), 600);
262 
263   EXPECT_EQ(events[1].timestamp(), 2);
264   EXPECT_EQ(events[1].network_packet_bundle().ctx().uid(), 456);
265   EXPECT_THAT(events[1].network_packet_bundle().packet_lengths(),
266               testing::ElementsAre(400, 100));
267   EXPECT_THAT(events[1].network_packet_bundle().packet_timestamps(),
268               testing::ElementsAre(0, 2));
269 }
270 
TEST_F(NetworkTraceHandlerTest,DropLocalPort)271 TEST_F(NetworkTraceHandlerTest, DropLocalPort) {
272   NetworkPacketTraceConfig config;
273   config.set_drop_local_port(true);
274   config.set_aggregation_threshold(10);
275 
276   __be16 a = htons(10000);
277   __be16 b = htons(10001);
278   std::vector<PacketTrace> input = {
279       // Recall that local is `src` for egress and `dst` for ingress.
280       PacketTrace{.timestampNs = 1, .length = 2, .sport = a, .egress = true},
281       PacketTrace{.timestampNs = 2, .length = 4, .dport = a, .egress = false},
282       PacketTrace{.timestampNs = 3, .length = 6, .sport = b, .egress = true},
283       PacketTrace{.timestampNs = 4, .length = 8, .dport = b, .egress = false},
284   };
285 
286   // Set common fields.
287   for (PacketTrace& pkt : input) {
288     pkt.ipProto = IPPROTO_TCP;
289   }
290 
291   std::vector<TracePacket> events;
292   ASSERT_TRUE(TraceAndSortPackets(input, &events, config));
293   ASSERT_EQ(events.size(), 2);
294 
295   // Despite having different local ports, drop and bundle by remaining fields.
296   EXPECT_EQ(events[0].network_packet_bundle().ctx().direction(),
297             TrafficDirection::DIR_EGRESS);
298   EXPECT_THAT(events[0].network_packet_bundle().packet_lengths(),
299               testing::ElementsAre(2, 6));
300 
301   EXPECT_EQ(events[1].network_packet_bundle().ctx().direction(),
302             TrafficDirection::DIR_INGRESS);
303   EXPECT_THAT(events[1].network_packet_bundle().packet_lengths(),
304               testing::ElementsAre(4, 8));
305 
306   // Local port shouldn't be in output.
307   EXPECT_FALSE(events[0].network_packet_bundle().ctx().has_local_port());
308   EXPECT_FALSE(events[1].network_packet_bundle().ctx().has_local_port());
309 }
310 
TEST_F(NetworkTraceHandlerTest,DropRemotePort)311 TEST_F(NetworkTraceHandlerTest, DropRemotePort) {
312   NetworkPacketTraceConfig config;
313   config.set_drop_remote_port(true);
314   config.set_aggregation_threshold(10);
315 
316   __be16 a = htons(443);
317   __be16 b = htons(80);
318   std::vector<PacketTrace> input = {
319       // Recall that remote is `dst` for egress and `src` for ingress.
320       PacketTrace{.timestampNs = 1, .length = 2, .dport = a, .egress = true},
321       PacketTrace{.timestampNs = 2, .length = 4, .sport = a, .egress = false},
322       PacketTrace{.timestampNs = 3, .length = 6, .dport = b, .egress = true},
323       PacketTrace{.timestampNs = 4, .length = 8, .sport = b, .egress = false},
324   };
325 
326   // Set common fields.
327   for (PacketTrace& pkt : input) {
328     pkt.ipProto = IPPROTO_TCP;
329   }
330 
331   std::vector<TracePacket> events;
332   ASSERT_TRUE(TraceAndSortPackets(input, &events, config));
333   ASSERT_EQ(events.size(), 2);
334 
335   // Despite having different remote ports, drop and bundle by remaining fields.
336   EXPECT_EQ(events[0].network_packet_bundle().ctx().direction(),
337             TrafficDirection::DIR_EGRESS);
338   EXPECT_THAT(events[0].network_packet_bundle().packet_lengths(),
339               testing::ElementsAre(2, 6));
340 
341   EXPECT_EQ(events[1].network_packet_bundle().ctx().direction(),
342             TrafficDirection::DIR_INGRESS);
343   EXPECT_THAT(events[1].network_packet_bundle().packet_lengths(),
344               testing::ElementsAre(4, 8));
345 
346   // Remote port shouldn't be in output.
347   EXPECT_FALSE(events[0].network_packet_bundle().ctx().has_remote_port());
348   EXPECT_FALSE(events[1].network_packet_bundle().ctx().has_remote_port());
349 }
350 
TEST_F(NetworkTraceHandlerTest,DropTcpFlags)351 TEST_F(NetworkTraceHandlerTest, DropTcpFlags) {
352   NetworkPacketTraceConfig config;
353   config.set_drop_tcp_flags(true);
354   config.set_aggregation_threshold(10);
355 
356   std::vector<PacketTrace> input = {
357       PacketTrace{.timestampNs = 1, .length = 1, .uid = 123, .tcpFlags = 1},
358       PacketTrace{.timestampNs = 2, .length = 2, .uid = 123, .tcpFlags = 2},
359       PacketTrace{.timestampNs = 3, .length = 3, .uid = 456, .tcpFlags = 1},
360       PacketTrace{.timestampNs = 4, .length = 4, .uid = 456, .tcpFlags = 2},
361   };
362 
363   // Set common fields.
364   for (PacketTrace& pkt : input) {
365     pkt.ipProto = IPPROTO_TCP;
366   }
367 
368   std::vector<TracePacket> events;
369   ASSERT_TRUE(TraceAndSortPackets(input, &events, config));
370 
371   ASSERT_EQ(events.size(), 2);
372 
373   // Despite having different tcp flags, drop and bundle by remaining fields.
374   EXPECT_EQ(events[0].network_packet_bundle().ctx().uid(), 123);
375   EXPECT_THAT(events[0].network_packet_bundle().packet_lengths(),
376               testing::ElementsAre(1, 2));
377 
378   EXPECT_EQ(events[1].network_packet_bundle().ctx().uid(), 456);
379   EXPECT_THAT(events[1].network_packet_bundle().packet_lengths(),
380               testing::ElementsAre(3, 4));
381 
382   // Tcp flags shouldn't be in output.
383   EXPECT_FALSE(events[0].network_packet_bundle().ctx().has_tcp_flags());
384   EXPECT_FALSE(events[1].network_packet_bundle().ctx().has_tcp_flags());
385 }
386 
TEST_F(NetworkTraceHandlerTest,Interning)387 TEST_F(NetworkTraceHandlerTest, Interning) {
388   NetworkPacketTraceConfig config;
389   config.set_intern_limit(2);
390 
391   // The test writes 4 packets coming from three sources (uids). With an intern
392   // limit of 2, the first two sources should be interned. This test splits this
393   // into individual writes since internally an unordered map is used and would
394   // otherwise non-deterministically choose what to intern (this is fine for
395   // real use, but not good for test assertions).
396   std::vector<std::vector<PacketTrace>> inputs = {
397       {PacketTrace{.timestampNs = 1, .uid = 123}},
398       {PacketTrace{.timestampNs = 2, .uid = 456}},
399       {PacketTrace{.timestampNs = 3, .uid = 789}},
400       {PacketTrace{.timestampNs = 4, .uid = 123}},
401   };
402 
403   auto session = StartTracing(config);
404 
405   NetworkTraceHandler::Trace([&](NetworkTraceHandler::TraceContext ctx) {
406     ctx.GetDataSourceLocked()->Write(inputs[0], ctx);
407     ctx.GetDataSourceLocked()->Write(inputs[1], ctx);
408     ctx.GetDataSourceLocked()->Write(inputs[2], ctx);
409     ctx.GetDataSourceLocked()->Write(inputs[3], ctx);
410     ctx.Flush();
411   });
412 
413   std::vector<TracePacket> events;
414   ASSERT_TRUE(StopTracing(session.get(), &events));
415 
416   ASSERT_EQ(events.size(), 4);
417 
418   // First time seen, emit new interned data, bundle uses iid instead of ctx.
419   EXPECT_EQ(events[0].network_packet_bundle().iid(), 1);
420   ASSERT_EQ(events[0].interned_data().packet_context().size(), 1);
421   EXPECT_EQ(events[0].interned_data().packet_context(0).iid(), 1);
422   EXPECT_EQ(events[0].interned_data().packet_context(0).ctx().uid(), 123);
423   EXPECT_EQ(events[0].sequence_flags(),
424             TracePacket::SEQ_INCREMENTAL_STATE_CLEARED);
425 
426   // First time seen, emit new interned data, bundle uses iid instead of ctx.
427   EXPECT_EQ(events[1].network_packet_bundle().iid(), 2);
428   ASSERT_EQ(events[1].interned_data().packet_context().size(), 1);
429   EXPECT_EQ(events[1].interned_data().packet_context(0).iid(), 2);
430   EXPECT_EQ(events[1].interned_data().packet_context(0).ctx().uid(), 456);
431   EXPECT_EQ(events[1].sequence_flags(),
432             TracePacket::SEQ_NEEDS_INCREMENTAL_STATE);
433 
434   // Not enough room in intern table (limit 2), inline the context.
435   EXPECT_EQ(events[2].network_packet_bundle().ctx().uid(), 789);
436   EXPECT_EQ(events[2].interned_data().packet_context().size(), 0);
437   EXPECT_EQ(events[2].sequence_flags(),
438             TracePacket::SEQ_NEEDS_INCREMENTAL_STATE);
439 
440   // Second time seen, no need to re-emit interned data, only record iid.
441   EXPECT_EQ(events[3].network_packet_bundle().iid(), 1);
442   EXPECT_EQ(events[3].interned_data().packet_context().size(), 0);
443   EXPECT_EQ(events[3].sequence_flags(),
444             TracePacket::SEQ_NEEDS_INCREMENTAL_STATE);
445 }
446 
447 }  // namespace bpf
448 }  // namespace android
449