// Copyright 2024 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
#include "net/test/win/fake_network_cost_manager.h"
#include <netlistmgr.h>
#include <wrl/implements.h>
#include <map>
#include "base/task/sequenced_task_runner.h"
#include "net/base/network_cost_change_notifier_win.h"
using Microsoft::WRL::ClassicCom;
using Microsoft::WRL::ComPtr;
using Microsoft::WRL::RuntimeClass;
using Microsoft::WRL::RuntimeClassFlags;
namespace net {
namespace {
DWORD NlmConnectionCostFlagsFromConnectionCost(
NetworkChangeNotifier::ConnectionCost source_cost) {
switch (source_cost) {
case NetworkChangeNotifier::ConnectionCost::CONNECTION_COST_UNMETERED:
return (NLM_CONNECTION_COST_UNRESTRICTED | NLM_CONNECTION_COST_CONGESTED);
case NetworkChangeNotifier::ConnectionCost::CONNECTION_COST_METERED:
return (NLM_CONNECTION_COST_VARIABLE | NLM_CONNECTION_COST_ROAMING |
NLM_CONNECTION_COST_APPROACHINGDATALIMIT);
case NetworkChangeNotifier::ConnectionCost::CONNECTION_COST_UNKNOWN:
default:
return NLM_CONNECTION_COST_UNKNOWN;
}
}
void DispatchCostChangedEvent(ComPtr<INetworkCostManagerEvents> event_target,
DWORD cost) {
std::ignore =
event_target->CostChanged(cost, /*destination_address=*/nullptr);
}
} // namespace
// A fake implementation of `INetworkCostManager` that can simulate costs,
// changed costs and errors.
class FakeNetworkCostManager final
: public RuntimeClass<RuntimeClassFlags<ClassicCom>,
INetworkCostManager,
IConnectionPointContainer,
IConnectionPoint> {
public:
FakeNetworkCostManager(NetworkChangeNotifier::ConnectionCost connection_cost,
NetworkCostManagerStatus error_status)
: error_status_(error_status), connection_cost_(connection_cost) {}
// For each event sink in `event_sinks_`, call
// `INetworkCostManagerEvents::CostChanged()` with `changed_cost` on the event
// sink's task runner.
void PostCostChangedEvents(
NetworkChangeNotifier::ConnectionCost changed_cost) {
DWORD cost_for_changed_event;
std::map</*event_sink_cookie=*/DWORD, EventSinkRegistration>
event_sinks_for_changed_event;
{
base::AutoLock auto_lock(member_lock_);
connection_cost_ = changed_cost;
cost_for_changed_event =
NlmConnectionCostFlagsFromConnectionCost(changed_cost);
// Get the snapshot of event sinks to notify. The snapshot collection
// creates a new `ComPtr` for each event sink, which increments each the
// event sink's reference count, ensuring that each event sink
// remains alive to receive the cost changed event notification.
event_sinks_for_changed_event = event_sinks_;
}
for (const auto& pair : event_sinks_for_changed_event) {
const auto& registration = pair.second;
registration.event_sink_task_runner_->PostTask(
FROM_HERE,
base::BindOnce(&DispatchCostChangedEvent, registration.event_sink_,
cost_for_changed_event));
}
}
// Implement the `INetworkCostManager` interface.
HRESULT
__stdcall GetCost(DWORD* cost,
NLM_SOCKADDR* destination_ip_address) override {
if (error_status_ == NetworkCostManagerStatus::kErrorGetCostFailed) {
return E_FAIL;
}
if (destination_ip_address != nullptr) {
NOTIMPLEMENTED();
return E_NOTIMPL;
}
{
base::AutoLock auto_lock(member_lock_);
*cost = NlmConnectionCostFlagsFromConnectionCost(connection_cost_);
}
return S_OK;
}
HRESULT __stdcall GetDataPlanStatus(
NLM_DATAPLAN_STATUS* data_plan_status,
NLM_SOCKADDR* destination_ip_address) override {
NOTIMPLEMENTED();
return E_NOTIMPL;
}
HRESULT __stdcall SetDestinationAddresses(
UINT32 length,
NLM_SOCKADDR* destination_ip_address_list,
VARIANT_BOOL append) override {
NOTIMPLEMENTED();
return E_NOTIMPL;
}
// Implement the `IConnectionPointContainer` interface.
HRESULT __stdcall FindConnectionPoint(REFIID connection_point_id,
IConnectionPoint** result) override {
if (error_status_ ==
NetworkCostManagerStatus::kErrorFindConnectionPointFailed) {
return E_ABORT;
}
if (connection_point_id != IID_INetworkCostManagerEvents) {
return E_NOINTERFACE;
}
*result = static_cast<IConnectionPoint*>(this);
AddRef();
return S_OK;
}
HRESULT __stdcall EnumConnectionPoints(
IEnumConnectionPoints** results) override {
NOTIMPLEMENTED();
return E_NOTIMPL;
}
// Implement the `IConnectionPoint` interface.
HRESULT __stdcall Advise(IUnknown* event_sink,
DWORD* event_sink_cookie) override {
if (error_status_ == NetworkCostManagerStatus::kErrorAdviseFailed) {
return E_NOT_VALID_STATE;
}
ComPtr<INetworkCostManagerEvents> cost_manager_event_sink;
HRESULT hr =
event_sink->QueryInterface(IID_PPV_ARGS(&cost_manager_event_sink));
if (hr != S_OK) {
return hr;
}
base::AutoLock auto_lock(member_lock_);
event_sinks_[next_event_sink_cookie_] = {
cost_manager_event_sink,
base::SequencedTaskRunner::GetCurrentDefault()};
*event_sink_cookie = next_event_sink_cookie_;
++next_event_sink_cookie_;
return S_OK;
}
HRESULT __stdcall Unadvise(DWORD event_sink_cookie) override {
base::AutoLock auto_lock(member_lock_);
auto it = event_sinks_.find(event_sink_cookie);
if (it == event_sinks_.end()) {
return ERROR_NOT_FOUND;
}
event_sinks_.erase(it);
return S_OK;
}
HRESULT __stdcall GetConnectionInterface(IID* result) override {
NOTIMPLEMENTED();
return E_NOTIMPL;
}
HRESULT __stdcall GetConnectionPointContainer(
IConnectionPointContainer** result) override {
NOTIMPLEMENTED();
return E_NOTIMPL;
}
HRESULT __stdcall EnumConnections(IEnumConnections** result) override {
NOTIMPLEMENTED();
return E_NOTIMPL;
}
// Implement the `IUnknown` interface.
HRESULT __stdcall QueryInterface(REFIID interface_id,
void** result) override {
if (error_status_ == NetworkCostManagerStatus::kErrorQueryInterfaceFailed) {
return E_NOINTERFACE;
}
return RuntimeClass<RuntimeClassFlags<ClassicCom>, INetworkCostManager,
IConnectionPointContainer,
IConnectionPoint>::QueryInterface(interface_id, result);
}
FakeNetworkCostManager(const FakeNetworkCostManager&) = delete;
FakeNetworkCostManager& operator=(const FakeNetworkCostManager&) = delete;
private:
// The error state for this `FakeNetworkCostManager` to simulate. Cannot be
// changed.
const NetworkCostManagerStatus error_status_;
// Synchronizes access to all members below.
base::Lock member_lock_;
NetworkChangeNotifier::ConnectionCost connection_cost_
GUARDED_BY(member_lock_);
DWORD next_event_sink_cookie_ GUARDED_BY(member_lock_) = 0;
struct EventSinkRegistration {
ComPtr<INetworkCostManagerEvents> event_sink_;
scoped_refptr<base::SequencedTaskRunner> event_sink_task_runner_;
};
std::map</*event_sink_cookie=*/DWORD, EventSinkRegistration> event_sinks_
GUARDED_BY(member_lock_);
};
FakeNetworkCostManagerEnvironment::FakeNetworkCostManagerEnvironment() {
// Set up `NetworkCostChangeNotifierWin` to use the fake OS APIs.
NetworkCostChangeNotifierWin::OverrideCoCreateInstanceForTesting(
base::BindRepeating(
&FakeNetworkCostManagerEnvironment::FakeCoCreateInstance,
base::Unretained(this)));
}
FakeNetworkCostManagerEnvironment::~FakeNetworkCostManagerEnvironment() {
// Restore `NetworkCostChangeNotifierWin` to use the real OS APIs.
NetworkCostChangeNotifierWin::OverrideCoCreateInstanceForTesting(
base::BindRepeating(&CoCreateInstance));
}
HRESULT FakeNetworkCostManagerEnvironment::FakeCoCreateInstance(
REFCLSID class_id,
LPUNKNOWN outer_aggregate,
DWORD context_flags,
REFIID interface_id,
LPVOID* result) {
NetworkChangeNotifier::ConnectionCost connection_cost_for_new_instance;
NetworkCostManagerStatus error_status_for_new_instance;
{
base::AutoLock auto_lock(member_lock_);
connection_cost_for_new_instance = connection_cost_;
error_status_for_new_instance = error_status_;
}
if (error_status_for_new_instance ==
NetworkCostManagerStatus::kErrorCoCreateInstanceFailed) {
return E_ACCESSDENIED;
}
if (class_id != CLSID_NetworkListManager) {
return E_NOINTERFACE;
}
if (interface_id != IID_INetworkCostManager) {
return E_NOINTERFACE;
}
ComPtr<FakeNetworkCostManager> instance =
Microsoft::WRL::Make<FakeNetworkCostManager>(
connection_cost_for_new_instance, error_status_for_new_instance);
{
base::AutoLock auto_lock(member_lock_);
fake_network_cost_managers_.push_back(instance);
}
*result = instance.Detach();
return S_OK;
}
void FakeNetworkCostManagerEnvironment::SetCost(
NetworkChangeNotifier::ConnectionCost value) {
// Update the cost for each `INetworkCostManager` instance in
// `fake_network_cost_managers_`.
std::vector<Microsoft::WRL::ComPtr<FakeNetworkCostManager>>
fake_network_cost_managers_for_change_event;
{
base::AutoLock auto_lock(member_lock_);
connection_cost_ = value;
fake_network_cost_managers_for_change_event = fake_network_cost_managers_;
}
for (const auto& network_cost_manager :
fake_network_cost_managers_for_change_event) {
network_cost_manager->PostCostChangedEvents(/*connection_cost=*/value);
}
}
void FakeNetworkCostManagerEnvironment::SimulateError(
NetworkCostManagerStatus error_status) {
base::AutoLock auto_lock(member_lock_);
error_status_ = error_status;
}
} // namespace net
Codebase Browser
chromium