Merge "Trace Redaction - "PID connected to UID" query" into main

12 files changed, 71 insertions, 109 deletions

diff --git a/src/trace_redaction/collect_frame_cookies.cc b/src/trace_redaction/collect_frame_cookies.cc
index 13e553975..04c5eb3c5 100644
--- a/src/trace_redaction/collect_frame_cookies.cc
+++ b/src/trace_redaction/collect_frame_cookies.cc
@@ -148,18 +148,12 @@ base::Status ReduceFrameCookies::Build(Context* context) const {
     return base::OkStatus();
   }
 
-  const auto* timeline = context->timeline.get();
-  auto uid = context->package_uid.value();
-
-  auto& package_frame_cookies = context->package_frame_cookies;
-
   // Filter the global cookies down to cookies that belong to the target package
   // (uid).
   for (const auto& cookie : context->global_frame_cookies) {
-    auto cookie_slice = timeline->Search(cookie.ts, cookie.pid);
-
-    if (cookie_slice.uid == uid) {
-      package_frame_cookies.insert(cookie.cookie);
+    if (context->timeline->PidConnectsToUid(cookie.ts, cookie.pid,
+                                            *context->package_uid)) {
+      context->package_frame_cookies.insert(cookie.cookie);
     }
   }
 
diff --git a/src/trace_redaction/filter_print_events.cc b/src/trace_redaction/filter_print_events.cc
index 6c1843b50..869c7a0fe 100644
--- a/src/trace_redaction/filter_print_events.cc
+++ b/src/trace_redaction/filter_print_events.cc
@@ -41,9 +41,6 @@ bool FilterPrintEvents::KeepEvent(const Context& context,
   PERFETTO_DCHECK(context.timeline);
   PERFETTO_DCHECK(context.package_uid.has_value());
 
-  const auto* timeline = context.timeline.get();
-  auto package_uid = context.package_uid;
-
   protozero::ProtoDecoder event(bytes);
 
   // This is not a print packet. Keep the packet.
@@ -56,9 +53,9 @@ bool FilterPrintEvents::KeepEvent(const Context& context,
       event.FindField(protos::pbzero::FtraceEvent::kTimestampFieldNumber);
   auto pid = event.FindField(protos::pbzero::FtraceEvent::kPidFieldNumber);
 
-  // Pid + Time --> UID, if the uid matches the target package, keep the event.
   return pid.valid() && time.valid() &&
-         timeline->Search(time.as_uint64(), pid.as_int32()).uid == package_uid;
+         context.timeline->PidConnectsToUid(time.as_uint64(), pid.as_int32(),
+                                            *context.package_uid);
 }
 
 }  // namespace perfetto::trace_redaction
diff --git a/src/trace_redaction/filter_sched_waking_events.cc b/src/trace_redaction/filter_sched_waking_events.cc
index e08a9d6ef..287090dad 100644
--- a/src/trace_redaction/filter_sched_waking_events.cc
+++ b/src/trace_redaction/filter_sched_waking_events.cc
@@ -59,14 +59,9 @@ bool FilterSchedWakingEvents::KeepEvent(const Context& context,
   auto outer_pid =
       event_decoder.FindField(protos::pbzero::FtraceEvent::kPidFieldNumber);
 
-  if (!outer_pid.valid()) {
-    return false;  // Remove
-  }
-
-  auto outer_slice = context.timeline->Search(
-      timestamp.as_uint64(), static_cast<int32_t>(outer_pid.as_uint32()));
-
-  if (outer_slice.uid != context.package_uid.value()) {
+  if (!outer_pid.valid() ||
+      !context.timeline->PidConnectsToUid(
+          timestamp.as_uint64(), outer_pid.as_int32(), *context.package_uid)) {
     return false;  // Remove
   }
 
@@ -75,13 +70,13 @@ bool FilterSchedWakingEvents::KeepEvent(const Context& context,
   auto inner_pid = waking_decoder.FindField(
       protos::pbzero::SchedWakingFtraceEvent::kPidFieldNumber);
 
-  if (!inner_pid.valid()) {
+  if (!inner_pid.valid() ||
+      !context.timeline->PidConnectsToUid(
+          timestamp.as_uint64(), inner_pid.as_int32(), *context.package_uid)) {
     return false;  // Remove
   }
 
-  auto inner_slice =
-      context.timeline->Search(timestamp.as_uint64(), inner_pid.as_int32());
-  return inner_slice.uid == context.package_uid.value();
+  return true;
 }
 
 }  // namespace perfetto::trace_redaction
diff --git a/src/trace_redaction/filter_task_rename.cc b/src/trace_redaction/filter_task_rename.cc
index 65f84fba0..467e5b5bb 100644
--- a/src/trace_redaction/filter_task_rename.cc
+++ b/src/trace_redaction/filter_task_rename.cc
@@ -67,8 +67,8 @@ bool FilterTaskRename::KeepEvent(const Context& context,
     return false;
   }
 
-  auto slice = context.timeline->Search(timestamp.as_uint64(), pid.as_int32());
-  return slice.uid == context.package_uid.value();
+  return context.timeline->PidConnectsToUid(
+      timestamp.as_uint64(), pid.as_int32(), *context.package_uid);
 }
 
 }  // namespace perfetto::trace_redaction
diff --git a/src/trace_redaction/process_thread_timeline.cc b/src/trace_redaction/process_thread_timeline.cc
index c1c3984f5..43d36372e 100644
--- a/src/trace_redaction/process_thread_timeline.cc
+++ b/src/trace_redaction/process_thread_timeline.cc
@@ -45,8 +45,9 @@ void ProcessThreadTimeline::Sort() {
   mode_ = Mode::kRead;
 }
 
-ProcessThreadTimeline::Slice ProcessThreadTimeline::Search(uint64_t ts,
-                                                           int32_t pid) const {
+bool ProcessThreadTimeline::PidConnectsToUid(uint64_t ts,
+                                             int32_t pid,
+                                             uint64_t uid) const {
   PERFETTO_DCHECK(mode_ == Mode::kRead);
 
   auto event = FindPreviousEvent(ts, pid);
@@ -54,26 +55,23 @@ ProcessThreadTimeline::Slice ProcessThreadTimeline::Search(uint64_t ts,
   for (size_t d = 0; d < kMaxSearchDepth; ++d) {
     // The thread/process was freed. It won't exist until a new open event.
     if (event.type != Event::Type::kOpen) {
-      return {pid, Event::kUnknownUid};
+      return false;
     }
 
-    // System processes all have uids equal to zero, so everything eventually
-    // has a uid. This means that all threads should find a process and a uid.
-    // However, if a thread does not have a process (this should not happen)
-    // that thread will be treated as invalid.
-    if (event.uid != Event::kUnknownUid) {
-      return {pid, event.uid};
+    // TODO(vaage): Normalize the uid values.
+    if (event.uid == uid) {
+      return true;
     }
 
-    // If there is no parent, there is no reason to keep searching.
+    // If there is no parent, there is no way to keep searching.
     if (event.ppid == Event::kUnknownPid) {
-      return {pid, Event::kUnknownUid};
+      return false;
     }
 
     event = FindPreviousEvent(ts, event.ppid);
   }
 
-  return {pid, Event::kUnknownUid};
+  return false;
 }
 
 ProcessThreadTimeline::Event ProcessThreadTimeline::FindPreviousEvent(
diff --git a/src/trace_redaction/process_thread_timeline.h b/src/trace_redaction/process_thread_timeline.h
index b176121eb..fdb70ba2d 100644
--- a/src/trace_redaction/process_thread_timeline.h
+++ b/src/trace_redaction/process_thread_timeline.h
@@ -82,15 +82,6 @@ class ProcessThreadTimeline {
     uint64_t uid = kUnknownUid;
   };
 
-  // The state of a process at a specific point in time.
-  struct Slice {
-    int32_t pid = -1;
-
-    // It is safe to use 0 as the invalid value because that's effectively
-    // what happening in the trace.
-    uint64_t uid = 0;
-  };
-
   ProcessThreadTimeline() = default;
 
   ProcessThreadTimeline(const ProcessThreadTimeline&) = delete;
@@ -105,12 +96,8 @@ class ProcessThreadTimeline {
   // subset of events will, on average, be trivally small.
   void Sort();
 
-  // Returns a snapshot that contains a process's pid and ppid, but contains the
-  // first uid found in its parent-child chain. If a uid cannot be found, uid=0
-  // is returned.
-  //
-  // `Sort()` must be called before this.
-  Slice Search(uint64_t ts, int32_t pid) const;
+  // Returns true if a process/thread is connected to a package.
+  bool PidConnectsToUid(uint64_t ts, int32_t pid, uint64_t uid) const;
 
   // Finds the pid's last event before ts.
   Event FindPreviousEvent(uint64_t ts, int32_t pid) const;
diff --git a/src/trace_redaction/process_thread_timeline_unittest.cc b/src/trace_redaction/process_thread_timeline_unittest.cc
index 23a97274a..4e86a8ab8 100644
--- a/src/trace_redaction/process_thread_timeline_unittest.cc
+++ b/src/trace_redaction/process_thread_timeline_unittest.cc
@@ -193,35 +193,29 @@ class ProcessThreadTimelineIsConnectedTest : public testing::Test {
 
 // PID A is directly connected to UID A.
 TEST_F(ProcessThreadTimelineIsConnectedTest, DirectPidAndUid) {
-  auto slice = timeline_.Search(kTimeB, kPidA);
-
-  ASSERT_EQ(slice.pid, kPidA);
-  ASSERT_EQ(slice.uid, kUidA);
+  ASSERT_TRUE(timeline_.PidConnectsToUid(kTimeB, kPidA, kUidA));
 }
 
 // PID B is indirectly connected to UID A through PID A.
 TEST_F(ProcessThreadTimelineIsConnectedTest, IndirectPidAndUid) {
-  auto slice = timeline_.Search(kTimeB, kPidB);
+  ASSERT_TRUE(timeline_.PidConnectsToUid(kTimeB, kPidB, kUidA));
+}
 
-  ASSERT_EQ(slice.pid, kPidB);
-  ASSERT_EQ(slice.uid, kUidA);
+// UID A and UID C are valid packages. However, PID B is connected to UID A, not
+// UID C.
+TEST_F(ProcessThreadTimelineIsConnectedTest, NotConnectedToOtherUid) {
+  ASSERT_FALSE(timeline_.PidConnectsToUid(kTimeB, kPidB, kUidC));
 }
 
 // PID D is not in the timeline, so it shouldn't be connected to anything.
 TEST_F(ProcessThreadTimelineIsConnectedTest, MissingPid) {
-  auto slice = timeline_.Search(kTimeB, kPidD);
-
-  ASSERT_EQ(slice.pid, kPidD);
-  ASSERT_EQ(slice.uid, ProcessThreadTimeline::Event::kUnknownUid);
+  ASSERT_FALSE(timeline_.PidConnectsToUid(kTimeB, kPidD, kUidA));
 }
 
 // Even through there is a connection between PID A and UID A, the query is too
 // soon (events are at TIME B, but the query is at TIME A).
 TEST_F(ProcessThreadTimelineIsConnectedTest, PrematureDirectPidAndUid) {
-  auto slice = timeline_.Search(kTimeA, kPidA);
-
-  ASSERT_EQ(slice.pid, kPidA);
-  ASSERT_EQ(slice.uid, ProcessThreadTimeline::Event::kUnknownUid);
+  ASSERT_FALSE(timeline_.PidConnectsToUid(kTimeA, kPidA, kUidA));
 }
 
 }  // namespace perfetto::trace_redaction
diff --git a/src/trace_redaction/redact_sched_switch.cc b/src/trace_redaction/redact_sched_switch.cc
index 8dca0301f..83cb36f79 100644
--- a/src/trace_redaction/redact_sched_switch.cc
+++ b/src/trace_redaction/redact_sched_switch.cc
@@ -26,11 +26,13 @@ namespace perfetto::trace_redaction {
 
 namespace {
 
-protozero::ConstChars SanitizeCommValue(const Context& context,
-                                        ProcessThreadTimeline::Slice slice,
-                                        protozero::Field field) {
-  if (NormalizeUid(slice.uid) == NormalizeUid(context.package_uid.value())) {
-    return field.as_string();
+// TODO(vaage): Merge with RedactComm in redact_task_newtask.cc.
+protozero::ConstChars RedactComm(const Context& context,
+                                 uint64_t ts,
+                                 int32_t pid,
+                                 protozero::ConstChars comm) {
+  if (context.timeline->PidConnectsToUid(ts, pid, *context.package_uid)) {
+    return comm;
   }
 
   return {};
@@ -105,22 +107,19 @@ base::Status RedactSchedSwitch::Redact(
   // Avoid making the message until we know that we have prev and next pids.
   auto sched_switch_message = event_message->set_sched_switch();
 
-  auto prev_slice =
-      context.timeline->Search(timestamp.as_uint64(), prev_pid.as_int32());
-  auto next_slice =
-      context.timeline->Search(timestamp.as_uint64(), next_pid.as_int32());
-
   for (auto field = sched_switch_decoder.ReadField(); field.valid();
        field = sched_switch_decoder.ReadField()) {
     switch (field.id()) {
       case protos::pbzero::SchedSwitchFtraceEvent::kNextCommFieldNumber:
         sched_switch_message->set_next_comm(
-            SanitizeCommValue(context, next_slice, field));
+            RedactComm(context, timestamp.as_uint64(), next_pid.as_int32(),
+                       field.as_string()));
         break;
 
       case protos::pbzero::SchedSwitchFtraceEvent::kPrevCommFieldNumber:
         sched_switch_message->set_prev_comm(
-            SanitizeCommValue(context, prev_slice, field));
+            RedactComm(context, timestamp.as_uint64(), prev_pid.as_int32(),
+                       field.as_string()));
         break;
 
       default:
diff --git a/src/trace_redaction/redact_task_newtask.cc b/src/trace_redaction/redact_task_newtask.cc
index c9426121e..fb0cfe8ba 100644
--- a/src/trace_redaction/redact_task_newtask.cc
+++ b/src/trace_redaction/redact_task_newtask.cc
@@ -26,18 +26,19 @@ namespace perfetto::trace_redaction {
 
 namespace {
 
-protozero::ConstChars SanitizeCommValue(const Context& context,
-                                        ProcessThreadTimeline::Slice slice,
-                                        protozero::Field field) {
-  if (NormalizeUid(slice.uid) == NormalizeUid(context.package_uid.value())) {
-    return field.as_string();
+// TODO(vaage): Merge with RedactComm in redact_sched_switch.cc.
+protozero::ConstChars RedactComm(const Context& context,
+                                 uint64_t ts,
+                                 int32_t pid,
+                                 protozero::ConstChars comm) {
+  if (context.timeline->PidConnectsToUid(ts, pid, *context.package_uid)) {
+    return comm;
   }
 
   return {};
 }
 
 }  // namespace
-
 // Redact sched switch trace events in an ftrace event bundle:
 //
 // event {
@@ -93,8 +94,6 @@ base::Status RedactTaskNewTask::Redact(
   // Avoid making the message until we know that we have prev and next pids.
   auto* new_task_message = event_message->set_task_newtask();
 
-  auto slice = context.timeline->Search(timestamp.as_uint64(), pid.as_int32());
-
   for (auto field = new_task_decoder.ReadField(); field.valid();
        field = new_task_decoder.ReadField()) {
     // Perfetto view (ui.perfetto.dev) crashes if the comm value is missing.
@@ -102,7 +101,8 @@ base::Status RedactTaskNewTask::Redact(
     // This appears to work.
     if (field.id() ==
         protos::pbzero::TaskNewtaskFtraceEvent::kCommFieldNumber) {
-      new_task_message->set_comm(SanitizeCommValue(context, slice, field));
+      new_task_message->set_comm(RedactComm(context, timestamp.as_uint64(),
+                                            pid.as_int32(), field.as_string()));
     } else {
       proto_util::AppendField(field, new_task_message);
     }
diff --git a/src/trace_redaction/remap_scheduling_events.cc b/src/trace_redaction/remap_scheduling_events.cc
index 4817f1972..be0973e12 100644
--- a/src/trace_redaction/remap_scheduling_events.cc
+++ b/src/trace_redaction/remap_scheduling_events.cc
@@ -32,14 +32,18 @@ int32_t RemapPid(const Context& context,
   PERFETTO_DCHECK(context.package_uid.value());
   PERFETTO_DCHECK(cpu < context.synthetic_threads->tids.size());
 
-  auto slice = context.timeline->Search(timestamp, pid);
+  // PID 0 is used for CPU idle. If it was to get re-mapped, threading
+  // information get corrupted.
+  if (pid == 0) {
+    return 0;
+  }
 
-  auto expected_uid = NormalizeUid(slice.uid);
-  auto actual_uid = NormalizeUid(context.package_uid.value());
+  if (context.timeline->PidConnectsToUid(timestamp, pid,
+                                         *context.package_uid)) {
+    return pid;
+  }
 
-  return !pid || expected_uid == actual_uid
-             ? pid
-             : context.synthetic_threads->tids[cpu];
+  return context.synthetic_threads->tids[cpu];
 }
 }  // namespace
 
diff --git a/src/trace_redaction/scrub_process_stats.cc b/src/trace_redaction/scrub_process_stats.cc
index 19d404898..78c983ae1 100644
--- a/src/trace_redaction/scrub_process_stats.cc
+++ b/src/trace_redaction/scrub_process_stats.cc
@@ -57,9 +57,6 @@ base::Status ScrubProcessStats::Transform(const Context& context,
   PERFETTO_DCHECK(time_field.valid());
   auto time = time_field.as_uint64();
 
-  auto* timeline = context.timeline.get();
-  auto uid = context.package_uid.value();
-
   for (auto packet_field = packet_decoder.ReadField(); packet_field.valid();
        packet_field = packet_decoder.ReadField()) {
     if (packet_field.id() !=
@@ -83,8 +80,9 @@ base::Status ScrubProcessStats::Transform(const Context& context,
         protozero::ProtoDecoder process_decoder(process_stats_field.as_bytes());
         auto pid = process_decoder.FindField(
             protos::pbzero::ProcessStats::Process::kPidFieldNumber);
-        keep_field =
-            pid.valid() && timeline->Search(time, pid.as_int32()).uid == uid;
+
+        keep_field = context.timeline->PidConnectsToUid(time, pid.as_int32(),
+                                                        *context.package_uid);
       } else {
         keep_field = true;
       }
diff --git a/src/trace_redaction/scrub_process_trees.cc b/src/trace_redaction/scrub_process_trees.cc
index 989d0942f..a3f1f3b46 100644
--- a/src/trace_redaction/scrub_process_trees.cc
+++ b/src/trace_redaction/scrub_process_trees.cc
@@ -43,14 +43,10 @@ void TryAppendPid(const Context& context,
     return;
   }
 
-  auto slice = context.timeline->Search(timestamp.as_uint64(), pid.as_int32());
-
-  // Only keep the target process cmdline.
-  if (NormalizeUid(slice.uid) != NormalizeUid(context.package_uid.value())) {
-    return;
+  if (context.timeline->PidConnectsToUid(timestamp.as_uint64(), pid.as_int32(),
+                                         *context.package_uid)) {
+    proto_util::AppendField(value, message);
   }
-
-  proto_util::AppendField(value, message);
 }
 
 }  // namespace