/*
 * Copyright (C) 2024 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 "src/trace_redaction/redact_process_trees.h"

#include <string>

#include "perfetto/base/status.h"
#include "perfetto/protozero/field.h"
#include "perfetto/protozero/scattered_heap_buffer.h"
#include "src/trace_processor/util/status_macros.h"
#include "src/trace_redaction/proto_util.h"
#include "src/trace_redaction/trace_redaction_framework.h"

#include "protos/perfetto/trace/ps/process_tree.pbzero.h"
#include "protos/perfetto/trace/trace_packet.pbzero.h"

namespace perfetto::trace_redaction {

ProcessTreeFilter::~ProcessTreeFilter() = default;

ProcessTreeModifier::~ProcessTreeModifier() = default;

bool ProcessTreeFilterAllowAll::Filter(const Context&,
                                       uint64_t,
                                       int32_t) const {
  return true;
}

bool ProcessTreeFilterConnectedToPackage::Filter(const Context& context,
                                                 uint64_t ts,
                                                 int32_t pid) const {
  PERFETTO_DCHECK(context.timeline);
  PERFETTO_DCHECK(context.package_uid.has_value());

  // Checking the uid directly is an option, but the timeline handles multiple
  // instances of the same app.
  return context.timeline->PidConnectsToUid(ts, pid, *context.package_uid);
}

base::Status ProcessTreeDoNothing::Modify(const Context&,
                                          protos::pbzero::ProcessTree*) const {
  return base::OkStatus();
}

base::Status ProcessTreeCreateSynthThreads::Modify(
    const Context& context,
    protos::pbzero::ProcessTree* message) const {
  PERFETTO_DCHECK(message);

  if (!context.synthetic_threads.has_value()) {
    return base::ErrStatus(
        "ProcessTreeCreateSynthThreads: missing synthetic thread group");
  }

  // There should be one thread per cpu. If there are no threads, it means there
  // were no cpus. That wrong.
  if (context.synthetic_threads->tids.empty()) {
    return base::ErrStatus(
        "ProcessTreeCreateSynthThreads: missing synthetic threads");
  }

  // uid 0 and ppid 2 means it is a system process.
  auto* process = message->add_processes();
  process->set_uid(0);
  process->set_ppid(2);
  process->set_pid(context.synthetic_threads->tgid);
  process->add_cmdline("MergedThreads");

  for (auto tid : context.synthetic_threads->tids) {
    auto* thread = message->add_threads();
    thread->set_tgid(context.synthetic_threads->tgid);
    thread->set_tid(tid);
    thread->set_name("CPU");
  }

  return base::OkStatus();
}

base::Status RedactProcessTrees::Transform(const Context& context,
                                           std::string* packet) const {
  PERFETTO_DCHECK(packet);

  if (!context.package_uid.has_value()) {
    return base::ErrStatus("RedactProcessTrees: missing package uid.");
  }

  if (!context.timeline) {
    return base::ErrStatus("RedactProcessTrees: missing timeline.");
  }

  if (!context.synthetic_threads.has_value()) {
    return base::ErrStatus("RedactProcessTrees: missing synthentic threads.");
  }

  protozero::ProtoDecoder decoder(*packet);

  auto tree =
      decoder.FindField(protos::pbzero::TracePacket::kProcessTreeFieldNumber);

  if (!tree.valid()) {
    return base::OkStatus();
  }

  // This has been verified by the verify primitive.
  auto timestamp =
      decoder.FindField(protos::pbzero::TracePacket::kTimestampFieldNumber);

  protozero::HeapBuffered<protos::pbzero::TracePacket> message;

  for (auto it = decoder.ReadField(); it.valid(); it = decoder.ReadField()) {
    if (it.id() == tree.id()) {
      RETURN_IF_ERROR(OnProcessTree(context, timestamp.as_uint64(),
                                    it.as_bytes(),
                                    message->set_process_tree()));
    } else {
      proto_util::AppendField(it, message.get());
    }
  }

  packet->assign(message.SerializeAsString());

  return base::OkStatus();
}

base::Status RedactProcessTrees::OnProcessTree(
    const Context& context,
    uint64_t ts,
    protozero::ConstBytes bytes,
    protos::pbzero::ProcessTree* message) const {
  protozero::ProtoDecoder decoder(bytes);

  for (auto it = decoder.ReadField(); it.valid(); it = decoder.ReadField()) {
    switch (it.id()) {
      case protos::pbzero::ProcessTree::kProcessesFieldNumber:
        RETURN_IF_ERROR(OnProcess(context, ts, it, message));
        break;
      case protos::pbzero::ProcessTree::kThreadsFieldNumber:
        RETURN_IF_ERROR(OnThread(context, ts, it, message));
        break;
      default:
        proto_util::AppendField(it, message);
        break;
    }
  }

  PERFETTO_DCHECK(modifier_);
  return modifier_->Modify(context, message);
}

base::Status RedactProcessTrees::OnProcess(
    const Context& context,
    uint64_t ts,
    protozero::Field field,
    protos::pbzero::ProcessTree* message) const {
  protozero::ProtoDecoder decoder(field.as_bytes());

  auto pid =
      decoder.FindField(protos::pbzero::ProcessTree::Process::kPidFieldNumber);
  if (!pid.valid()) {
    return base::ErrStatus("RedactProcessTrees: process with no pid");
  }

  PERFETTO_DCHECK(filter_);

  if (filter_->Filter(context, ts, pid.as_int32())) {
    proto_util::AppendField(field, message);
  }

  return base::OkStatus();
}

base::Status RedactProcessTrees::OnThread(
    const Context& context,
    uint64_t ts,
    protozero::Field field,
    protos::pbzero::ProcessTree* message) const {
  protozero::ProtoDecoder decoder(field.as_bytes());

  auto tid =
      decoder.FindField(protos::pbzero::ProcessTree::Thread::kTidFieldNumber);
  if (!tid.valid()) {
    return base::ErrStatus("RedactProcessTrees: thread with no tid");
  }

  PERFETTO_DCHECK(filter_);

  if (filter_->Filter(context, ts, tid.as_int32())) {
    proto_util::AppendField(field, message);
  }

  return base::OkStatus();
}

}  // namespace perfetto::trace_redaction