factory_installer: Patch config of custom process

Patch the config from cros_payload for factory shim to use.

BUG=b:329025415
TEST=use image_tool to set config, and execute
"cros_payload_patch_custom_process" on shim

Cq-Depend: chromium:5392635
Change-Id: I65aaff6bbde0146c090b9963fbed5206d108c839
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/factory_installer/+/5392636
Commit-Queue: Chih-Yao Chuang <jasonchuang@google.com>
Tested-by: Chih-Yao Chuang <jasonchuang@google.com>
Reviewed-by: Yu-An Wang <wyuang@google.com>
diff --git a/factory_cros_payload.sh b/factory_cros_payload.sh
index fdfe6e4..4f7335d 100644
--- a/factory_cros_payload.sh
+++ b/factory_cros_payload.sh
@@ -85,6 +85,44 @@
   rm "${temp_config_path}"
 }
 
+# Patch custom process in reset shim using toolkit_config payload.
+cros_payload_patch_custom_process() {
+  local real_usb_dev
+  real_usb_dev="$(findLSBValue REAL_USB_DEV)"
+  [ -n "${real_usb_dev}" ] || return 0
+  local stateful_dev="${real_usb_dev%%[0-9]*}1"
+  local temp_config_path
+  temp_config_path="$(mktemp)"
+  local mount_point
+  mount_point="$(mktemp -d)"
+  mount -o ro "${stateful_dev}" "${mount_point}"
+
+  echo 'Patching custom process config...'
+
+  local json_path
+  json_path="$(cros_payload_metadata "${mount_point}")"
+  local json_url="${mount_point}/${json_path}"
+  if [ -n "${json_url}" ]; then
+    # Get toolkit config in cros_payload.
+    cros_payload install "${json_url}" "${temp_config_path}" toolkit_config ||
+      true
+    # Get custom process config from toolkit config.
+    local config=""
+    if [ -s "${temp_config_path}" ]; then
+      config="$(
+        jq -s '.[].custom_reset_process | select(. != null)' "${temp_config_path}")"
+    fi
+    if [ -n "${config}" ]; then
+      local config_path="/mnt/stateful_partition/dev_image/etc/custom-process.json"
+      echo "${config}" >"${config_path}"
+    fi
+  fi
+
+  umount "${stateful_dev}"
+  rmdir "${mount_point}"
+  rm "${temp_config_path}"
+}
+
 # Install dev_image/etc/description from description payload.
 cros_payload_install_description() {
   local real_usb_dev="$(findLSBValue REAL_USB_DEV)"
diff --git a/factory_shim_service.sh b/factory_shim_service.sh
index eed0e9a..fa32d85 100644
--- a/factory_shim_service.sh
+++ b/factory_shim_service.sh
@@ -57,6 +57,7 @@
   # Patch lsb-factory and cutoff config.
   cros_payload_patch_lsb_factory
   cros_payload_patch_cutoff_config
+  cros_payload_patch_custom_process
 
   cros_payload_install_description
 
diff --git a/rust/src/bin/factory_installer.rs b/rust/src/bin/factory_installer.rs
index 33133d3..aa04d31 100644
--- a/rust/src/bin/factory_installer.rs
+++ b/rust/src/bin/factory_installer.rs
@@ -16,6 +16,7 @@
     // TODO(b/320872246): No need to patch after the issue is solved.
     script_wrapper::call_script_wrapper(vec!["cros_payload_patch_cutoff_config"], &mut context)?;
     script_wrapper::call_script_wrapper(vec!["cros_payload_patch_lsb_factory"], &mut context)?;
+    script_wrapper::call_script_wrapper(vec!["cros_payload_patch_custom_process"], &mut context)?;
     match args.action {
         Action::FAI {
             output_path,