aconfig storage daemon: support local flag override

1, local flag override will be stored in a separate protobuf file. See
the proto definition on aconfigd.proto file.

2, when flag override source is local, write to the pb file instead of
mapped storage file. But mark the flag as both has override is sticky in
the mapped flag info file. has a local override does not prevent server
side flag override get written to the memory mapped flag value files.

3, each time when the boot flag value copy is made, the local overrides
in the protobuf are applied onto the boot copy.

Bug: b/312444587
Test: atest aconfigd_test
Change-Id: Ied30d8c85d082e99895374e632ce32d9b0aab14e

8 files changed, 845 insertions, 300 deletions

diff --git a/aconfigd/Android.bp b/aconfigd/Android.bp
index 0f128a3..e8aadd6 100644
--- a/aconfigd/Android.bp
+++ b/aconfigd/Android.bp
@@ -5,6 +5,7 @@ cc_binary {
     "aconfigd.proto",
     "aconfigd_main.cpp",
     "aconfigd_util.cpp",
+    "aconfigd_mapped_file.cpp"
   ],
   static_libs: [
     "libaconfig_new_storage_flags",
diff --git a/aconfigd/aconfigd.cpp b/aconfigd/aconfigd.cpp
index e54fc49..96d3490 100644
--- a/aconfigd/aconfigd.cpp
+++ b/aconfigd/aconfigd.cpp
@@ -16,17 +16,13 @@
 
 #include <string>
 #include <unordered_map>
-
 #include <dirent.h>
 
 #include <android-base/file.h>
 #include <android-base/logging.h>
-#include <cutils/sockets.h>
-
-#include <aconfig_storage/aconfig_storage_read_api.hpp>
-#include <aconfig_storage/aconfig_storage_write_api.hpp>
 #include <protos/aconfig_storage_metadata.pb.h>
 
+#include "aconfigd_mapped_file.h"
 #include "aconfigd_util.h"
 #include "aconfigd.h"
 
@@ -46,6 +42,7 @@ struct StorageRecord {
   std::string flag_map;
   std::string flag_val;
   std::string flag_info;
+  std::string local_overrides;
   int timestamp;
 
   StorageRecord() = default;
@@ -57,6 +54,7 @@ struct StorageRecord {
       , flag_map(entry.flag_map())
       , flag_val(entry.flag_val())
       , flag_info(entry.flag_info())
+      , local_overrides(entry.local_overrides())
       , timestamp(entry.timestamp())
   {}
 };
@@ -67,8 +65,8 @@ using StorageRecords = std::unordered_map<std::string, StorageRecord>;
 /// In memory storage file records. Parsed from the pb.
 static StorageRecords persist_storage_records;
 
-/// In memory cache for package to container mapping
-static std::unordered_map<std::string, std::string> container_map;
+/// Mapped files manager
+static MappedFilesManager mapped_files_manager;
 
 namespace {
 
@@ -83,12 +81,48 @@ Result<void> WritePersistentStorageRecordsToFile() {
     record_pb->set_flag_map(entry.flag_map);
     record_pb->set_flag_val(entry.flag_val);
     record_pb->set_flag_info(entry.flag_info);
+    record_pb->set_local_overrides(entry.local_overrides);
     record_pb->set_timestamp(entry.timestamp);
   }
 
-  return WriteStorageRecordsPbToFile(records_pb, kPersistentStorageRecordsFileName);
+  return WritePbToFile<storage_records_pb>(records_pb, kPersistentStorageRecordsFileName);
+}
+
+Result<void> ApplyLocalOverridesToBootCopy(const std::string& container,
+                                           const std::string& flag_value_file) {
+  auto const& entry = persist_storage_records[container];
+  auto overrides_pb = ReadPbFromFile<LocalFlagOverrides>(entry.local_overrides);
+  if (!overrides_pb.ok()) {
+    return Error() << "Unable to read local overrides pb: " << overrides_pb.error();
+  }
+
+  // change boot flag value file to 0644 to allow write
+  if (chmod(flag_value_file.c_str(), 0644) == -1) {
+    return base::ErrnoError() << "chmod() failed";
+  };
+
+  auto& mapped_files = mapped_files_manager.get_mapped_files(container);
+  auto applied_pb = mapped_files.ApplyLocalOverride(flag_value_file, *overrides_pb);
+  if (!applied_pb.ok()) {
+    return base::Error() << "Failed to apply local override: " << applied_pb.error();
+  }
+
+  // change boot flag value file back to 0444
+  if (chmod(flag_value_file.c_str(), 0444) == -1) {
+    return base::ErrnoError() << "chmod() failed";
+  };
+
+  if (overrides_pb->overrides_size() != applied_pb->overrides_size()) {
+    auto result = WritePbToFile<LocalFlagOverrides>(*applied_pb, entry.local_overrides);
+    if (!result.ok()) {
+      return base::Error() << result.error();
+    }
+  }
+
+  return {};
 }
 
+
 /// Create boot flag value copy for a container
 Result<void> CreateBootSnapshotForContainer(const std::string& container) {
   // check existence persistent storage copy
@@ -117,6 +151,11 @@ Result<void> CreateBootSnapshotForContainer(const std::string& container) {
                    << copy_result.error();
   }
 
+  auto apply_result = ApplyLocalOverridesToBootCopy(container, dst_value_file);
+  if (!apply_result.ok()) {
+    return Error() << "Failed to apply local overrides: " << apply_result.error();
+  }
+
   copy_result = CopyFile(src_info_file, dst_info_file, 0444);
   if (!copy_result.ok()) {
     return Error() << "CopyFile failed for " << src_info_file << " :"
@@ -125,7 +164,7 @@ Result<void> CreateBootSnapshotForContainer(const std::string& container) {
 
   // update available storage records pb
   auto const& entry = persist_storage_records[container];
-  auto records_pb = ReadStorageRecordsPb(kAvailableStorageRecordsFileName);
+  auto records_pb = ReadPbFromFile<storage_records_pb>(kAvailableStorageRecordsFileName);
   if (!records_pb.ok()) {
     return Error() << "Unable to read available storage records: "
                    << records_pb.error();
@@ -140,7 +179,7 @@ Result<void> CreateBootSnapshotForContainer(const std::string& container) {
   record_pb->set_flag_info(dst_info_file);
   record_pb->set_timestamp(entry.timestamp);
 
-  auto write_result =  WriteStorageRecordsPbToFile(
+  auto write_result =  WritePbToFile<storage_records_pb>(
       *records_pb, kAvailableStorageRecordsFileName);
   if (!write_result.ok()) {
     return Error() << "Failed to write available storage records: "
@@ -150,11 +189,11 @@ Result<void> CreateBootSnapshotForContainer(const std::string& container) {
   return {};
 }
 
-/// Handle container update, returns if container has been updated
-Result<bool> HandleContainerUpdate(const std::string& container,
-                                   const std::string& package_file,
-                                   const std::string& flag_file,
-                                   const std::string& value_file) {
+/// Add a container storage if not existed, otherwise update if needed
+Result<bool> AddOrUpdateStorageForContainer(const std::string& container,
+                                            const std::string& package_file,
+                                            const std::string& flag_file,
+                                            const std::string& value_file) {
   auto timestamp = GetFileTimeStamp(value_file);
   if (!timestamp.ok()) {
     return Error() << "Failed to get timestamp of " << value_file
@@ -209,237 +248,145 @@ Result<bool> HandleContainerUpdate(const std::string& container,
   return false;
 }
 
-/// Find the container name given flag package name
-Result<std::string> FindContainer(const std::string& package) {
-  if (container_map.count(package)) {
-    return container_map[package];
-  }
-
-  auto records_pb = ReadStorageRecordsPb(kAvailableStorageRecordsFileName);
-  if (!records_pb.ok()) {
-    return Error() << "Unable to read available storage records: "
-                   << records_pb.error();
+/// Handle new storage request
+void HandleNewStorage(const StorageRequestMessage::NewStorageMessage& msg,
+                      StorageReturnMessage& return_msg) {
+  auto updated = AddOrUpdateStorageForContainer(
+      msg.container(), msg.package_map(), msg.flag_map(), msg.flag_value());
+  if (!updated.ok()) {
+    auto* errmsg = return_msg.mutable_error_message();
+    *errmsg = "Failed to add or update container " + msg.container()
+              + ": " + updated.error().message();
+    return;
   }
 
-  for (auto& entry : records_pb->files()) {
-    auto mapped_file = get_mapped_file(entry.container(), StorageFileType::package_map);
-    if (!mapped_file.ok()) {
-      return Error() << "Failed to map file for container " << entry.container()
-                     << ": " << mapped_file.error();
-    }
-
-    auto context= get_package_read_context(*mapped_file, package);
-    if (!context.ok()) {
-      return Error() << "Failed to get offset for package " << package
-                     << " from package map of " << entry.container() << " :"
-                     << context.error();
-    }
-
-    if (context->package_exists) {
-      container_map[package] = entry.container();
-      return entry.container();
-    }
+  auto copy = CreateBootSnapshotForContainer(msg.container());
+  if (!copy.ok()) {
+    auto* errmsg = return_msg.mutable_error_message();
+    *errmsg = "Failed to make a boot copy for " + msg.container()
+              + ": " + copy.error().message();
+    return;
   }
 
-  return Error() << "package not found";
+  return_msg.mutable_new_storage_message();
 }
 
-/// Find boolean flag offset in flag value file
-Result<std::pair<FlagValueType, uint32_t>> FindFlagContext(
-    const std::string& container,
-    const std::string& package,
-    const std::string& flag) {
-
-  auto package_map = get_mapped_file(container, StorageFileType::package_map);
-  if (!package_map.ok()) {
-    return Error() << "Failed to map package map file for " << container
-                   << ": " << package_map.error();
-  }
-
-  auto package_context = get_package_read_context(*package_map, package);
-  if (!package_context.ok()) {
-    return Error() << "Failed to get package offset of " << package
-                   << " in " << container  << " :" << package_context.error();
-  }
-
-  if (!package_context->package_exists) {
-    return Error() << package << " is not found in " << container;
-  }
-
-  uint32_t package_id = package_context->package_id;
-  uint32_t package_start_index = package_context->boolean_start_index;
-
-  auto flag_map = get_mapped_file(container, StorageFileType::flag_map);
-  if (!flag_map.ok()) {
-    return Error() << "Failed to map flag map file for " << container
-                   << ": " << flag_map.error();
-  }
-
-  auto flag_context = get_flag_read_context(*flag_map, package_id, flag);
-  if (!flag_context.ok()) {
-    return Error() << "Failed to get flag offset of " << flag
-                   << " in " << container  << " :" << flag_context.error();
+/// Handle a local flag override request
+Result<void> HandleLocalFlagOverride(const std::string& package,
+                                     const std::string& flag,
+                                     const std::string& flag_value) {
+  auto container = mapped_files_manager.GetContainer(package);
+  if (!container.ok()) {
+      return container.error();
   }
 
-  if (!flag_context->flag_exists) {
-    return Error() << flag << " is not found in " << container;
+  auto pb_file = persist_storage_records[*container].local_overrides;
+  auto pb = ReadPbFromFile<LocalFlagOverrides>(pb_file);
+  if (!pb.ok()) {
+    return Error() << "Failed to read pb from " << pb_file << ": " << pb.error();
   }
 
-  auto value_type = map_to_flag_value_type(flag_context->flag_type);
-  if (!value_type.ok()) {
-    return Error() << "Failed to get flag value type :" << value_type.error();
+  bool exist = false;
+  for (auto& entry : *(pb->mutable_overrides())) {
+    if (entry.package_name() == package && entry.flag_name() == flag) {
+      if (entry.flag_value() == flag_value) {
+        return {};
+      }
+      exist = true;
+      entry.set_flag_value(flag_value);
+      break;
+    }
   }
 
-  auto index = package_start_index + flag_context->flag_index;
-  return std::make_pair(*value_type, index);
-}
+  if (!exist) {
+    // add a new override entry to pb
+    auto new_override = pb->add_overrides();
+    new_override->set_package_name(package);
+    new_override->set_flag_name(flag);
+    new_override->set_flag_value(flag_value);
 
-/// Add a new storage
-Result<void> AddNewStorage(const std::string& container,
-                           const std::string& package_map,
-                           const std::string& flag_map,
-                           const std::string& flag_val) {
-  auto updated_result = HandleContainerUpdate(
-      container, package_map, flag_map, flag_val);
-  if (!updated_result.ok()) {
-    return Error() << "Failed to update container " << container
-                   << ":" << updated_result.error();
+    // mark override sticky
+    auto& mapped_files = mapped_files_manager.get_mapped_files(*container);
+    auto update = mapped_files.MarkStickyOverride(package, flag);
+    if (!update.ok()) {
+      return Error() << "Failed to mark flag " << package + "." + flag << " sticky: "
+                     << update.error();
+    }
   }
 
-  auto copy_result = CreateBootSnapshotForContainer(container);
-  if (!copy_result.ok()) {
-    return Error() << "Failed to make a boot copy: " << copy_result.error();
+  auto write = WritePbToFile<LocalFlagOverrides>(*pb, pb_file);
+  if (!write.ok()) {
+    return Error() << "Failed to write pb to " << pb_file << ": " << write.error();
   }
 
   return {};
 }
 
-/// Update persistent flag value
-Result<void> UpdatePersistentFlagValue(const std::string& package_name,
-                                       const std::string& flag_name,
-                                       const std::string& flag_value) {
-  auto container_result = FindContainer(package_name);
-  if (!container_result.ok()) {
-    return Error() << "Failed for find container for package " << package_name
-                   << ": " << container_result.error();
-  }
-  auto container = *container_result;
-
-  auto context_result = FindFlagContext(container, package_name, flag_name);
-  if (!context_result.ok()) {
-    return Error() << "Failed to obtain " << package_name << "."
-                   << flag_name << " flag value offset: " << context_result.error();
-  }
-
-  auto mapped_value_file = get_mutable_mapped_file(container, StorageFileType::flag_val);
-  if (!mapped_value_file.ok()) {
-    return Error() << "Failed to map flag value file for " << container
-                   << ": " << mapped_value_file.error();
-  }
-
-  auto mapped_info_file = get_mutable_mapped_file(container, StorageFileType::flag_info);
-  if (!mapped_info_file.ok()) {
-    return Error() << "Failed to map flag info file for " << container
-                   << ": " << mapped_info_file.error();
+/// Handle a server flag override request
+Result<void> HandleServerFlagOverride(const std::string& package,
+                                      const std::string& flag,
+                                      const std::string& flag_value) {
+  auto container = mapped_files_manager.GetContainer(package);
+  if (!container.ok()) {
+      return container.error();
   }
+  auto& mapped_files = mapped_files_manager.get_mapped_files(*container);
+  return  mapped_files.UpdatePersistFlag(package, flag, flag_value);
+}
 
-  auto value_type = context_result->first;
-  auto flag_index = context_result->second;
-
-  switch (value_type) {
-    case FlagValueType::Boolean: {
-      if (flag_value != "true" && flag_value != "false") {
-        return Error() << "Invalid boolean flag value, it should be true|false";
-      }
-      auto update_result = set_boolean_flag_value(
-          *mapped_value_file, flag_index, flag_value == "true");
-      if (!update_result.ok()) {
-        return Error() << "Failed to update flag value: " << update_result.error();
-      }
-      update_result = set_flag_has_override(
-          *mapped_info_file, value_type, flag_index, true);
-      if (!update_result.ok()) {
-        return Error() << "Failed to update flag has override: " << update_result.error();
-      }
-      break;
-    }
-    default:
-      return Error() << "Unsupported flag value type";
+/// Handle a flag override request
+void HandleFlagOverride(const StorageRequestMessage::FlagOverrideMessage& msg,
+                        StorageReturnMessage& return_msg) {
+  auto result = Result<void>();
+  if (msg.is_local()) {
+    result = HandleLocalFlagOverride(msg.package_name(),
+                                     msg.flag_name(),
+                                     msg.flag_value());
+  } else {
+    result = HandleServerFlagOverride(msg.package_name(),
+                                      msg.flag_name(),
+                                      msg.flag_value());
+  }
+
+  if (!result.ok()) {
+    auto* errmsg = return_msg.mutable_error_message();
+    *errmsg = result.error().message();
+  } else {
+    return_msg.mutable_flag_override_message();
   }
-
-  return {};
 }
 
-/// Query persistent flag value and info
-Result<std::pair<std::string, uint8_t>> QueryPersistentFlag(
-    const std::string& package_name,
-    const std::string& flag_name) {
-  auto container = FindContainer(package_name);
+/// Handle a flag query request
+void HandlePersistFlagQuery(const StorageRequestMessage::FlagQueryMessage& msg,
+                            StorageReturnMessage& return_msg) {
+  auto container = mapped_files_manager.GetContainer(msg.package_name());
   if (!container.ok()) {
-    return Error() << "Failed for find container for package " << package_name
-                   << ": " << container.error();
-  }
-
-  auto context = FindFlagContext(*container, package_name, flag_name);
-  if (!context.ok()) {
-    return Error() << "Failed to obtain " << package_name << "."
-                   << flag_name << " flag value offset: " << context.error();
+    auto* errmsg = return_msg.mutable_error_message();
+    *errmsg = container.error().message();
+    return;
   }
-  auto value_type = context->first;
-  auto flag_index = context->second;
 
-  auto value_file = get_mutable_mapped_file(*container, StorageFileType::flag_val);
-  if (!value_file.ok()) {
-    return Error() << "Failed to map flag value file for " << *container
-                   << ": " << value_file.error();
-  }
-
-  auto info_file = get_mutable_mapped_file(*container, StorageFileType::flag_info);
-  if (!info_file.ok()) {
-    return Error() << "Failed to map flag info file for " << *container
-                   << ": " << info_file.error();
-  }
-
-  // return value
-  auto flag_value = std::string();
-  uint8_t flag_info = 0;
-
-  switch (value_type) {
-    case FlagValueType::Boolean: {
-      // get flag value
-      auto ro_value_file = MappedStorageFile();
-      ro_value_file.file_ptr = value_file->file_ptr;
-      ro_value_file.file_size = value_file->file_size;
-      auto value = get_boolean_flag_value(ro_value_file, flag_index);
-      if (!value.ok()) {
-        return Error() << "Failed to get flag value: " << value.error();
-      }
-      flag_value = *value ? "true" : "false";
-
-      // get flag attribute
-      auto ro_info_file = MappedStorageFile();
-      ro_info_file.file_ptr = info_file->file_ptr;
-      ro_info_file.file_size = info_file->file_size;
-      auto attribute = get_flag_attribute(ro_info_file, value_type, flag_index);
-      if (!attribute.ok()) {
-        return Error() << "Failed to get flag info: " << attribute.error();
-      }
-      flag_info = *attribute;
+  auto& mapped_files = mapped_files_manager.get_mapped_files(*container);
+  auto result = mapped_files.GetPersistFlagValueAndInfo(
+      msg.package_name(), msg.flag_name());
 
-      break;
-    }
-    default:
-      return Error() << "Unsupported flag value type";
+  if (!result.ok()) {
+    auto* errmsg = return_msg.mutable_error_message();
+    *errmsg = result.error().message();
+  } else {
+    auto result_msg = return_msg.mutable_flag_query_message();
+    result_msg->set_flag_value(result->first);
+    result_msg->set_is_sticky(result->second & FlagInfoBit::IsSticky);
+    result_msg->set_is_readwrite(result->second & FlagInfoBit::IsReadWrite);
+    result_msg->set_has_override(result->second & FlagInfoBit::HasOverride);
   }
-
-  return std::make_pair(flag_value, flag_info);
 }
 
 } // namespace
 
 /// Initialize in memory aconfig storage records
 Result<void> InitializeInMemoryStorageRecords() {
-  auto records_pb = ReadStorageRecordsPb(kPersistentStorageRecordsFileName);
+  auto records_pb = ReadPbFromFile<storage_records_pb>(kPersistentStorageRecordsFileName);
   if (!records_pb.ok()) {
     return Error() << "Unable to read persistent storage records: "
                    << records_pb.error();
@@ -470,7 +417,7 @@ Result<void> InitializePlatformStorage() {
       continue;
     }
 
-    auto updated = HandleContainerUpdate(
+    auto updated = AddOrUpdateStorageForContainer(
         container, package_file, flag_file, value_file);
     if (!updated.ok()) {
       return Error() << updated.error();
@@ -492,48 +439,19 @@ void HandleSocketRequest(const StorageRequestMessage& message,
     case StorageRequestMessage::kNewStorageMessage: {
       LOG(INFO) << "received a new storage request";
       auto msg = message.new_storage_message();
-      auto result = AddNewStorage(msg.container(),
-                                  msg.package_map(),
-                                  msg.flag_map(),
-                                  msg.flag_value());
-      if (!result.ok()) {
-        auto* errmsg = return_message.mutable_error_message();
-        *errmsg = result.error().message();
-      } else {
-        return_message.mutable_new_storage_message();
-      }
+      HandleNewStorage(msg, return_message);
       break;
     }
-
     case StorageRequestMessage::kFlagOverrideMessage: {
       LOG(INFO) << "received a flag override request";
       auto msg = message.flag_override_message();
-      auto result = UpdatePersistentFlagValue(msg.package_name(),
-                                              msg.flag_name(),
-                                              msg.flag_value());
-      if (!result.ok()) {
-        auto* errmsg = return_message.mutable_error_message();
-        *errmsg = result.error().message();
-      } else {
-        return_message.mutable_flag_override_message();
-      }
+      HandleFlagOverride(msg, return_message);
       break;
     }
-
     case StorageRequestMessage::kFlagQueryMessage: {
       LOG(INFO) << "received a flag query request";
       auto msg = message.flag_query_message();
-      auto result = QueryPersistentFlag(msg.package_name(), msg.flag_name());
-      if (!result.ok()) {
-        auto* errmsg = return_message.mutable_error_message();
-        *errmsg = result.error().message();
-      } else {
-        auto return_msg = return_message.mutable_flag_query_message();
-        return_msg->set_flag_value(result->first);
-        return_msg->set_is_sticky(result->second & FlagInfoBit::IsSticky);
-        return_msg->set_is_readwrite(result->second & FlagInfoBit::IsReadWrite);
-        return_msg->set_has_override(result->second & FlagInfoBit::HasOverride);
-      }
+      HandlePersistFlagQuery(msg, return_message);
       break;
     }
     default:
diff --git a/aconfigd/aconfigd.proto b/aconfigd/aconfigd.proto
index eb1d788..4adfa26 100644
--- a/aconfigd/aconfigd.proto
+++ b/aconfigd/aconfigd.proto
@@ -32,6 +32,7 @@ message StorageRequestMessage {
     optional string package_name = 1;
     optional string flag_name = 2;
     optional string flag_value = 3;
+    optional bool is_local = 4;
   }
 
   // query persistent flag value and info
@@ -75,3 +76,13 @@ message StorageReturnMessage {
 message StorageReturnMessages {
   repeated StorageReturnMessage msgs = 1;
 }
+
+message LocalFlagOverride {
+  optional string package_name = 1;
+  optional string flag_name = 2;
+  optional string flag_value = 3;
+}
+
+message LocalFlagOverrides {
+  repeated LocalFlagOverride overrides = 1;
+}
diff --git a/aconfigd/aconfigd_mapped_file.cpp b/aconfigd/aconfigd_mapped_file.cpp
new file mode 100644
index 0000000..ad16a30
--- /dev/null
+++ b/aconfigd/aconfigd_mapped_file.cpp
@@ -0,0 +1,459 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <protos/aconfig_storage_metadata.pb.h>
+
+#include "aconfigd.h"
+#include "aconfigd_util.h"
+#include "aconfigd_mapped_file.h"
+
+using namespace aconfig_storage;
+
+namespace android {
+  namespace aconfigd {
+
+  /// get mapped files for a container
+  MappedFiles& MappedFilesManager::get_mapped_files(
+      const std::string& container) {
+    if (mapped_files_.count(container) == 0) {
+      mapped_files_[container] = std::make_unique<MappedFiles>(container);
+    }
+    return *(mapped_files_[container]);
+  }
+
+  /// get container name given flag package name
+  base::Result<std::string> MappedFilesManager::GetContainer(
+      const std::string& package) {
+    if (package_to_container_.count(package)) {
+      return package_to_container_[package];
+    }
+
+    auto records_pb = ReadPbFromFile<aconfig_storage_metadata::storage_files>(
+        kAvailableStorageRecordsFileName);
+    if (!records_pb.ok()) {
+      return base::Error() << "Unable to read available storage records: "
+                           << records_pb.error();
+    }
+
+    for (auto& entry : records_pb->files()) {
+      auto& mapped_file = get_mapped_files(entry.container());
+      auto has_flag = mapped_file.HasPackage(package);
+      if (!has_flag.ok()) {
+        return base::Error() << has_flag.error();
+      }
+
+      if (*has_flag) {
+        package_to_container_[package] = entry.container();
+        return entry.container();
+      }
+    }
+
+    return base::Error() << "container not found";
+  }
+
+  /// constructor
+  MappedFiles::MappedFiles(const std::string& container)
+      : container_(container)
+      , package_map_(nullptr)
+      , flag_map_(nullptr)
+      , persist_flag_val_(nullptr)
+      , persist_flag_info_(nullptr) {
+  }
+
+  /// move constructor
+  MappedFiles::MappedFiles(MappedFiles&& rhs) {
+    if (this != &rhs) {
+      *this = std::move(rhs);
+    }
+  }
+
+  /// move assignment
+  MappedFiles& MappedFiles::operator=(MappedFiles&& rhs) {
+    if (this != &rhs) {
+      container_ = rhs.container_;
+      package_map_ = std::move(rhs.package_map_);
+      flag_map_ = std::move(rhs.flag_map_);
+      persist_flag_val_ = std::move(rhs.persist_flag_val_);
+      persist_flag_info_ = std::move(rhs.persist_flag_info_);
+    }
+    return *this;
+  }
+
+  /// map a storage file
+  base::Result<MappedStorageFile> MappedFiles::MapStorageFile(StorageFileType file_type) {
+    switch (file_type) {
+      case StorageFileType::package_map:
+      case StorageFileType::flag_map:
+      case StorageFileType::flag_info:
+        return get_mapped_file(container_, file_type);
+        break;
+      default:
+        return base::Error() << "Unsupported storage file type to map";
+    }
+  }
+
+  /// map a mutable storage file
+  base::Result<MutableMappedStorageFile> MappedFiles::MapMutableStorageFile(
+      StorageFileType file_type) {
+    switch (file_type) {
+      case StorageFileType::flag_val:
+      case StorageFileType::flag_info:
+        return get_mutable_mapped_file(container_, file_type);
+        break;
+      default:
+        return base::Error() << "Unsupported storage file type to map";
+    }
+  }
+
+  /// get package map
+  base::Result<const MappedStorageFile*> MappedFiles::get_package_map() {
+    if (!package_map_) {
+      auto package_map = MapStorageFile(StorageFileType::package_map);
+      if (!package_map.ok()) {
+        return base::Error() << "Failed to map package map file for " << container_
+                             << ": " << package_map.error();
+      }
+      package_map_.reset(new MappedStorageFile(*package_map));
+    }
+    return package_map_.get();
+  }
+
+  /// get flag map
+  base::Result<const MappedStorageFile*> MappedFiles::get_flag_map() {
+    if (!flag_map_) {
+      auto flag_map = MapStorageFile(StorageFileType::flag_map);
+      if (!flag_map.ok()) {
+        return base::Error() << "Failed to map flag map file for " << container_
+                             << ": " << flag_map.error();
+      }
+      flag_map_.reset(new MappedStorageFile(*flag_map));
+    }
+    return flag_map_.get();
+  }
+
+  /// get persist flag val
+  base::Result<const MutableMappedStorageFile*> MappedFiles::get_persist_flag_val() {
+    if (!persist_flag_val_) {
+      auto flag_val = MapMutableStorageFile(StorageFileType::flag_val);
+      if (!flag_val.ok()) {
+        return base::Error() << "Failed to map persist flag value file for " << container_
+                             << ": " << flag_val.error();
+      }
+      persist_flag_val_.reset(new MutableMappedStorageFile(*flag_val));
+    }
+    return persist_flag_val_.get();
+  }
+
+  /// get persist flag info
+  base::Result<const MutableMappedStorageFile*> MappedFiles::get_persist_flag_info() {
+    if (!persist_flag_info_) {
+      auto flag_info = MapMutableStorageFile(StorageFileType::flag_info);
+      if (!flag_info.ok()) {
+        return base::Error() << "Failed to map persist flag info file for " << container_
+                             << ": " << flag_info.error();
+      }
+      persist_flag_info_.reset(new MutableMappedStorageFile(*flag_info));
+    }
+    return persist_flag_info_.get();
+  }
+
+  /// Find flag value type and global index
+  base::Result<MappedFiles::FlagTypeAndIndex> MappedFiles::GetFlagTypeAndIndex(
+      const std::string& package,
+      const std::string& flag) {
+    auto result = FlagTypeAndIndex();
+
+    auto package_map = get_package_map();
+    if (!package_map.ok()) {
+      return base::Error() << package_map.error();
+    }
+
+    auto package_context = get_package_read_context(**package_map, package);
+    if (!package_context.ok()) {
+      return base::Error() << "Failed to get package context for " << package
+                           << " in " << container_  << " :" << package_context.error();
+    }
+
+    if (!package_context->package_exists) {
+      result.flag_exists = false;
+      return result;
+    }
+
+    uint32_t package_id = package_context->package_id;
+    uint32_t package_start_index = package_context->boolean_start_index;
+
+    auto flag_map = get_flag_map();
+    if (!flag_map.ok()) {
+      return base::Error() << flag_map.error();
+    }
+
+    auto flag_context = get_flag_read_context(**flag_map, package_id, flag);
+    if (!flag_context.ok()) {
+      return base::Error() << "Failed to get flag context of " << flag
+                           << " in " << container_  << " :" << flag_context.error();
+    }
+
+    if (!flag_context->flag_exists) {
+      result.flag_exists = false;
+      return result;
+    }
+
+    StoredFlagType stored_type = flag_context->flag_type;
+    uint16_t flag_index = flag_context->flag_index;
+
+    auto value_type = map_to_flag_value_type(stored_type);
+    if (!value_type.ok()) {
+      return base::Error() << "Failed to get flag value type :" << value_type.error();
+    }
+
+    result.flag_exists = true;
+    result.value_type = *value_type;
+    result.flag_index = package_start_index + flag_index;
+    return result;
+  }
+
+  /// check if has package
+  base::Result<bool> MappedFiles::HasPackage(const std::string& package) {
+    auto package_map_file = get_package_map();
+    if (!package_map_file.ok()) {
+      return base::Error() << package_map_file.error();
+    }
+
+    auto context= get_package_read_context(**package_map_file, package);
+    if (!context.ok()) {
+      return base::Error() << "Failed to get context for package " << package
+                           << " in container " << container_ << ": " << context.error();
+    }
+
+    return context->package_exists;
+  }
+
+  /// server flag override, update persistent flag value and info
+  base::Result<void> MappedFiles::UpdatePersistFlag(const std::string& package,
+                                                    const std::string& flag,
+                                                    const std::string& flag_value) {
+
+    // find flag value type and index
+    auto type_and_index = GetFlagTypeAndIndex(package, flag);
+    if (!type_and_index.ok()) {
+      return base::Error() << "Failed to find flag " << flag << ": "
+                           << type_and_index.error();
+    }
+    if (!type_and_index->flag_exists) {
+      return base::Error() << "Failed to find flag " << flag;
+    }
+    auto value_type = type_and_index->value_type;
+    auto flag_index = type_and_index->flag_index;
+
+    auto flag_value_file = get_persist_flag_val();
+    if (!flag_value_file.ok()) {
+      return base::Error() << flag_value_file.error();
+    }
+
+    auto flag_info_file = get_persist_flag_info();
+    if (!flag_info_file.ok()) {
+      return base::Error() << flag_info_file.error();
+    }
+
+    switch (value_type) {
+      case FlagValueType::Boolean: {
+        // validate value
+        if (flag_value != "true" && flag_value != "false") {
+          return base::Error() << "Invalid boolean flag value, it should be true|false";
+        }
+
+        // update flag value
+        auto update_result = set_boolean_flag_value(
+            **flag_value_file, flag_index, flag_value == "true");
+        if (!update_result.ok()) {
+          return base::Error() << "Failed to update flag value: " << update_result.error();
+        }
+
+        // update flag info
+        update_result = set_flag_has_override(
+            **flag_info_file, value_type, flag_index, true);
+        if (!update_result.ok()) {
+          return base::Error() << "Failed to update flag has override: " << update_result.error();
+        }
+        break;
+      }
+      default:
+        return base::Error() << "Unsupported flag value type";
+    }
+
+    return {};
+  }
+
+  /// mark this flag has sticky override
+  base::Result<void> MappedFiles::MarkStickyOverride(const std::string& package,
+                                                     const std::string& flag) {
+    // find flag value type and index
+    auto type_and_index = GetFlagTypeAndIndex(package, flag);
+    if (!type_and_index.ok()) {
+      return base::Error() << "Failed to find flag " << flag << ": "
+                           << type_and_index.error();
+    }
+    if (!type_and_index->flag_exists) {
+      return base::Error() << "Failed to find flag " << flag;
+
+    }
+
+    auto value_type = type_and_index->value_type;
+    auto flag_index = type_and_index->flag_index;
+
+    auto flag_info_file = get_persist_flag_info();
+    if (!flag_info_file.ok()) {
+      return base::Error() << flag_info_file.error();
+    }
+
+    // update flag info, has override
+    auto update_result = set_flag_has_override(
+        **flag_info_file, value_type, flag_index, true);
+    if (!update_result.ok()) {
+      return base::Error() << "Failed to update flag has override: " << update_result.error();
+    }
+
+    // update flag info, is sticky
+    update_result = set_flag_is_sticky(
+        **flag_info_file, value_type, flag_index, true);
+    if (!update_result.ok()) {
+      return base::Error() << "Failed to update flag is sticky: " << update_result.error();
+    }
+
+    return {};
+  }
+
+  /// get persistent flag value and info
+  base::Result<std::pair<std::string, uint8_t>> MappedFiles::GetPersistFlagValueAndInfo(
+      const std::string& package,
+      const std::string& flag) {
+
+    // find flag value type and index
+    auto type_and_index = GetFlagTypeAndIndex(package, flag);
+    if (!type_and_index.ok()) {
+      return base::Error() << "Failed to find flag " << flag << ": "
+                           << type_and_index.error();
+    }
+    if (!type_and_index->flag_exists) {
+      return base::Error() << "Failed to find flag " << flag;
+
+    }
+    auto value_type = type_and_index->value_type;
+    auto flag_index = type_and_index->flag_index;
+
+    auto flag_value_file = get_persist_flag_val();
+    if (!flag_value_file.ok()) {
+      return base::Error() << flag_value_file.error();
+    }
+
+    auto flag_info_file = get_persist_flag_info();
+    if (!flag_info_file.ok()) {
+      return base::Error() << flag_info_file.error();
+    }
+
+    // return value
+    auto flag_value = std::string();
+    uint8_t flag_attribute = 0;
+
+    switch (value_type) {
+      case FlagValueType::Boolean: {
+        // get flag value
+        auto ro_value_file = MappedStorageFile();
+        ro_value_file.file_ptr = (*flag_value_file)->file_ptr;
+        ro_value_file.file_size = (*flag_value_file)->file_size;
+        auto value = get_boolean_flag_value(ro_value_file, flag_index);
+        if (!value.ok()) {
+          return base::Error() << "Failed to get flag value: " << value.error();
+        }
+        flag_value = *value ? "true" : "false";
+
+        // get flag attribute
+        auto ro_info_file = MappedStorageFile();
+        ro_info_file.file_ptr = (*flag_info_file)->file_ptr;
+        ro_info_file.file_size = (*flag_info_file)->file_size;
+        auto attribute = get_flag_attribute(ro_info_file, value_type, flag_index);
+        if (!attribute.ok()) {
+          return base::Error() << "Failed to get flag info: " << attribute.error();
+        }
+        flag_attribute = *attribute;
+
+        break;
+      }
+      default:
+        return base::Error() << "Unsupported flag value type";
+    }
+
+    return std::make_pair(flag_value, flag_attribute);
+  }
+
+  /// apply local update to boot flag value copy
+  base::Result<LocalFlagOverrides> MappedFiles::ApplyLocalOverride(
+      const std::string& flag_value_file,
+      const LocalFlagOverrides& local_overrides) {
+    auto applied_overrides = LocalFlagOverrides();
+    auto mutable_flag_value_file = map_mutable_storage_file(flag_value_file);
+    if (!mutable_flag_value_file.ok()) {
+      return base::Error() << "Failed to map flag value file for local override: "
+                           << mutable_flag_value_file.error();
+    }
+
+    for (auto& entry : local_overrides.overrides()) {
+
+      // find flag value type and index
+      auto type_and_index = GetFlagTypeAndIndex(entry.package_name(), entry.flag_name());
+      if (!type_and_index.ok()) {
+        return base::Error() << "Failed to find flag: " << type_and_index.error();
+      }
+      if (!type_and_index->flag_exists) {
+        continue;
+      }
+
+      auto value_type = type_and_index->value_type;
+      auto flag_index = type_and_index->flag_index;
+
+      // apply a local override
+      switch (value_type) {
+        case FlagValueType::Boolean: {
+          // validate value
+          if (entry.flag_value() != "true" && entry.flag_value() != "false") {
+            return base::Error() << "Invalid boolean flag value, it should be true|false";
+          }
+
+          // update flag value
+          auto update_result = set_boolean_flag_value(
+              *mutable_flag_value_file, flag_index, entry.flag_value() == "true");
+          if (!update_result.ok()) {
+            return base::Error() << "Failed to update flag value: " << update_result.error();
+          }
+
+          break;
+        }
+        default:
+          return base::Error() << "Unsupported flag value type";
+      }
+
+      // mark it applied
+      auto new_applied = applied_overrides.add_overrides();
+      new_applied->set_package_name(entry.package_name());
+      new_applied->set_flag_name(entry.flag_name());
+      new_applied->set_flag_value(entry.flag_value());
+    }
+
+    return applied_overrides;
+  }
+
+  } // namespace aconfigd
+} // namespace android
diff --git a/aconfigd/aconfigd_mapped_file.h b/aconfigd/aconfigd_mapped_file.h
new file mode 100644
index 0000000..b615690
--- /dev/null
+++ b/aconfigd/aconfigd_mapped_file.h
@@ -0,0 +1,159 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <string>
+#include <memory>
+#include <unordered_map>
+
+#include <android-base/result.h>
+
+#include <aconfigd.pb.h>
+#include <aconfig_storage/aconfig_storage_read_api.hpp>
+#include <aconfig_storage/aconfig_storage_write_api.hpp>
+
+namespace android {
+  namespace aconfigd {
+
+    /// Mapped files for a container
+    class MappedFiles {
+      public:
+
+      /// constructor
+      MappedFiles(const std::string& container);
+
+      /// destructor
+      ~MappedFiles() = default;
+
+      /// no copy
+      MappedFiles(const MappedFiles&) = delete;
+      MappedFiles& operator=(const MappedFiles&) = delete;
+
+      /// move constructor and assignment
+      MappedFiles(MappedFiles&& rhs);
+      MappedFiles& operator=(MappedFiles&& rhs);
+
+      /// check if has package
+      base::Result<bool> HasPackage(const std::string& package);
+
+      /// server flag override, update persistent flag value and info
+      base::Result<void> UpdatePersistFlag(const std::string& package,
+                                           const std::string& flag,
+                                           const std::string& value);
+
+      /// mark this flag has sticky override
+      base::Result<void> MarkStickyOverride(const std::string& package,
+                                            const std::string& flag);
+
+      /// get persistent flag value and info
+      base::Result<std::pair<std::string, uint8_t>> GetPersistFlagValueAndInfo(
+          const std::string& package,
+          const std::string& flag);
+
+      /// apply local update to boot flag value copy, return stale local overrides
+      base::Result<LocalFlagOverrides> ApplyLocalOverride(
+          const std::string& flag_value_file,
+          const LocalFlagOverrides& pb);
+
+      private:
+
+      /// map a storage file
+      base::Result<aconfig_storage::MappedStorageFile> MapStorageFile(
+          aconfig_storage::StorageFileType file_type);
+
+      /// map a mutable storage file
+      base::Result<aconfig_storage::MutableMappedStorageFile> MapMutableStorageFile(
+          aconfig_storage::StorageFileType file_type);
+
+      /// get package map
+      base::Result<const aconfig_storage::MappedStorageFile*> get_package_map();
+
+      /// get flag map
+      base::Result<const aconfig_storage::MappedStorageFile*> get_flag_map();
+
+      /// get persist flag val
+      base::Result<const aconfig_storage::MutableMappedStorageFile*> get_persist_flag_val();
+
+      /// get persist flag info
+      base::Result<const aconfig_storage::MutableMappedStorageFile*> get_persist_flag_info();
+
+      /// return result for flag type and index query
+      struct FlagTypeAndIndex {
+        bool flag_exists;
+        aconfig_storage::FlagValueType value_type;
+        uint32_t flag_index;
+      };
+
+      /// Find flag value type and global index
+      base::Result<FlagTypeAndIndex> GetFlagTypeAndIndex(
+          const std::string& package, const std::string& flag);
+
+      private:
+
+      /// container name
+      std::string container_;
+
+      /// mapped package map file
+      std::unique_ptr<aconfig_storage::MappedStorageFile> package_map_;
+
+      /// mapped flag map file
+      std::unique_ptr<aconfig_storage::MappedStorageFile> flag_map_;
+
+      /// mapped mutable flag value file
+      std::unique_ptr<aconfig_storage::MutableMappedStorageFile> persist_flag_val_;
+
+      /// mapped mutable flag info file
+      std::unique_ptr<aconfig_storage::MutableMappedStorageFile> persist_flag_info_;
+
+    }; // class MappedFiles
+
+    /// Manager all mapped files across different containers
+    class MappedFilesManager {
+      public:
+
+      /// constructor
+      MappedFilesManager() = default;
+
+      /// destructor
+      ~MappedFilesManager() = default;
+
+      /// no copy
+      MappedFilesManager(const MappedFilesManager&) = delete;
+      MappedFilesManager& operator=(const MappedFilesManager&) = delete;
+
+      /// move constructor and assignment
+      MappedFilesManager(MappedFilesManager&& rhs) = default;
+      MappedFilesManager& operator=(MappedFilesManager&& rhs) = default;
+
+      /// get mapped files for a container
+      MappedFiles& get_mapped_files(const std::string& container);
+
+      /// get container name given flag package name
+      base::Result<std::string> GetContainer(const std::string& package);
+
+      private:
+
+      /// a hash table from container name to mapped files
+      std::unordered_map<std::string, std::unique_ptr<MappedFiles>> mapped_files_;
+
+      /// a hash table from package name to container name
+      std::unordered_map<std::string, std::string> package_to_container_;
+
+    }; // class MappedFilesManager
+
+  } // namespace aconfigd
+} // namespace android
diff --git a/aconfigd/aconfigd_test.cpp b/aconfigd/aconfigd_test.cpp
index 4ee6754..c1026a1 100644
--- a/aconfigd/aconfigd_test.cpp
+++ b/aconfigd/aconfigd_test.cpp
@@ -108,13 +108,15 @@ base::Result<StorageReturnMessages> send_new_storage_message() {
 
 base::Result<StorageReturnMessages> send_flag_override_message(const std::string& package,
                                                                const std::string& flag,
-                                                               const std::string& value) {
+                                                               const std::string& value,
+                                                               bool is_local) {
   auto messages = StorageRequestMessages{};
   auto* message = messages.add_msgs();
   auto* msg = message->mutable_flag_override_message();
   msg->set_package_name(package);
   msg->set_flag_name(flag);
   msg->set_flag_value(value);
+  msg->set_is_local(is_local);
   return send_message(messages);
 }
 
@@ -159,11 +161,12 @@ TEST(aconfigd_socket, flag_override_message) {
   ASSERT_TRUE(return_message.has_new_storage_message());
 
   auto flag_override_result = send_flag_override_message(
-      "com.android.aconfig.storage.test_1", "enabled_rw", "true");
+      "com.android.aconfig.storage.test_1", "enabled_rw", "true", false);
   ASSERT_TRUE(flag_override_result.ok()) << flag_override_result.error();
   ASSERT_EQ(flag_override_result->msgs_size(), 1);
   return_message = flag_override_result->msgs(0);
-  ASSERT_TRUE(return_message.has_flag_override_message());
+  ASSERT_TRUE(return_message.has_flag_override_message())
+      << return_message.error_message();
 
   auto flag_query_result = send_flag_query_message(
       "com.android.aconfig.storage.test_1", "enabled_rw");
@@ -178,7 +181,7 @@ TEST(aconfigd_socket, flag_override_message) {
   ASSERT_EQ(query.has_override(), true);
 
   flag_override_result = send_flag_override_message(
-      "com.android.aconfig.storage.test_1", "enabled_rw", "false");
+      "com.android.aconfig.storage.test_1", "enabled_rw", "false", false);
   ASSERT_TRUE(flag_override_result.ok()) << flag_override_result.error();
   ASSERT_EQ(flag_override_result->msgs_size(), 1);
   return_message = flag_override_result->msgs(0);
@@ -205,13 +208,13 @@ TEST(aconfigd_socket, nonexist_flag_override_message) {
   ASSERT_TRUE(return_message.has_new_storage_message());
 
   auto flag_override_result = send_flag_override_message(
-      "com.android.aconfig.storage.test_1", "unknown", "true");
+      "com.android.aconfig.storage.test_1", "unknown", "true", false);
   ASSERT_TRUE(flag_override_result.ok()) << flag_override_result.error();
   ASSERT_EQ(flag_override_result->msgs_size(), 1);
   return_message = flag_override_result->msgs(0);
   ASSERT_TRUE(return_message.has_error_message());
   auto errmsg = return_message.error_message();
-  ASSERT_TRUE(errmsg.find("unknown is not found in mockup") != std::string::npos);
+  ASSERT_TRUE(errmsg.find("Failed to find flag unknown") != std::string::npos);
 }
 
 TEST(aconfigd_socket, nonexist_flag_query_message) {
@@ -229,7 +232,7 @@ TEST(aconfigd_socket, nonexist_flag_query_message) {
   ASSERT_TRUE(return_message.has_error_message());
   auto query = return_message.flag_query_message();
   auto errmsg = return_message.error_message();
-  ASSERT_TRUE(errmsg.find("unknown is not found in mockup") != std::string::npos);
+  ASSERT_TRUE(errmsg.find("Failed to find flag unknown") != std::string::npos);
 }
 
 } // namespace aconfigd
diff --git a/aconfigd/aconfigd_util.cpp b/aconfigd/aconfigd_util.cpp
index 72a0b10..607b48b 100644
--- a/aconfigd/aconfigd_util.cpp
+++ b/aconfigd/aconfigd_util.cpp
@@ -16,12 +16,11 @@
 
 #include <memory>
 #include <vector>
+#include <sys/sendfile.h>
+#include <fts.h>
 
 #include <android-base/logging.h>
 #include <android-base/unique_fd.h>
-#include <android-base/file.h>
-#include <sys/sendfile.h>
-#include <fts.h>
 
 #include "aconfigd_util.h"
 
@@ -104,42 +103,5 @@ bool FileExists(const std::string& file) {
   return stat(file.c_str(), &st) == 0 ? true : false;
 }
 
-/// Read persistent aconfig storage records pb file
-Result<aconfig_storage_metadata::storage_files> ReadStorageRecordsPb(
-    const std::string& pb_file) {
-  auto records = aconfig_storage_metadata::storage_files();
-  if (FileExists(pb_file)) {
-    auto content = std::string();
-    if (!ReadFileToString(pb_file, &content)) {
-      return ErrnoError() << "ReadFileToString failed";
-    }
-
-    if (!records.ParseFromString(content)) {
-      return ErrnoError() << "Unable to parse storage records protobuf";
-    }
-  }
-  return records;
-}
-
-/// Write aconfig storage records protobuf to file
-Result<void> WriteStorageRecordsPbToFile(
-    const aconfig_storage_metadata::storage_files& records_pb,
-    const std::string& file_name) {
-  auto content = std::string();
-  if (!records_pb.SerializeToString(&content)) {
-    return ErrnoError() << "Unable to serialize storage records protobuf";
-  }
-
-  if (!WriteStringToFile(content, file_name)) {
-    return ErrnoError() << "WriteStringToFile failed";
-  }
-
-  if (chmod(file_name.c_str(), 0644) == -1) {
-    return ErrnoError() << "chmod() failed";
-  };
-
-  return {};
-}
-
 } // namespace aconfig
 } // namespace android
diff --git a/aconfigd/aconfigd_util.h b/aconfigd/aconfigd_util.h
index e25fd8a..104981c 100644
--- a/aconfigd/aconfigd_util.h
+++ b/aconfigd/aconfigd_util.h
@@ -15,9 +15,10 @@
  */
 
 #include <string>
-#include <android-base/result.h>
 #include <sys/stat.h>
-#include <protos/aconfig_storage_metadata.pb.h>
+
+#include <android-base/result.h>
+#include <android-base/file.h>
 
 namespace android {
   namespace aconfigd {
@@ -26,7 +27,9 @@ namespace android {
   base::Result<void> RemoveFilesInDir(const std::string& dir);
 
   /// Copy file
-  base::Result<void> CopyFile(const std::string& src, const std::string& dst, mode_t mode);
+  base::Result<void> CopyFile(const std::string& src,
+                              const std::string& dst,
+                              mode_t mode);
 
   /// Get a file's timestamp
   base::Result<int> GetFileTimeStamp(const std::string& file);
@@ -34,14 +37,43 @@ namespace android {
   /// Check if file exists
   bool FileExists(const std::string& file);
 
-  /// Read persistent aconfig storage records pb file
-  base::Result<aconfig_storage_metadata::storage_files> ReadStorageRecordsPb(
-      const std::string& pb_file);
+  /// Read protobuf from file
+  template <typename T>
+  base::Result<T> ReadPbFromFile(const std::string& pb_file) {
+    auto pb = T();
+    if (FileExists(pb_file)) {
+      auto content = std::string();
+      if (!base::ReadFileToString(pb_file, &content)) {
+        return base::ErrnoError() << "ReadFileToString() failed";
+      }
+
+      if (!pb.ParseFromString(content)) {
+        return base::ErrnoError() << "Unable to parse to protobuf";
+      }
+    }
+    return pb;
+  }
+
+  /// Write protobuf to file
+  template <typename T>
+  base::Result<void> WritePbToFile(const T& pb,
+                                   const std::string& file_name,
+                                   mode_t mode = 0644) {
+    auto content = std::string();
+    if (!pb.SerializeToString(&content)) {
+      return base::ErrnoError() << "Unable to serialize protobuf to string";
+    }
+
+    if (!base::WriteStringToFile(content, file_name)) {
+      return base::ErrnoError() << "WriteStringToFile() failed";
+    }
+
+    if (chmod(file_name.c_str(), mode) == -1) {
+      return base::ErrnoError() << "chmod() failed";
+    };
 
-  /// Write aconfig storage records protobuf to file
-  base::Result<void> WriteStorageRecordsPbToFile(
-      const aconfig_storage_metadata::storage_files& records_pb,
-      const std::string& file_name);
+    return {};
+  }
 
   }// namespace aconfig
 } // namespace android