source/LifecycleMethod.cpp (161 lines of code) (raw):

/* * Copyright (c) Meta Platforms, Inc. and affiliates. * * This source code is licensed under the MIT license found in the * LICENSE file in the root directory of this source tree. */ #include <Creators.h> #include <SpartaWorkQueue.h> #include <mariana-trench/ClassHierarchies.h> #include <mariana-trench/JsonValidation.h> #include <mariana-trench/LifecycleMethod.h> #include <mariana-trench/Log.h> #include <mariana-trench/Methods.h> namespace marianatrench { LifecycleMethodCall LifecycleMethodCall::from_json(const Json::Value& value) { auto method_name = JsonValidation::string(value, "method_name"); auto return_type = JsonValidation::string(value, "return_type"); std::vector<std::string> argument_types; for (const auto& argument_type : JsonValidation::null_or_array(value, "argument_types")) { argument_types.emplace_back(JsonValidation::string(argument_type)); } return LifecycleMethodCall(method_name, return_type, argument_types); } DexMethodRef* MT_NULLABLE LifecycleMethodCall::get_dex_method(DexType* klass) const { const auto* return_type = DexType::get_type(DexString::make_string(return_type_)); if (return_type == nullptr) { ERROR(1, "Could not find return type `{}`.", return_type_); return nullptr; } const auto* argument_types = get_argument_types(); if (argument_types == nullptr) { return nullptr; } return DexMethod::get_method( /* type */ klass, /* name */ DexString::make_string(method_name_), /* proto */ DexProto::make_proto(return_type, argument_types)); } const DexTypeList* MT_NULLABLE LifecycleMethodCall::get_argument_types() const { DexTypeList::ContainerType argument_types; for (const auto& argument_type : argument_types_) { auto* type = DexType::get_type(argument_type); if (type == nullptr) { ERROR(1, "Could not find argument type `{}`.", argument_type); return nullptr; } argument_types.push_back(type); } return DexTypeList::make_type_list(std::move(argument_types)); } bool LifecycleMethodCall::operator==(const LifecycleMethodCall& other) const { return method_name_ == other.method_name_ && return_type_ == other.return_type_ && argument_types_ == other.argument_types_; } LifecycleMethod LifecycleMethod::from_json(const Json::Value& value) { auto base_class_name = JsonValidation::string(value, "base_class_name"); auto method_name = JsonValidation::string(value, "method_name"); std::vector<LifecycleMethodCall> callees; for (const auto& callee : JsonValidation::nonempty_array(value, "callees")) { callees.emplace_back(LifecycleMethodCall::from_json(callee)); } return LifecycleMethod(base_class_name, method_name, callees); } void LifecycleMethod::create_methods( const ClassHierarchies& class_hierarchies, Methods& methods) const { // All DexMethods created by `LifecycleMethod` have the same signature: // void <method_name_>(<arguments>) // The arguments are determined by the callees' arguments. This creates the // map of argument type -> location/position (first argument is at index 1). // The position corresponds to the register location containing the argument // in the DexMethod's code. The register location will be used to create the // invoke operation for methods that take a given DexType* as its argument. TypeIndexMap type_index_map; for (const auto& callee : callees_) { const auto* type_list = callee.get_argument_types(); if (type_list == nullptr) { ERROR(1, "Callee `{}` has invalid argument types.", callee.to_string()); continue; } for (auto* type : *type_list) { type_index_map.emplace(type, type_index_map.size() + 1); } } auto* MT_NULLABLE base_class_type = DexType::get_type(base_class_name_); if (!base_class_type) { WARNING( 1, "Could not find type for base class name `{}`. Will skip creating life-cycle methods.", base_class_name_); return; } std::atomic<int> methods_created_count = 0; const auto& children = class_hierarchies.extends(base_class_type); LOG(3, "Found {} child(ren) for type `{}`.", children.size(), base_class_name_); auto queue = sparta::work_queue<DexType*>([&](DexType* child) { const auto* method = create_dex_method(child, type_index_map); if (method != nullptr) { ++methods_created_count; methods.create(method); } }); for (const auto* child : children) { queue.add_item(const_cast<DexType*>(child)); } queue.run_all(); LOG(1, "Created {} life-cycle methods for classes inheriting from `{}`", methods_created_count, base_class_name_); } bool LifecycleMethod::operator==(const LifecycleMethod& other) const { return base_class_name_ == other.base_class_name_ && method_name_ == other.method_name_ && callees_ == other.callees_; } const DexMethod* MT_NULLABLE LifecycleMethod::create_dex_method( DexType* klass, const TypeIndexMap& type_index_map) const { auto method = MethodCreator( /* class */ klass, /* name */ DexString::make_string(method_name_), /* proto */ DexProto::make_proto(type::_void(), get_argument_types(type_index_map)), /* access */ DexAccessFlags::ACC_PRIVATE); auto this_location = method.get_local(0); auto* main_block = method.get_main_block(); mt_assert(main_block != nullptr); int callee_count = 0; for (const auto& callee : callees_) { auto* dex_method = callee.get_dex_method(klass); if (!dex_method) { // This can be null if `klass` does not override the method, in which // case, it will not be invoked. continue; } ++callee_count; std::vector<Location> invoke_with_registers{this_location}; auto* type_list = callee.get_argument_types(); // This should have been verified at the start of `create_methods` mt_assert(type_list != nullptr); for (auto* type : *type_list) { auto argument_register = method.get_local(type_index_map.at(type)); invoke_with_registers.push_back(argument_register); } main_block->invoke( IROpcode::OPCODE_INVOKE_DIRECT, dex_method, invoke_with_registers); } if (callee_count < 2) { // The point of life-cycle methods is to find flows where tainted member // variables flow from one callee into another. If only one (or no) method // is overridden, there is no need to create the artificial method. LOG(5, "Skipped creating life-cycle method for class `{}`. Reason: Insufficient callees.", klass->str()); return nullptr; } // The CFG needs to be built for the call graph to be constructed later. auto* new_method = method.create(); mt_assert(new_method != nullptr && new_method->get_code() != nullptr); new_method->get_code()->build_cfg(); LOG(5, "Created life-cycle method `{}`", new_method->str()); return new_method; } const DexTypeList* LifecycleMethod::get_argument_types( const TypeIndexMap& type_index_map) const { // While the register locations for the arguments start at 1, the actual // argument index for the method's prototype start at index 0. int num_args = type_index_map.size(); DexTypeList::ContainerType argument_types(num_args, nullptr); for (const auto& [type, pos] : type_index_map) { mt_assert(pos > 0); // 0 is for "this", should not be in the map. mt_assert(static_cast<size_t>(pos) - 1 < argument_types.max_size()); argument_types[pos - 1] = type; } return DexTypeList::make_type_list(std::move(argument_types)); } } // namespace marianatrench