| // Copyright 2018 The ChromiumOS Authors |
| // Use of this source code is governed by a BSD-style license that can be |
| // found in the LICENSE file. |
| |
| #include "runtime_probe/probe_function.h" |
| |
| #include <memory> |
| #include <optional> |
| #include <utility> |
| |
| #include <base/check.h> |
| #include <base/json/json_reader.h> |
| #include <base/json/json_writer.h> |
| #include <base/logging.h> |
| |
| #include "runtime_probe/functions/all_functions.h" |
| #include "runtime_probe/system/context.h" |
| |
| namespace runtime_probe { |
| |
| using DataType = typename ProbeFunction::DataType; |
| |
| auto ProbeFunction::registered_functions_ = |
| AllFunctions::ConstructRegisteredFunctionTable(); |
| |
| ProbeFunction::ProbeFunction() = default; |
| |
| ProbeFunction::~ProbeFunction() = default; |
| |
| std::unique_ptr<ProbeFunction> ProbeFunction::FromValue(const base::Value& dv) { |
| if (!dv.is_dict()) { |
| LOG(ERROR) << "ProbeFunction::FromValue takes a dictionary as parameter"; |
| return nullptr; |
| } |
| const auto& dict = dv.GetDict(); |
| |
| auto size = dict.size(); |
| if (size == 0) { |
| LOG(ERROR) << "No function name found in the ProbeFunction dictionary"; |
| return nullptr; |
| } |
| |
| if (size > 1) { |
| LOG(ERROR) << "More than 1 function names specified in the ProbeFunction" |
| " dictionary"; |
| return nullptr; |
| } |
| |
| const auto& it = dict.begin(); |
| |
| // function_name is the only key exists in the dictionary */ |
| const auto& function_name = it->first; |
| const auto& kwargs = it->second; |
| |
| if (registered_functions_.find(function_name) == |
| registered_functions_.end()) { |
| // TODO(stimim): Should report an error. |
| LOG(ERROR) << "Function \"" << function_name << "\" not found"; |
| return nullptr; |
| } |
| |
| if (!kwargs.is_dict()) { |
| LOG(ERROR) << "Function argument should be a dictionary"; |
| return nullptr; |
| } |
| |
| return static_cast<std::unique_ptr<ProbeFunction>>( |
| registered_functions_[function_name](kwargs.GetDict())); |
| } |
| |
| void ProbeFunction::Eval(base::OnceCallback<void(DataType)> callback) const { |
| EvalAsyncImpl(std::move(callback)); |
| } |
| |
| int ProbeFunction::EvalInHelper(std::string* /*output*/) const { |
| LOG(ERROR) << "Probe function \"" << GetFunctionName() |
| << "\" cannot be invoked in helper."; |
| return -1; |
| } |
| |
| void ProbeFunction::RegisterArgumentParser(const std::string field_name, |
| ArgumentParser* parser) { |
| CHECK(!argument_parsers_.count(field_name)) |
| << "Register duplicated argument " << field_name; |
| argument_parsers_[field_name] = parser; |
| } |
| |
| bool ProbeFunction::ParseArguments(const base::Value::Dict& arguments) { |
| arguments_ = arguments.Clone(); |
| auto arguments_clone = arguments.Clone(); |
| bool success = true; |
| for (const auto& [field_name, parser] : argument_parsers_) { |
| auto value = arguments_clone.Extract(field_name); |
| std::string err; |
| if (parser->Parse(value, err)) { |
| continue; |
| } |
| success = false; |
| LOG(ERROR) << "ProbeFunction \"" << GetFunctionName() |
| << "\" failed to parse argument \"" << field_name |
| << "\": " << err; |
| } |
| if (!arguments_clone.empty()) { |
| success = false; |
| for (const auto& [field_name, unused_value] : arguments_clone) { |
| LOG(ERROR) << "ProbeFunction \"" << GetFunctionName() |
| << "\" got unexpected argument \"" << field_name << "\""; |
| } |
| } |
| return success ? PostParseArguments() : false; |
| } |
| |
| DataType ProbeFunction::EvalImpl() const { |
| NOTREACHED_NORETURN() |
| << "Either |EvalImpl| or |EvalAsyncImpl| should be implemented."; |
| } |
| |
| void ProbeFunction::EvalAsyncImpl( |
| base::OnceCallback<void(DataType)> callback) const { |
| std::move(callback).Run(EvalImpl()); |
| } |
| |
| bool PrivilegedProbeFunction::InvokeHelper(std::string* result) const { |
| base::Value::Dict probe_statement; |
| probe_statement.Set(GetFunctionName(), arguments().Clone()); |
| std::string probe_statement_str; |
| base::JSONWriter::Write(probe_statement, &probe_statement_str); |
| |
| return Context::Get()->helper_invoker()->Invoke( |
| /*probe_function=*/this, probe_statement_str, result); |
| } |
| |
| std::optional<base::Value> PrivilegedProbeFunction::InvokeHelperToJSON() const { |
| std::string raw_output; |
| if (!InvokeHelper(&raw_output)) { |
| LOG(ERROR) << "Failed to invoke helper."; |
| return std::nullopt; |
| } |
| |
| auto json_output = base::JSONReader::Read(raw_output); |
| if (!json_output) { |
| LOG(ERROR) << "Failed to parse output into json format."; |
| VLOG(3) << "InvokeHelper raw output:\n" << raw_output; |
| return std::nullopt; |
| } |
| |
| return json_output; |
| } |
| |
| int PrivilegedProbeFunction::EvalInHelper(std::string* output) const { |
| DLOG(INFO) << "Invoking probe function \"" << GetFunctionName() |
| << "\" in helper."; |
| base::Value result{EvalImpl()}; |
| if (base::JSONWriter::Write(result, output)) |
| return 0; |
| LOG(ERROR) << "Failed to serialize probed result to json string"; |
| return -1; |
| } |
| |
| void PrivilegedProbeFunction::Eval( |
| base::OnceCallback<void(DataType)> callback) const { |
| auto json_output = InvokeHelperToJSON(); |
| if (!json_output) { |
| std::move(callback).Run({}); |
| return; |
| } |
| if (!json_output->is_list()) { |
| LOG(ERROR) << "Failed to parse json output as list."; |
| VLOG(3) << "InvokeHelper output:\n" << *json_output; |
| std::move(callback).Run({}); |
| return; |
| } |
| |
| DataType result = std::move(json_output->GetList()); |
| PostHelperEvalImpl(&result); |
| VLOG(3) << GetFunctionName() << " Eval output:\n" << result; |
| std::move(callback).Run(std::move(result)); |
| } |
| |
| } // namespace runtime_probe |