From 5b144a10611350b3854a0806a5638b6cfa9b484a Mon Sep 17 00:00:00 2001
From: Ken Chen <cken@google.com>
Date: Thu, 26 Oct 2023 14:00:23 +0800
Subject: [PATCH] Add Data Saver utils in Firewall test helper class

Resolv_integration_test needs the function to toggle Data Saver settings
in BPF map.

Bug: 288340533
Test: atest resolv_integration_test
Change-Id: I1725d3c8c6fc641b855c935d7289a99f78dd5933
---
 tests/native/utilities/firewall.cpp | 20 ++++++++++++++++++++
 tests/native/utilities/firewall.h   |  4 +++-
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/tests/native/utilities/firewall.cpp b/tests/native/utilities/firewall.cpp
index e4669cbb91..6e35d07efd 100644
--- a/tests/native/utilities/firewall.cpp
+++ b/tests/native/utilities/firewall.cpp
@@ -27,6 +27,9 @@ Firewall::Firewall() {
 
     result = mUidOwnerMap.init(UID_OWNER_MAP_PATH);
     EXPECT_RESULT_OK(result) << "init mUidOwnerMap failed";
+
+    result = mDataSaverEnabledMap.init(DATA_SAVER_ENABLED_MAP_PATH);
+    EXPECT_RESULT_OK(result) << "init mDataSaverEnabledMap failed";
 }
 
 Firewall* Firewall::getInstance() {
@@ -116,3 +119,20 @@ Result<void> Firewall::removeUidInterfaceRules(const std::vector<int32_t>& uids)
     }
     return {};
 }
+
+Result<bool> Firewall::getDataSaverSetting() {
+    std::lock_guard guard(mMutex);
+    auto dataSaverSetting = mDataSaverEnabledMap.readValue(DATA_SAVER_ENABLED_KEY);
+    if (!dataSaverSetting.ok()) {
+        return Errorf("Cannot read the data saver setting: {}", dataSaverSetting.error().message());
+    }
+    return dataSaverSetting;
+}
+
+Result<void> Firewall::setDataSaver(bool enabled) {
+    std::lock_guard guard(mMutex);
+    auto res = mDataSaverEnabledMap.writeValue(DATA_SAVER_ENABLED_KEY, enabled, BPF_EXIST);
+    if (!res.ok()) return Errorf("Failed to set data saver: {}", res.error().message());
+
+    return {};
+}
diff --git a/tests/native/utilities/firewall.h b/tests/native/utilities/firewall.h
index 1e7e98720d..b3d69bfa80 100644
--- a/tests/native/utilities/firewall.h
+++ b/tests/native/utilities/firewall.h
@@ -33,9 +33,11 @@ class Firewall {
     Result<void> removeRule(uint32_t uid, UidOwnerMatchType match) EXCLUDES(mMutex);
     Result<void> addUidInterfaceRules(const std::string& ifName, const std::vector<int32_t>& uids);
     Result<void> removeUidInterfaceRules(const std::vector<int32_t>& uids);
-
+    Result<bool> getDataSaverSetting();
+    Result<void> setDataSaver(bool enabled);
   private:
     BpfMap<uint32_t, uint32_t> mConfigurationMap GUARDED_BY(mMutex);
     BpfMap<uint32_t, UidOwnerValue> mUidOwnerMap GUARDED_BY(mMutex);
+    BpfMap<uint32_t, bool> mDataSaverEnabledMap GUARDED_BY(mMutex);
     std::mutex mMutex;
 };