Add ancestry annotator

Source/common/BUILD

@@ -1,3 +1,4 @@
+load("@com_google_protobuf//bazel:proto_library.bzl", "proto_library")
 load("@rules_cc//cc:defs.bzl", "cc_proto_library")
 load("//:helper.bzl", "santa_unit_test")
 

Source/common/santa.proto

@@ -22,12 +22,6 @@ message GroupInfo {
   optional string name = 2;
 }
 
-// A process is uniquely identified on macOS by its pid and pidversion
-message ProcessID {
-  optional int32 pid = 1;
-  optional int32 pidversion = 2;
-}
-
 // Code signature information
 message CodeSignature {
   // The code directory hash identifies a specific version of a program
@@ -129,13 +123,13 @@ message FileDescriptor {
 // Process information
 message ProcessInfo {
   // Process ID of the process
-  optional ProcessID id = 1;
+  optional santa.pb.v1.process_tree.ProcessID id = 1;
 
   // Process ID of the parent process
-  optional ProcessID parent_id = 2;
+  optional santa.pb.v1.process_tree.ProcessID parent_id = 2;
 
   // Process ID of the process responsible for this one
-  optional ProcessID responsible_id = 3;
+  optional santa.pb.v1.process_tree.ProcessID responsible_id = 3;
 
   // Original parent ID, remains stable in the event a process is reparented
   optional int32 original_parent_pid = 4;
@@ -181,10 +175,10 @@ message ProcessInfo {
 // Light variant of ProcessInfo message to help minimize on-disk/on-wire sizes
 message ProcessInfoLight {
   // Process ID of the process
-  optional ProcessID id = 1;
+  optional santa.pb.v1.process_tree.ProcessID id = 1;
 
   // Process ID of the parent process
-  optional ProcessID parent_id = 2;
+  optional santa.pb.v1.process_tree.ProcessID parent_id = 2;
 
   // Original parent ID, remains stable in the event a process is reparented
   optional int32 original_parent_pid = 3;

Source/santad/BUILD

@@ -802,6 +802,7 @@ objc_library(
         "//Source/common:Unit",
         "//Source/santad/ProcessTree:process_tree",
         "//Source/santad/ProcessTree/annotations:originator",
+        "//Source/santad/ProcessTree/annotations:ancestry",
         "@MOLXPCConnection",
     ],
 )
@@ -840,7 +841,6 @@ macos_bundle(
     ],
     entitlements = select({
         "//:adhoc_build": "com.google.santa.daemon.systemextension-adhoc.entitlements",
-        # Non-adhoc builds get their entitlements from the provisioning profile.
         "//conditions:default": None,
     }),
     infoplists = ["Info.plist"],

Source/santad/EventProviders/EndpointSecurity/Message.h

@@ -41,7 +41,7 @@ class Message {
   Message(const Message& other);
   Message& operator=(const Message& other) = delete;
 
-  void SetProcessToken(santa::santad::process_tree::ProcessToken tok);
+  void SetProcessToken(santa::santad::process_tree::ProcessToken& tok);
 
   // Operators to access underlying es_message_t
   inline const es_message_t* operator->() const { return es_msg_; }

Source/santad/Logs/EndpointSecurity/Serializers/Protobuf.mm

@@ -103,7 +103,8 @@ static inline void EncodeTimestamp(Timestamp *timestamp, struct timeval tv) {
   EncodeTimestamp(timestamp, (struct timespec){tv.tv_sec, tv.tv_usec * 1000});
 }
 
-static inline void EncodeProcessID(pbv1::ProcessID *proc_id, const audit_token_t &tok) {
+static inline void EncodeProcessID(pbv1::process_tree::ProcessID *proc_id,
+                                   const audit_token_t &tok) {
   proc_id->set_pid(Pid(tok));
   proc_id->set_pidversion(Pidversion(tok));
 }

Source/santad/Logs/EndpointSecurity/Writers/FSSpool/BUILD

@@ -1,3 +1,4 @@
+load("@com_google_protobuf//bazel:proto_library.bzl", "proto_library")
 load("@rules_cc//cc:defs.bzl", "cc_proto_library")
 load("//:helper.bzl", "santa_unit_test")
 

Source/santad/ProcessTree/BUILD

@@ -1,3 +1,4 @@
+load("@com_google_protobuf//bazel:proto_library.bzl", "proto_library")
 load("@rules_cc//cc:defs.bzl", "cc_proto_library")
 load("//:helper.bzl", "santa_unit_test")
 

Source/santad/ProcessTree/annotations/BUILD

@@ -25,6 +25,19 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "ancestry",
+    srcs = ["ancestry.cc"],
+    hdrs = ["ancestry.h"],
+    deps = [
+        ":annotator",
+        "//Source/santad/ProcessTree:process",
+        "//Source/santad/ProcessTree:process_tree",
+        "//Source/santad/ProcessTree:process_tree_cc_proto",
+        "@com_google_absl//absl/container:flat_hash_map",
+    ],
+)
+
 santa_unit_test(
     name = "originator_test",
     srcs = ["originator_test.mm"],
@@ -35,3 +48,14 @@ santa_unit_test(
         "//Source/santad/ProcessTree:process_tree_test_helpers",
     ],
 )
+
+santa_unit_test(
+    name = "ancestry_test",
+    srcs = ["ancestry_test.mm"],
+    deps = [
+        ":ancestry",
+        "//Source/santad/ProcessTree:process",
+        "//Source/santad/ProcessTree:process_tree_cc_proto",
+        "//Source/santad/ProcessTree:process_tree_test_helpers",
+    ],
+)

Source/santad/ProcessTree/annotations/ancestry.cc

@@ -0,0 +1,74 @@
+/// Copyright 2023 Google LLC
+///
+/// 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
+///
+///     https://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 "Source/santad/ProcessTree/annotations/ancestry.h"
+
+#include "Source/santad/ProcessTree/process.h"
+#include "Source/santad/ProcessTree/process_tree.h"
+#include "Source/santad/ProcessTree/process_tree.pb.h"
+#include "absl/container/flat_hash_map.h"
+
+namespace ptpb = ::santa::pb::v1::process_tree;
+
+namespace santa::santad::process_tree {
+
+::santa::pb::v1::process_tree::Annotations::Ancestry
+AncestryAnnotator::getAncestry() const {
+  ::santa::pb::v1::process_tree::Annotations::Ancestry ancestry;
+  ancestry.CopyFrom(ancestry_);
+  return ancestry;
+}
+
+void AncestryAnnotator::AddProcessToAncestry(
+    ptpb::Annotations::Ancestry &ancestry, const Process &process) {
+  ptpb::ProcessID *ancestor = ancestry.add_ancestor();
+  ancestor->set_pid(process.pid_.pid);
+  ancestor->set_pidversion(process.pid_.pidversion);
+}
+
+void AncestryAnnotator::AnnotateFork(ProcessTree &tree, const Process &parent,
+                                     const Process &child) {
+  ptpb::Annotations::Ancestry ancestry;
+
+  // If parent process has ancestry annotation, copy and add parent.
+  if (auto parent_annotation = tree.GetAnnotation<AncestryAnnotator>(parent)) {
+    ancestry.CopyFrom((*parent_annotation)->getAncestry());
+    AddProcessToAncestry(ancestry, parent);
+    // Otherwise, get all ancestors of the child and add them.
+  } else {
+    std::vector<ptpb::ProcessID> ancestors = tree.GetAncestors(child);
+    // Add ancestors starting from the root process
+    for (auto it = ancestors.rbegin(); it != ancestors.rend(); it++) {
+      ptpb::ProcessID *ancestor = ancestry.add_ancestor();
+      ancestor->set_pid(it->pid());
+      ancestor->set_pidversion(it->pidversion());
+    }
+  }
+  tree.AnnotateProcess(child, std::make_shared<AncestryAnnotator>(ancestry));
+}
+
+void AncestryAnnotator::AnnotateExec(ProcessTree &tree,
+                                     const Process &orig_process,
+                                     const Process &new_process) {
+  // Do not annotate process on exec
+  return;
+}
+
+std::optional<ptpb::Annotations> AncestryAnnotator::Proto() const {
+  auto annotation = ptpb::Annotations();
+  auto *ancestry_ptr = annotation.mutable_ancestry();
+  ancestry_ptr->CopyFrom(AncestryAnnotator::getAncestry());
+  return annotation;
+}
+
+}  // namespace santa::santad::process_tree

Source/santad/ProcessTree/annotations/ancestry.h

@@ -0,0 +1,50 @@
+/// Copyright 2023 Google LLC
+///
+/// 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
+///
+///     https://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.
+#ifndef SANTA__SANTAD_PROCESSTREE_ANNOTATIONS_ANCESTRY_H
+#define SANTA__SANTAD_PROCESSTREE_ANNOTATIONS_ANCESTRY_H
+
+#include <optional>
+
+#include "Source/santad/ProcessTree/annotations/annotator.h"
+#include "Source/santad/ProcessTree/process.h"
+#include "Source/santad/ProcessTree/process_tree.pb.h"
+
+namespace santa::santad::process_tree {
+
+class AncestryAnnotator : public Annotator {
+ public:
+  // clang-format off
+  AncestryAnnotator() {}
+ explicit AncestryAnnotator(
+            ::santa::pb::v1::process_tree::Annotations::Ancestry ancestry)
+            : ancestry_(ancestry) {};
+  // clang-format on
+  void AnnotateFork(ProcessTree &tree, const Process &parent,
+                    const Process &child) override;
+  void AnnotateExec(ProcessTree &tree, const Process &orig_process,
+                    const Process &new_process) override;
+  std::optional<::santa::pb::v1::process_tree::Annotations> Proto()
+      const override;
+  ::santa::pb::v1::process_tree::Annotations::Ancestry getAncestry() const;
+
+ private:
+  void AddProcessToAncestry(
+      ::santa::pb::v1::process_tree::Annotations::Ancestry &ancestry,
+      const Process &process);
+  ::santa::pb::v1::process_tree::Annotations::Ancestry ancestry_;
+};
+
+}  // namespace santa::santad::process_tree
+
+#endif

Source/santad/ProcessTree/annotations/ancestry_test.mm

@@ -0,0 +1,98 @@
+/// Copyright 2023 Google LLC
+///
+/// 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
+///
+///     https://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.
+
+#import <Foundation/Foundation.h>
+#import <XCTest/XCTest.h>
+
+#include "Source/santad/ProcessTree/annotations/ancestry.h"
+#include "Source/santad/ProcessTree/process.h"
+#include "Source/santad/ProcessTree/process_tree.pb.h"
+#include "Source/santad/ProcessTree/process_tree_test_helpers.h"
+
+using namespace santa::santad::process_tree;
+namespace ptpb = ::santa::pb::v1::process_tree;
+
+@interface AncestryAnnotatorTest : XCTestCase
+@property std::shared_ptr<ProcessTreeTestPeer> tree;
+@property std::shared_ptr<const Process> initProc;
+@end
+
+@implementation AncestryAnnotatorTest
+
+- (void)setUp {
+  std::vector<std::unique_ptr<Annotator>> annotators;
+  annotators.emplace_back(std::make_unique<AncestryAnnotator>());
+  self.tree = std::make_shared<ProcessTreeTestPeer>(std::move(annotators));
+  self.initProc = self.tree->InsertInit();
+}
+
+- (void)testSingleFork_childHasAncestryAnnotation {
+  uint64_t event_id = 1;
+  // PID 1.1: fork() -> PID 2.2
+  const struct Pid child_pid = {.pid = 2, .pidversion = 2};
+  self.tree->HandleFork(event_id++, *self.initProc, child_pid);
+
+  auto child = *self.tree->Get(child_pid);
+  auto annotation_opt = self.tree->GetAnnotation<AncestryAnnotator>(*child);
+  XCTAssertTrue(annotation_opt.has_value());
+  auto proto_opt = (*annotation_opt)->Proto();
+
+  XCTAssertTrue(proto_opt.has_value());
+  XCTAssertEqual(proto_opt->ancestry().ancestor_size(), 1);
+  XCTAssertEqual(proto_opt->ancestry().ancestor().Get(0).pid(), 1);
+  XCTAssertEqual(proto_opt->ancestry().ancestor().Get(0).pidversion(), 1);
+}
+
+- (void)testDoubleFork_grandchildHasAncestryAnnotation {
+  uint64_t event_id = 1;
+  // PID 1.1: fork() -> PID 2.2 fork() -> PID 3.3
+  const struct Pid child_pid = {.pid = 2, .pidversion = 2};
+  const struct Pid grandchild_pid = {.pid = 3, .pidversion = 3};
+
+  self.tree->HandleFork(event_id++, *self.initProc, child_pid);
+  auto child = *self.tree->Get(child_pid);
+  self.tree->HandleFork(event_id++, *child, grandchild_pid);
+
+  auto grandchild = *self.tree->Get(grandchild_pid);
+  auto annotation_opt = self.tree->GetAnnotation<AncestryAnnotator>(*grandchild);
+  XCTAssertTrue(annotation_opt.has_value());
+  auto grandchild_proto_opt = (*annotation_opt)->Proto();
+  XCTAssertTrue(grandchild_proto_opt.has_value());
+  auto grandchild_proto = *grandchild_proto_opt;
+  XCTAssertEqual(grandchild_proto.ancestry().ancestor_size(), 2);
+  XCTAssertEqual(grandchild_proto.ancestry().ancestor().Get(0).pid(), 1);
+  XCTAssertEqual(grandchild_proto.ancestry().ancestor().Get(0).pidversion(), 1);
+  XCTAssertEqual(grandchild_proto.ancestry().ancestor().Get(1).pid(), 2);
+  XCTAssertEqual(grandchild_proto.ancestry().ancestor().Get(1).pidversion(), 2);
+}
+
+- (void)testExec_noEffectOnAncestryAnnotation() {
+  uint64_t event_id = 1;
+
+  // PID 1.1: fork() -> PID 2.2
+  const struct Pid child_pid = {.pid = 2, .pidversion = 2};
+  self.tree->HandleFork(event_id++, *self.initProc, child_pid);
+  auto child = *self.tree->Get(child_pid);
+  auto annotation_opt = self.tree->GetAnnotation<AncestryAnnotator>(*child);
+  XCTAssertTrue(annotation_opt.has_value());
+  auto proto_opt = (*annotation_opt)->Proto();
+
+  // PID 2.2: exec("/usr/bin/login") -> PID 2.3
+  const struct Pid login_exec_pid = {.pid = 2, .pidversion = 3};
+  const struct Program login_prog = {.executable = "/usr/bin/login", .arguments = {}};
+  auto login = *self.tree->Get(login_pid);
+  self.tree->HandleExec(event_id++, *login, login_exec_pid, login_prog, cred);
+}
+
+@end

Source/santad/ProcessTree/process_tree.cc

@@ -67,6 +67,31 @@ void ProcessTree::BackfillInsertChildren(
   }
 }
 
+std::vector<::santa::pb::v1::process_tree::ProcessID> ProcessTree::GetAncestors(
+    const Process &process) {
+  std::vector<::santa::pb::v1::process_tree::ProcessID> ancestors;
+  if (!process.parent_) {
+    return ancestors;
+  }
+  {
+    absl::MutexLock lock(&mtx_);
+    std::optional<std::shared_ptr<Process>> chainProcess =
+        GetLocked(process.parent_->pid_);
+    while (chainProcess != std::nullopt) {
+      ::santa::pb::v1::process_tree::ProcessID process_id;
+      process_id.set_pid((*chainProcess)->pid_.pid);
+      process_id.set_pidversion((*chainProcess)->pid_.pidversion);
+      ancestors.push_back(process_id);
+      if ((*chainProcess)->parent_) {
+        chainProcess = GetLocked((*chainProcess)->parent_->pid_);
+      } else {
+        chainProcess = std::nullopt;
+      }
+    }
+  }
+  return ancestors;
+}
+
 void ProcessTree::HandleFork(uint64_t timestamp, const Process &parent,
                              const Pid new_pid) {
   if (Step(timestamp)) {

Source/santad/ProcessTree/process_tree.h

@@ -43,6 +43,9 @@ class ProcessTree {
   // Initialize the tree with the processes currently running on the system.
   absl::Status Backfill();
 
+  std::vector<::santa::pb::v1::process_tree::ProcessID> GetAncestors(
+      const Process &process);
+
   // Inform the tree of a fork event, in which the parent process spawns a child
   // with the only difference between the two being the pid.
   void HandleFork(uint64_t timestamp, const Process &parent,