7 files changed, 32 insertions, 0 deletions

diff --git a/.clang-format b/.clang-format
index 4a3bd2fc3..6fa717cfb 100644
--- a/.clang-format
+++ b/.clang-format
@@ -36,6 +36,7 @@ AllowShortIfStatementsOnASingleLine: true
 ColumnLimit: 100
 CommentPragmas: NOLINT:.*
 DerivePointerAlignment: false
+IncludeBlocks: Preserve
 IndentWidth: 2
 PointerAlignment: Left
 TabWidth: 2
diff --git a/edify/include/edify/updater_runtime_interface.h b/edify/include/edify/updater_runtime_interface.h
index d3d26da64..bdd6aecc8 100644
--- a/edify/include/edify/updater_runtime_interface.h
+++ b/edify/include/edify/updater_runtime_interface.h
@@ -71,4 +71,7 @@ class UpdaterRuntimeInterface {
   virtual bool MapPartitionOnDeviceMapper(const std::string& partition_name, std::string* path) = 0;
   virtual bool UnmapPartitionOnDeviceMapper(const std::string& partition_name) = 0;
   virtual bool UpdateDynamicPartitions(const std::string_view op_list_value) = 0;
+
+  // On devices supports A/B, add current slot suffix to arg. Otherwise, return |arg| as is.
+  virtual std::string AddSlotSuffix(const std::string_view arg) const = 0;
 };
diff --git a/updater/include/updater/simulator_runtime.h b/updater/include/updater/simulator_runtime.h
index 9f7847b4f..fa878db33 100644
--- a/updater/include/updater/simulator_runtime.h
+++ b/updater/include/updater/simulator_runtime.h
@@ -53,6 +53,7 @@ class SimulatorRuntime : public UpdaterRuntimeInterface {
   bool MapPartitionOnDeviceMapper(const std::string& partition_name, std::string* path) override;
   bool UnmapPartitionOnDeviceMapper(const std::string& partition_name) override;
   bool UpdateDynamicPartitions(const std::string_view op_list_value) override;
+  std::string AddSlotSuffix(const std::string_view arg) const override;
 
  private:
   std::string FindBlockDeviceName(const std::string_view name) const override;
diff --git a/updater/include/updater/updater_runtime.h b/updater/include/updater/updater_runtime.h
index 8fc066f6a..b943dfcf1 100644
--- a/updater/include/updater/updater_runtime.h
+++ b/updater/include/updater/updater_runtime.h
@@ -56,6 +56,7 @@ class UpdaterRuntime : public UpdaterRuntimeInterface {
   bool MapPartitionOnDeviceMapper(const std::string& partition_name, std::string* path) override;
   bool UnmapPartitionOnDeviceMapper(const std::string& partition_name) override;
   bool UpdateDynamicPartitions(const std::string_view op_list_value) override;
+  std::string AddSlotSuffix(const std::string_view arg) const override;
 
  private:
   struct selabel_handle* sehandle_{ nullptr };
diff --git a/updater/install.cpp b/updater/install.cpp
index 7608dc3cd..afa5195d0 100644
--- a/updater/install.cpp
+++ b/updater/install.cpp
@@ -852,6 +852,20 @@ Value* Tune2FsFn(const char* name, State* state, const std::vector<std::unique_p
   return StringValue("t");
 }
 
+Value* AddSlotSuffixFn(const char* name, State* state,
+                       const std::vector<std::unique_ptr<Expr>>& argv) {
+  if (argv.size() != 1) {
+    return ErrorAbort(state, kArgsParsingFailure, "%s() expects 1 arg, got %zu", name, argv.size());
+  }
+  std::vector<std::string> args;
+  if (!ReadArgs(state, argv, &args)) {
+    return ErrorAbort(state, kArgsParsingFailure, "%s() Failed to parse the argument(s)", name);
+  }
+  const std::string& arg = args[0];
+  auto updater_runtime = state->updater->GetRuntime();
+  return StringValue(updater_runtime->AddSlotSuffix(arg));
+}
+
 void RegisterInstallFunctions() {
   RegisterFunction("mount", MountFn);
   RegisterFunction("is_mounted", IsMountedFn);
@@ -885,4 +899,6 @@ void RegisterInstallFunctions() {
 
   RegisterFunction("enable_reboot", EnableRebootFn);
   RegisterFunction("tune2fs", Tune2FsFn);
+
+  RegisterFunction("add_slot_suffix", AddSlotSuffixFn);
 }
diff --git a/updater/simulator_runtime.cpp b/updater/simulator_runtime.cpp
index 3ed7bf337..57dfb32d4 100644
--- a/updater/simulator_runtime.cpp
+++ b/updater/simulator_runtime.cpp
@@ -130,3 +130,8 @@ bool SimulatorRuntime::UpdateDynamicPartitions(const std::string_view op_list_va
   }
   return true;
 }
+
+std::string SimulatorRuntime::AddSlotSuffix(const std::string_view arg) const {
+  LOG(INFO) << "Skip adding slot suffix to " << arg;
+  return std::string(arg);
+}
diff --git a/updater/updater_runtime.cpp b/updater/updater_runtime.cpp
index b1b8863fd..e93830505 100644
--- a/updater/updater_runtime.cpp
+++ b/updater/updater_runtime.cpp
@@ -28,6 +28,7 @@
 #include <android-base/strings.h>
 #include <android-base/unique_fd.h>
 #include <ext4_utils/wipe.h>
+#include <fs_mgr.h>
 #include <selinux/label.h>
 #include <tune2fs.h>
 
@@ -186,3 +187,7 @@ int UpdaterRuntime::Tune2Fs(const std::vector<std::string>& args) const {
   // tune2fs changes the filesystem parameters on an ext2 filesystem; it returns 0 on success.
   return tune2fs_main(tune2fs_args.size() - 1, tune2fs_args.data());
 }
+
+std::string UpdaterRuntime::AddSlotSuffix(const std::string_view arg) const {
+  return std::string(arg) + fs_mgr_get_slot_suffix();
+}