| // Copyright 2021 The Chromium OS Authors. All rights reserved. |
| // Use of this source code is governed by a BSD-style license that can be |
| // found in the LICENSE file. |
| |
| #include "ml/process.h" |
| |
| #include <utility> |
| |
| #include <signal.h> |
| #include <sysexits.h> |
| #include <sys/mount.h> |
| #include <sys/types.h> |
| #include <sys/wait.h> |
| #include <unistd.h> |
| #include <pwd.h> |
| |
| #include <base/logging.h> |
| #include <base/process/process_metrics.h> |
| #include <base/strings/string_number_conversions.h> |
| #include <base/strings/string_util.h> |
| #include <libminijail.h> |
| #include <mojo/core/embedder/embedder.h> |
| #include <mojo/public/cpp/platform/platform_channel.h> |
| #include <mojo/public/cpp/system/invitation.h> |
| #include <scoped_minijail.h> |
| |
| #include "ml/daemon.h" |
| #include "ml/machine_learning_service_impl.h" |
| #include "ml/request_metrics.h" |
| #include "ml/time_metrics.h" |
| |
| namespace ml { |
| |
| namespace { |
| constexpr char kMojoBootstrapFdSwitchName[] = "mojo-bootstrap-fd"; |
| |
| constexpr char kInternalMojoPrimordialPipeName[] = "cros_ml"; |
| |
| constexpr char kDefaultMlServiceBinaryPath[] = "/usr/bin/ml_service"; |
| |
| constexpr uid_t kMlServiceDBusUid = 20177; |
| |
| std::string GetSeccompPolicyPath(const std::string& model_name) { |
| return "/usr/share/policy/ml_service-" + model_name + "-seccomp.policy"; |
| } |
| |
| std::string GetArgumentForWorkerProcess(int fd) { |
| std::string fd_argv = kMojoBootstrapFdSwitchName; |
| return "--" + fd_argv + "=" + std::to_string(fd); |
| } |
| |
| } // namespace |
| |
| // static |
| Process* Process::GetInstance() { |
| // This is thread-safe. |
| static base::NoDestructor<Process> instance; |
| return instance.get(); |
| } |
| |
| int Process::Run(int argc, char* argv[]) { |
| // Parses the command line and determines the process type. |
| base::CommandLine command_line(argc, argv); |
| std::string mojo_fd_string = |
| command_line.GetSwitchValueASCII(kMojoBootstrapFdSwitchName); |
| |
| if (mojo_fd_string.empty()) { |
| process_type_ = Type::kControl; |
| } else { |
| process_type_ = Type::kWorker; |
| } |
| |
| if (!command_line.GetArgs().empty()) { |
| LOG(ERROR) << "Unexpected command line arguments: " |
| << base::JoinString(command_line.GetArgs(), "\t"); |
| return ExitCode::kUnexpectedCommandLine; |
| } |
| |
| if (process_type_ == Type::kControl) { |
| ControlProcessRun(); |
| } else { |
| // The process type is either "control" or "worker". |
| DCHECK(GetType() == Type::kWorker); |
| const auto is_valid_fd_str = |
| base::StringToInt(mojo_fd_string, &mojo_bootstrap_fd_); |
| DCHECK(is_valid_fd_str) << "Invalid mojo bootstrap fd"; |
| WorkerProcessRun(); |
| } |
| |
| return ExitCode::kSuccess; |
| } |
| |
| Process::Type Process::GetType() { |
| return process_type_; |
| } |
| |
| bool Process::SpawnWorkerProcessAndGetPid(const mojo::PlatformChannel& channel, |
| const std::string& model_name, |
| pid_t* worker_pid) { |
| DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_); |
| DCHECK(worker_pid != nullptr); |
| // Should only be called by the control process. |
| DCHECK(IsControlProcess()) << "Should only be called by the control process"; |
| |
| // Start the process. |
| ScopedMinijail jail(minijail_new()); |
| |
| minijail_namespace_ipc(jail.get()); |
| minijail_namespace_uts(jail.get()); |
| minijail_namespace_net(jail.get()); |
| minijail_namespace_cgroups(jail.get()); |
| |
| // The following sandboxing makes unit test crash so we do not use them in |
| // unit tests. |
| if (process_type_ != Type::kControlForTest) { |
| minijail_namespace_pids(jail.get()); |
| minijail_namespace_vfs(jail.get()); |
| std::string seccomp_policy_path = GetSeccompPolicyPath(model_name); |
| minijail_parse_seccomp_filters(jail.get(), seccomp_policy_path.c_str()); |
| minijail_use_seccomp_filter(jail.get()); |
| } |
| |
| std::string fd_argv = kMojoBootstrapFdSwitchName; |
| // Use GetFD instead of TakeFD to non-destructively obtain the fd. |
| fd_argv = GetArgumentForWorkerProcess( |
| channel.remote_endpoint().platform_handle().GetFD().get()); |
| char* const argv[3] = {&ml_service_path_[0], &fd_argv[0], nullptr}; |
| |
| if (minijail_run_pid(jail.get(), &ml_service_path_[0], argv, worker_pid) != |
| 0) { |
| RecordProcessErrorEvent(ProcessError::kSpawnWorkerProcessFailed); |
| LOG(DFATAL) << "Failed to spawn worker process for " << model_name; |
| return false; |
| } |
| |
| return true; |
| } |
| |
| mojo::Remote<chromeos::machine_learning::mojom::MachineLearningService>& |
| Process::SendMojoInvitationAndGetRemote(pid_t worker_pid, |
| mojo::PlatformChannel channel, |
| const std::string& model_name) { |
| // Send the Mojo invitation to the worker process. |
| mojo::OutgoingInvitation invitation; |
| mojo::ScopedMessagePipeHandle pipe = |
| invitation.AttachMessagePipe(kInternalMojoPrimordialPipeName); |
| |
| mojo::Remote<chromeos::machine_learning::mojom::MachineLearningService> |
| remote(mojo::PendingRemote< |
| chromeos::machine_learning::mojom::MachineLearningService>( |
| std::move(pipe), 0u /* version */)); |
| |
| mojo::OutgoingInvitation::Send(std::move(invitation), worker_pid, |
| channel.TakeLocalEndpoint()); |
| |
| remote.set_disconnect_handler( |
| base::BindOnce(&Process::InternalPrimordialMojoPipeDisconnectHandler, |
| base::Unretained(this), worker_pid)); |
| |
| DCHECK(worker_pid_info_map_.find(worker_pid) == worker_pid_info_map_.end()) |
| << "Worker pid already exists"; |
| |
| WorkerInfo worker_info; |
| worker_info.remote = std::move(remote); |
| worker_info.process_metrics = |
| base::ProcessMetrics::CreateProcessMetrics(worker_pid); |
| // Baseline the CPU usage counter in `process_metrics` to be zero as of now. |
| worker_info.process_metrics->GetPlatformIndependentCPUUsage(); |
| |
| worker_pid_info_map_[worker_pid] = std::move(worker_info); |
| |
| return worker_pid_info_map_[worker_pid].remote; |
| } |
| |
| void Process::UnregisterWorkerProcess(pid_t pid) { |
| DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_); |
| const auto iter = worker_pid_info_map_.find(pid); |
| DCHECK(iter != worker_pid_info_map_.end()) << "Pid is not registered"; |
| worker_pid_info_map_.erase(iter); |
| } |
| |
| Process::Process() |
| : process_type_(Type::kUnset), |
| mojo_bootstrap_fd_(-1), |
| ml_service_path_(kDefaultMlServiceBinaryPath) {} |
| Process::~Process() = default; |
| |
| void Process::ControlProcessRun() { |
| // We need to set euid to kMlServiceDBusUid to bootstrap DBus. Otherwise, DBus |
| // will block us because our euid inside of the userns is 0 but is 20106 |
| // outside of the userns. |
| if (seteuid(kMlServiceDBusUid) != 0) { |
| RecordProcessErrorEvent(ProcessError::kChangeEuidToMlServiceDBusFailed); |
| LOG(ERROR) << "Unable to change effective uid to " << kMlServiceDBusUid; |
| exit(EX_OSERR); |
| } |
| |
| ml::Daemon daemon; |
| daemon.Run(); |
| } |
| |
| void Process::WorkerProcessRun() { |
| brillo::BaseMessageLoop message_loop; |
| message_loop.SetAsCurrent(); |
| DETACH_FROM_SEQUENCE(sequence_checker_); |
| mojo::core::Init(); |
| mojo::core::ScopedIPCSupport ipc_support( |
| base::ThreadTaskRunnerHandle::Get(), |
| mojo::core::ScopedIPCSupport::ShutdownPolicy::FAST); |
| mojo::IncomingInvitation invitation; |
| { |
| WallTimeMetric walltime_metric( |
| "MachineLearningService.WorkerProcessAcceptMojoConnectionTime"); |
| invitation = mojo::IncomingInvitation::Accept(mojo::PlatformChannelEndpoint( |
| mojo::PlatformHandle(base::ScopedFD(mojo_bootstrap_fd_)))); |
| } |
| mojo::ScopedMessagePipeHandle pipe = |
| invitation.ExtractMessagePipe(kInternalMojoPrimordialPipeName); |
| // The worker process exits if it disconnects with the control process. |
| // This can be important because in the control process's disconnect handler |
| // function we will use waitpid to wait for this process to finish. So |
| // the exit here will make sure that the waitpid in control process |
| // won't hang. |
| MachineLearningServiceImpl machine_learning_service_impl( |
| mojo::PendingReceiver< |
| chromeos::machine_learning::mojom::MachineLearningService>( |
| std::move(pipe)), |
| message_loop.QuitClosure()); |
| message_loop.Run(); |
| } |
| |
| const std::unordered_map<pid_t, Process::WorkerInfo>& |
| Process::GetWorkerPidInfoMap() { |
| DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_); |
| return worker_pid_info_map_; |
| } |
| |
| void Process::SetTypeForTesting(Type type) { |
| process_type_ = type; |
| } |
| |
| void Process::SetMlServicePathForTesting(const std::string& path) { |
| ml_service_path_ = path; |
| } |
| |
| void Process::SetBeforeExitWorkerDisconnectHandlerHookForTesting( |
| base::RepeatingClosure hook) { |
| before_exit_worker_disconnect_handler_hook_ = std::move(hook); |
| } |
| |
| bool Process::IsControlProcess() { |
| return process_type_ == Type::kControl || |
| process_type_ == Type::kControlForTest; |
| } |
| |
| bool Process::IsWorkerProcess() { |
| return process_type_ == Type::kWorker || |
| process_type_ == Type::kSingleProcessForTest; |
| } |
| |
| void Process::InternalPrimordialMojoPipeDisconnectHandler(pid_t child_pid) { |
| WallTimeMetric walltime_metric( |
| "MachineLearningService.WorkerProcessCleanUpTime"); |
| |
| UnregisterWorkerProcess(child_pid); |
| // Reap the worker process. |
| int status; |
| pid_t ret_pid = waitpid(child_pid, &status, 0); |
| DCHECK(ret_pid == child_pid); |
| int exit_status = WEXITSTATUS(status); |
| if (exit_status != 0) { |
| RecordWorkerProcessExitStatus(WEXITSTATUS(status)); |
| } |
| |
| // Call the hooks used in testing. |
| if (process_type_ == Type::kControlForTest && |
| !before_exit_worker_disconnect_handler_hook_.is_null()) { |
| before_exit_worker_disconnect_handler_hook_.Run(); |
| } |
| } |
| |
| } // namespace ml |