Merge "gn2bp: Fix some protos include_dirs"
diff --git a/Cronet/tests/cts/Android.bp b/Cronet/tests/cts/Android.bp
index d969b54..7b440cd 100644
--- a/Cronet/tests/cts/Android.bp
+++ b/Cronet/tests/cts/Android.bp
@@ -64,6 +64,7 @@
         "junit",
         "hamcrest-library",
         "kotlin-test",
+        "mockito-target",
     ],
     libs: [
         "android.test.base",
diff --git a/Cronet/tests/cts/src/android/net/http/cts/HttpEngineTest.java b/Cronet/tests/cts/src/android/net/http/cts/HttpEngineTest.java
index d247201..816596c 100644
--- a/Cronet/tests/cts/src/android/net/http/cts/HttpEngineTest.java
+++ b/Cronet/tests/cts/src/android/net/http/cts/HttpEngineTest.java
@@ -29,6 +29,7 @@
 import static org.junit.Assert.assertTrue;
 
 import android.content.Context;
+import android.net.Network;
 import android.net.http.HttpEngine;
 import android.net.http.UrlRequest;
 import android.net.http.UrlResponseInfo;
@@ -43,6 +44,7 @@
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.Mockito;
 
 @RunWith(AndroidJUnit4.class)
 public class HttpEngineTest {
@@ -51,6 +53,7 @@
 
     private HttpEngine.Builder mEngineBuilder;
     private TestUrlRequestCallback mCallback;
+    private HttpCtsTestServer mTestServer;
     private UrlRequest mRequest;
     private HttpEngine mEngine;
     private Context mContext;
@@ -61,6 +64,7 @@
         skipIfNoInternetConnection(mContext);
         mEngineBuilder = new HttpEngine.Builder(mContext);
         mCallback = new TestUrlRequestCallback();
+        mTestServer = new HttpCtsTestServer(mContext);
     }
 
     @After
@@ -72,6 +76,9 @@
         if (mEngine != null) {
             mEngine.shutdown();
         }
+        if (mTestServer != null) {
+            mTestServer.shutdown();
+        }
     }
 
     private boolean isQuic(String negotiatedProtocol) {
@@ -253,4 +260,45 @@
                 .replaceFirst(".*<title>", "")
                 .replaceFirst("</title>.*", "");
     }
-}
+
+    @Test
+    public void testHttpEngine_bindToNetwork() throws Exception {
+        // Create a fake Android.net.Network. Since that network doesn't exist, binding to
+        // that should end up in a failed request.
+        Network mockNetwork = Mockito.mock(Network.class);
+        Mockito.when(mockNetwork.getNetworkHandle()).thenReturn(123L);
+        String url = mTestServer.getSuccessUrl();
+
+        mEngine = mEngineBuilder.build();
+        mEngine.bindToNetwork(mockNetwork);
+        UrlRequest.Builder builder =
+                mEngine.newUrlRequestBuilder(url, mCallback.getExecutor(), mCallback);
+        mRequest = builder.build();
+        mRequest.start();
+
+        mCallback.expectCallback(ResponseStep.ON_FAILED);
+    }
+
+    @Test
+    public void testHttpEngine_unbindFromNetwork() throws Exception {
+        // Create a fake Android.net.Network. Since that network doesn't exist, binding to
+        // that should end up in a failed request.
+        Network mockNetwork = Mockito.mock(Network.class);
+        Mockito.when(mockNetwork.getNetworkHandle()).thenReturn(123L);
+        String url = mTestServer.getSuccessUrl();
+
+        mEngine = mEngineBuilder.build();
+        // Bind to the fake network but then unbind. This should result in a successful
+        // request.
+        mEngine.bindToNetwork(mockNetwork);
+        mEngine.bindToNetwork(null);
+        UrlRequest.Builder builder =
+                mEngine.newUrlRequestBuilder(url, mCallback.getExecutor(), mCallback);
+        mRequest = builder.build();
+        mRequest.start();
+
+        mCallback.expectCallback(ResponseStep.ON_SUCCEEDED);
+        UrlResponseInfo info = mCallback.mResponseInfo;
+        assertOKStatusCode(info);
+    }
+}
\ No newline at end of file
diff --git a/Cronet/tests/cts/src/android/net/http/cts/UrlRequestTest.java b/Cronet/tests/cts/src/android/net/http/cts/UrlRequestTest.java
index a364e29..2ec035b 100644
--- a/Cronet/tests/cts/src/android/net/http/cts/UrlRequestTest.java
+++ b/Cronet/tests/cts/src/android/net/http/cts/UrlRequestTest.java
@@ -44,6 +44,8 @@
 import android.net.http.cts.util.TestUrlRequestCallback.ResponseStep;
 import android.net.http.cts.util.UploadDataProviders;
 
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 import androidx.test.core.app.ApplicationProvider;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 
@@ -310,7 +312,7 @@
         assertThat(mCallback.mResponseAsString).isEqualTo(body);
     }
 
-    private static class StubUrlRequestCallback extends UrlRequest.Callback {
+    private static class StubUrlRequestCallback implements UrlRequest.Callback {
 
         @Override
         public void onRedirectReceived(
@@ -338,6 +340,11 @@
         public void onFailed(UrlRequest request, UrlResponseInfo info, HttpException error) {
             throw new UnsupportedOperationException(error);
         }
+
+        @Override
+        public void onCanceled(@NonNull UrlRequest request, @Nullable UrlResponseInfo info) {
+            throw new UnsupportedOperationException();
+        }
     }
 
     private static class InMemoryUploadDataProvider extends UploadDataProvider {
diff --git a/Cronet/tests/cts/src/android/net/http/cts/util/TestUrlRequestCallback.java b/Cronet/tests/cts/src/android/net/http/cts/util/TestUrlRequestCallback.java
index efbcff6..28443b7 100644
--- a/Cronet/tests/cts/src/android/net/http/cts/util/TestUrlRequestCallback.java
+++ b/Cronet/tests/cts/src/android/net/http/cts/util/TestUrlRequestCallback.java
@@ -50,7 +50,7 @@
  * method to block thread until the request completes on another thread.
  * Allows us to cancel, block request or throw an exception from an arbitrary step.
  */
-public class TestUrlRequestCallback extends UrlRequest.Callback {
+public class TestUrlRequestCallback implements UrlRequest.Callback {
     private static final int TIMEOUT_MS = 12_000;
     public ArrayList<UrlResponseInfo> mRedirectResponseInfoList = new ArrayList<>();
     public ArrayList<String> mRedirectUrlList = new ArrayList<>();
diff --git a/Cronet/tools/import/copy.bara.sky b/Cronet/tools/import/copy.bara.sky
index 2acf8cd..c066fc0 100644
--- a/Cronet/tools/import/copy.bara.sky
+++ b/Cronet/tools/import/copy.bara.sky
@@ -19,6 +19,7 @@
 
     # Exclude existing *OWNERS files
     "**/*OWNERS",
+    "**/.git/**",
 ]
 
 cronet_origin_files = glob(
@@ -35,6 +36,8 @@
         "crypto/**",
         "ipc/**",
         "net/**",
+        # Note: Only used for tests.
+        "testing/**",
         "url/**",
         "LICENSE",
     ],
@@ -48,6 +51,8 @@
         "components/cronet/ios/**",
         "components/cronet/native/**",
 
+        # Per aosp/2399270
+        "testing/buildbot/**",
 
         # Exclude all third-party directories. Those are specified explicitly
         # below, so no dependency can accidentally creep in.
@@ -69,10 +74,10 @@
         "base/third_party/symbolize/**",
         "base/third_party/valgrind/**",
         "base/third_party/xdg_user_dirs/**",
-        # Not present in source repo; requires gclient sync.
         "buildtools/third_party/libc++/**",
-        # Not present in source repo; requires gclient sync.
         "buildtools/third_party/libc++abi/**",
+        # Note: Only used for tests.
+        "net/third_party/nist-pkits/**",
         "net/third_party/quiche/**",
         "net/third_party/uri_template/**",
         "third_party/abseil-cpp/**",
@@ -80,12 +85,21 @@
         "third_party/ashmem/**",
         "third_party/boringssl/**",
         "third_party/brotli/**",
-        # Not present in source repo; requires gclient sync.
+        # Note: Only used for tests.
+        "third_party/ced/**",
+        # Note: Only used for tests.
+        "third_party/googletest/**",
         "third_party/icu/**",
         "third_party/libevent/**",
+        # Note: Only used for tests.
+        "third_party/libxml/**",
+        # Note: Only used for tests.
+        "third_party/lss/**",
         "third_party/metrics_proto/**",
         "third_party/modp_b64/**",
         "third_party/protobuf/**",
+        # Note: Only used for tests.
+        "third_party/quic_trace/**",
         "third_party/zlib/**",
     ],
     exclude = common_excludes,
@@ -94,12 +108,8 @@
 core.workflow(
     name = "import_cronet",
     authoring = authoring.overwrite("Cronet Mainline Eng <cronet-mainline-eng+copybara@google.com>"),
-    origin = git.origin(
-        url = "rpc://chromium/chromium/src",
-        # Source ref is set by the invoking script.
-        ref = "overwritten-by-script",
-        partial_fetch = True,
-    ),
+    # Origin folder is specified via source_ref argument, see import_cronet.sh
+    origin = folder.origin(),
     origin_files = cronet_origin_files,
     destination = git.destination(
         # The destination URL is set by the invoking script.
diff --git a/Cronet/tools/import/import_cronet.sh b/Cronet/tools/import/import_cronet.sh
index eb82551..d0c8deb 100755
--- a/Cronet/tools/import/import_cronet.sh
+++ b/Cronet/tools/import/import_cronet.sh
@@ -33,6 +33,8 @@
     exit 1
 }
 
+COPYBARA_FOLDER_ORIGIN="/tmp/copybara-origin"
+
 #######################################
 # Create upstream-import branch in external/cronet.
 # Globals:
@@ -49,22 +51,59 @@
 }
 
 #######################################
+# Setup folder.origin for copybara inside /tmp
+# Globals:
+#   COPYBARA_FOLDER_ORIGIN
+# Arguments:
+#   new_rev, string
+#######################################
+setup_folder_origin() {
+    local _new_rev=$1
+    mkdir -p "${COPYBARA_FOLDER_ORIGIN}"
+    cd "${COPYBARA_FOLDER_ORIGIN}"
+
+    # For this to work _new_rev must be a branch or a tag.
+    git clone --depth=1 --branch "${_new_rev}" https://chromium.googlesource.com/chromium/src.git
+
+    cat <<EOF >.gclient
+solutions = [
+  {
+    "name": "src",
+    "url": "https://chromium.googlesource.com/chromium/src.git",
+    "managed": False,
+    "custom_deps": {},
+    "custom_vars": {},
+  },
+]
+target_os = ["android"]
+EOF
+    cd src
+    # Set appropriate gclient flags to speed up syncing.
+    gclient sync \
+        --no-history
+        --shallow
+}
+
+#######################################
 # Runs the copybara import of Chromium
 # Globals:
 #   ANDROID_BUILD_TOP
+#   COPYBARA_FOLDER_ORIGIN
 # Arguments:
-#   new_rev, string
 #   last_rev, string or empty
 #   force, string or empty
 #######################################
 do_run_copybara() {
-    local _new_rev=$1
-    local _last_rev=$2
-    local _force=$3
+    local _last_rev=$1
+    local _force=$2
 
     local -a flags
     flags+=(--git-destination-url="file://${ANDROID_BUILD_TOP}/external/cronet")
-    flags+=(--repo-timeout 3h)
+    flags+=(--repo-timeout 3m)
+
+    # buildtools/third_party/libc++ contains an invalid symlink
+    flags+=(--folder-origin-ignore-invalid-symlinks)
+    flags+=(--git-no-verify)
 
     if [ ! -z "${_force}" ]; then
         flags+=(--force)
@@ -77,7 +116,7 @@
     /google/bin/releases/copybara/public/copybara/copybara \
         "${flags[@]}" \
         "${ANDROID_BUILD_TOP}/packages/modules/Connectivity/Cronet/tools/import/copy.bara.sky" \
-        import_cronet "${_new_rev}"
+        import_cronet "${COPYBARA_FOLDER_ORIGIN}/src"
 }
 
 while getopts $OPTSTRING opt; do
@@ -96,5 +135,6 @@
 fi
 
 setup_upstream_import_branch
-do_run_copybara "${new_rev}" "${last_rev}" "${force}"
+setup_folder_origin "${new_rev}"
+do_run_copybara "${last_rev}" "${force}"
 
diff --git a/Tethering/common/TetheringLib/cronet_enabled/api/current.txt b/Tethering/common/TetheringLib/cronet_enabled/api/current.txt
index cf441ce..66a0295 100644
--- a/Tethering/common/TetheringLib/cronet_enabled/api/current.txt
+++ b/Tethering/common/TetheringLib/cronet_enabled/api/current.txt
@@ -241,14 +241,13 @@
     method @NonNull public abstract android.net.http.UrlRequest.Builder setUploadDataProvider(@NonNull android.net.http.UploadDataProvider, @NonNull java.util.concurrent.Executor);
   }
 
-  public abstract static class UrlRequest.Callback {
-    ctor public UrlRequest.Callback();
+  public static interface UrlRequest.Callback {
     method public void onCanceled(@NonNull android.net.http.UrlRequest, @Nullable android.net.http.UrlResponseInfo);
-    method public abstract void onFailed(@NonNull android.net.http.UrlRequest, @Nullable android.net.http.UrlResponseInfo, @NonNull android.net.http.HttpException);
-    method public abstract void onReadCompleted(@NonNull android.net.http.UrlRequest, @NonNull android.net.http.UrlResponseInfo, @NonNull java.nio.ByteBuffer) throws java.lang.Exception;
-    method public abstract void onRedirectReceived(@NonNull android.net.http.UrlRequest, @NonNull android.net.http.UrlResponseInfo, @NonNull String) throws java.lang.Exception;
-    method public abstract void onResponseStarted(@NonNull android.net.http.UrlRequest, @NonNull android.net.http.UrlResponseInfo) throws java.lang.Exception;
-    method public abstract void onSucceeded(@NonNull android.net.http.UrlRequest, @NonNull android.net.http.UrlResponseInfo);
+    method public void onFailed(@NonNull android.net.http.UrlRequest, @Nullable android.net.http.UrlResponseInfo, @NonNull android.net.http.HttpException);
+    method public void onReadCompleted(@NonNull android.net.http.UrlRequest, @NonNull android.net.http.UrlResponseInfo, @NonNull java.nio.ByteBuffer) throws java.lang.Exception;
+    method public void onRedirectReceived(@NonNull android.net.http.UrlRequest, @NonNull android.net.http.UrlResponseInfo, @NonNull String) throws java.lang.Exception;
+    method public void onResponseStarted(@NonNull android.net.http.UrlRequest, @NonNull android.net.http.UrlResponseInfo) throws java.lang.Exception;
+    method public void onSucceeded(@NonNull android.net.http.UrlRequest, @NonNull android.net.http.UrlResponseInfo);
   }
 
   public static class UrlRequest.Status {
diff --git a/service-t/native/libs/libnetworkstats/Android.bp b/service-t/native/libs/libnetworkstats/Android.bp
index f2c569f..f40d388 100644
--- a/service-t/native/libs/libnetworkstats/Android.bp
+++ b/service-t/native/libs/libnetworkstats/Android.bp
@@ -62,6 +62,7 @@
     header_libs: ["bpf_connectivity_headers"],
     srcs: [
         "BpfNetworkStatsTest.cpp",
+        "NetworkTraceHandlerTest.cpp",
         "NetworkTracePollerTest.cpp",
     ],
     cflags: [
@@ -74,6 +75,8 @@
         "libgmock",
         "libnetworkstats",
         "libperfetto_client_experimental",
+        "libprotobuf-cpp-lite",
+        "perfetto_trace_protos",
     ],
     shared_libs: [
         "libbase",
diff --git a/service-t/native/libs/libnetworkstats/NetworkTraceHandler.cpp b/service-t/native/libs/libnetworkstats/NetworkTraceHandler.cpp
index 8e70950..cd62bc5 100644
--- a/service-t/native/libs/libnetworkstats/NetworkTraceHandler.cpp
+++ b/service-t/native/libs/libnetworkstats/NetworkTraceHandler.cpp
@@ -57,11 +57,16 @@
 }
 
 // static
-NetworkTracePoller NetworkTraceHandler::sPoller([](const PacketTrace& pkt) {
-  NetworkTraceHandler::Trace([pkt](NetworkTraceHandler::TraceContext ctx) {
-    NetworkTraceHandler::Fill(pkt, *ctx.NewTracePacket());
-  });
-});
+NetworkTracePoller NetworkTraceHandler::sPoller(
+    [](const std::vector<PacketTrace>& packets) {
+      // Trace calls the provided callback for each active session. The context
+      // gets a reference to the NetworkTraceHandler instance associated with
+      // the session and delegates writing. The corresponding handler will write
+      // with the setting specified in the trace config.
+      NetworkTraceHandler::Trace([&](NetworkTraceHandler::TraceContext ctx) {
+        ctx.GetDataSourceLocked()->Write(packets, ctx);
+      });
+    });
 
 void NetworkTraceHandler::OnSetup(const SetupArgs& args) {
   const std::string& raw = args.config->network_packet_trace_config_raw();
@@ -72,6 +77,12 @@
     ALOGI("poll_ms is missing or below the 100ms minimum. Increasing to 100ms");
     mPollMs = 100;
   }
+
+  mInternLimit = config.intern_limit();
+  mAggregationThreshold = config.aggregation_threshold();
+  mDropLocalPort = config.drop_local_port();
+  mDropRemotePort = config.drop_remote_port();
+  mDropTcpFlags = config.drop_tcp_flags();
 }
 
 void NetworkTraceHandler::OnStart(const StartArgs&) {
@@ -83,6 +94,13 @@
   mStarted = false;
 }
 
+void NetworkTraceHandler::Write(const std::vector<PacketTrace>& packets,
+                                NetworkTraceHandler::TraceContext& ctx) {
+  for (const PacketTrace& pkt : packets) {
+    Fill(pkt, *ctx.NewTracePacket());
+  }
+}
+
 // static class method
 void NetworkTraceHandler::Fill(const PacketTrace& src, TracePacket& dst) {
   dst.set_timestamp(src.timestampNs);
diff --git a/service-t/native/libs/libnetworkstats/NetworkTraceHandlerTest.cpp b/service-t/native/libs/libnetworkstats/NetworkTraceHandlerTest.cpp
new file mode 100644
index 0000000..2318da5
--- /dev/null
+++ b/service-t/native/libs/libnetworkstats/NetworkTraceHandlerTest.cpp
@@ -0,0 +1,166 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+#include <vector>
+
+#include "netdbpf/NetworkTraceHandler.h"
+#include "protos/perfetto/trace/android/network_trace.pb.h"
+#include "protos/perfetto/trace/trace.pb.h"
+#include "protos/perfetto/trace/trace_packet.pb.h"
+
+namespace android {
+namespace bpf {
+using ::perfetto::protos::NetworkPacketEvent;
+using ::perfetto::protos::Trace;
+using ::perfetto::protos::TracePacket;
+using ::perfetto::protos::TrafficDirection;
+
+// This handler makes OnStart and OnStop a no-op so that tracing is not really
+// started on the device.
+class HandlerForTest : public NetworkTraceHandler {
+ public:
+  void OnStart(const StartArgs&) override {}
+  void OnStop(const StopArgs&) override {}
+};
+
+class NetworkTraceHandlerTest : public testing::Test {
+ protected:
+  // Starts a tracing session with the handler under test.
+  std::unique_ptr<perfetto::TracingSession> StartTracing() {
+    perfetto::TracingInitArgs args;
+    args.backends = perfetto::kInProcessBackend;
+    perfetto::Tracing::Initialize(args);
+
+    perfetto::DataSourceDescriptor dsd;
+    dsd.set_name("test.network_packets");
+    HandlerForTest::Register(dsd);
+
+    perfetto::TraceConfig cfg;
+    cfg.add_buffers()->set_size_kb(1024);
+    cfg.add_data_sources()->mutable_config()->set_name("test.network_packets");
+
+    auto session = perfetto::Tracing::NewTrace(perfetto::kInProcessBackend);
+    session->Setup(cfg);
+    session->StartBlocking();
+    return session;
+  }
+
+  // Stops the trace session and reports all relevant trace packets.
+  bool StopTracing(perfetto::TracingSession* session,
+                   std::vector<TracePacket>* output) {
+    session->StopBlocking();
+
+    Trace trace;
+    std::vector<char> raw_trace = session->ReadTraceBlocking();
+    if (!trace.ParseFromArray(raw_trace.data(), raw_trace.size())) {
+      ADD_FAILURE() << "trace.ParseFromArray failed";
+      return false;
+    }
+
+    // This is a real trace and includes irrelevant trace packets such as trace
+    // metadata. The following strips the results to just the packets we want.
+    for (const auto& pkt : trace.packet()) {
+      if (pkt.has_network_packet()) {
+        output->emplace_back(pkt);
+      }
+    }
+
+    return true;
+  }
+
+  // This runs a trace with a single call to Write.
+  bool TraceAndSortPackets(const std::vector<PacketTrace>& input,
+                           std::vector<TracePacket>* output) {
+    auto session = StartTracing();
+    HandlerForTest::Trace([&](HandlerForTest::TraceContext ctx) {
+      ctx.GetDataSourceLocked()->Write(input, ctx);
+      ctx.Flush();
+    });
+
+    if (!StopTracing(session.get(), output)) {
+      return false;
+    }
+
+    // Sort to provide deterministic ordering regardless of Perfetto internals
+    // or implementation-defined (e.g. hash map) reshuffling.
+    std::sort(output->begin(), output->end(),
+              [](const TracePacket& a, const TracePacket& b) {
+                return a.timestamp() < b.timestamp();
+              });
+
+    return true;
+  }
+};
+
+TEST_F(NetworkTraceHandlerTest, WriteBasicFields) {
+  std::vector<PacketTrace> input = {
+      PacketTrace{
+          .timestampNs = 1000,
+          .length = 100,
+          .uid = 10,
+          .tag = 123,
+          .ipProto = 6,
+          .tcpFlags = 1,
+      },
+  };
+
+  std::vector<TracePacket> events;
+  ASSERT_TRUE(TraceAndSortPackets(input, &events));
+
+  ASSERT_EQ(events.size(), 1);
+  EXPECT_THAT(events[0].timestamp(), 1000);
+  EXPECT_THAT(events[0].network_packet().length(), 100);
+  EXPECT_THAT(events[0].network_packet().uid(), 10);
+  EXPECT_THAT(events[0].network_packet().tag(), 123);
+  EXPECT_THAT(events[0].network_packet().ip_proto(), 6);
+  EXPECT_THAT(events[0].network_packet().tcp_flags(), 1);
+}
+
+TEST_F(NetworkTraceHandlerTest, WriteDirectionAndPorts) {
+  std::vector<PacketTrace> input = {
+      PacketTrace{
+          .timestampNs = 1,
+          .sport = htons(8080),
+          .dport = htons(443),
+          .egress = true,
+      },
+      PacketTrace{
+          .timestampNs = 2,
+          .sport = htons(443),
+          .dport = htons(8080),
+          .egress = false,
+      },
+  };
+
+  std::vector<TracePacket> events;
+  ASSERT_TRUE(TraceAndSortPackets(input, &events));
+
+  ASSERT_EQ(events.size(), 2);
+  EXPECT_THAT(events[0].network_packet().local_port(), 8080);
+  EXPECT_THAT(events[0].network_packet().remote_port(), 443);
+  EXPECT_THAT(events[0].network_packet().direction(),
+              TrafficDirection::DIR_EGRESS);
+  EXPECT_THAT(events[1].network_packet().local_port(), 8080);
+  EXPECT_THAT(events[1].network_packet().remote_port(), 443);
+  EXPECT_THAT(events[1].network_packet().direction(),
+              TrafficDirection::DIR_INGRESS);
+}
+
+}  // namespace bpf
+}  // namespace android
diff --git a/service-t/native/libs/libnetworkstats/NetworkTracePoller.cpp b/service-t/native/libs/libnetworkstats/NetworkTracePoller.cpp
index 34dbf9e..3abb49a 100644
--- a/service-t/native/libs/libnetworkstats/NetworkTracePoller.cpp
+++ b/service-t/native/libs/libnetworkstats/NetworkTracePoller.cpp
@@ -116,12 +116,16 @@
     return false;
   }
 
-  base::Result<int> ret = mRingBuffer->ConsumeAll(mCallback);
+  std::vector<PacketTrace> packets;
+  base::Result<int> ret = mRingBuffer->ConsumeAll(
+      [&](const PacketTrace& pkt) { packets.push_back(pkt); });
   if (!ret.ok()) {
     ALOGW("Failed to poll ringbuf: %s", ret.error().message().c_str());
     return false;
   }
 
+  mCallback(packets);
+
   return true;
 }
 
diff --git a/service-t/native/libs/libnetworkstats/NetworkTracePollerTest.cpp b/service-t/native/libs/libnetworkstats/NetworkTracePollerTest.cpp
index 28ec208..725cec1 100644
--- a/service-t/native/libs/libnetworkstats/NetworkTracePollerTest.cpp
+++ b/service-t/native/libs/libnetworkstats/NetworkTracePollerTest.cpp
@@ -100,7 +100,7 @@
 };
 
 TEST_F(NetworkTracePollerTest, PollWhileInactive) {
-  NetworkTracePoller handler([&](const PacketTrace& pkt) {});
+  NetworkTracePoller handler([&](const std::vector<PacketTrace>& pkt) {});
 
   // One succeed after start and before stop.
   EXPECT_FALSE(handler.ConsumeAll());
@@ -113,7 +113,7 @@
 TEST_F(NetworkTracePollerTest, ConcurrentSessions) {
   // Simulate two concurrent sessions (two starts followed by two stops). Check
   // that tracing is stopped only after both sessions finish.
-  NetworkTracePoller handler([&](const PacketTrace& pkt) {});
+  NetworkTracePoller handler([&](const std::vector<PacketTrace>& pkt) {});
 
   ASSERT_TRUE(handler.Start(kNeverPoll));
   EXPECT_TRUE(handler.ConsumeAll());
@@ -135,10 +135,12 @@
   // Record all packets with the bound address and current uid. This callback is
   // involked only within ConsumeAll, at which point the port should have
   // already been filled in and all packets have been processed.
-  NetworkTracePoller handler([&](const PacketTrace& pkt) {
-    if (pkt.sport != server_port && pkt.dport != server_port) return;
-    if (pkt.uid != getuid()) return;
-    packets.push_back(pkt);
+  NetworkTracePoller handler([&](const std::vector<PacketTrace>& pkts) {
+    for (const PacketTrace& pkt : pkts) {
+      if (pkt.sport != server_port && pkt.dport != server_port) return;
+      if (pkt.uid != getuid()) return;
+      packets.push_back(pkt);
+    }
   });
 
   ASSERT_TRUE(handler.Start(kNeverPoll));
diff --git a/service-t/native/libs/libnetworkstats/include/netdbpf/NetworkTraceHandler.h b/service-t/native/libs/libnetworkstats/include/netdbpf/NetworkTraceHandler.h
index 1266237..9520d7b 100644
--- a/service-t/native/libs/libnetworkstats/include/netdbpf/NetworkTraceHandler.h
+++ b/service-t/native/libs/libnetworkstats/include/netdbpf/NetworkTraceHandler.h
@@ -46,14 +46,26 @@
   void OnStart(const StartArgs&) override;
   void OnStop(const StopArgs&) override;
 
+  // Writes the packets as Perfetto TracePackets, creating packets as needed
+  // using the provided callback (which allows easy testing).
+  void Write(const std::vector<PacketTrace>& packets,
+             NetworkTraceHandler::TraceContext& ctx);
+
  private:
   // Convert a PacketTrace into a Perfetto trace packet.
   static void Fill(const PacketTrace& src,
                    ::perfetto::protos::pbzero::TracePacket& dst);
 
   static internal::NetworkTracePoller sPoller;
-  uint32_t mPollMs;
   bool mStarted;
+
+  // Values from config, see proto for details.
+  uint32_t mPollMs;
+  uint32_t mInternLimit;
+  uint32_t mAggregationThreshold;
+  bool mDropLocalPort;
+  bool mDropRemotePort;
+  bool mDropTcpFlags;
 };
 
 }  // namespace bpf
diff --git a/service-t/native/libs/libnetworkstats/include/netdbpf/NetworkTracePoller.h b/service-t/native/libs/libnetworkstats/include/netdbpf/NetworkTracePoller.h
index b0189a7..adde51e 100644
--- a/service-t/native/libs/libnetworkstats/include/netdbpf/NetworkTracePoller.h
+++ b/service-t/native/libs/libnetworkstats/include/netdbpf/NetworkTracePoller.h
@@ -38,9 +38,10 @@
 // it is not meant to be used elsewhere.
 class NetworkTracePoller {
  public:
+  using EventSink = std::function<void(const std::vector<PacketTrace>&)>;
+
   // Testonly: initialize with a callback capable of intercepting data.
-  NetworkTracePoller(std::function<void(const PacketTrace&)> callback)
-      : mCallback(std::move(callback)) {}
+  NetworkTracePoller(EventSink callback) : mCallback(std::move(callback)) {}
 
   // Starts tracing with the given poll interval.
   bool Start(uint32_t pollMs) EXCLUDES(mMutex);
@@ -67,7 +68,7 @@
   uint32_t mPollMs GUARDED_BY(mMutex);
 
   // The function to process PacketTrace, typically a Perfetto sink.
-  std::function<void(const PacketTrace&)> mCallback GUARDED_BY(mMutex);
+  EventSink mCallback GUARDED_BY(mMutex);
 
   // The BPF ring buffer handle.
   std::unique_ptr<BpfRingbuf<PacketTrace>> mRingBuffer GUARDED_BY(mMutex);