cros_generate_update_payload: support for user provided kernel/root paths

This is needed in order for a caller to be able to utilize the extracted
partitions after processing has completed, e.g. for verifying the
generated payloads.

BUG=chromium:243561
TEST=Partitions extracted to given paths; files preserved.

Change-Id: I6070596d5b7d6a16bbf8e1b37f7113a2c366bd8b
Reviewed-on: https://chromium-review.googlesource.com/181221
Tested-by: Gilad Arnold <garnold@chromium.org>
Reviewed-by: Don Garrett <dgarrett@chromium.org>
Commit-Queue: Gilad Arnold <garnold@chromium.org>
diff --git a/host/cros_generate_update_payload b/host/cros_generate_update_payload
index 9040731..54b829d 100755
--- a/host/cros_generate_update_payload
+++ b/host/cros_generate_update_payload
@@ -64,10 +64,18 @@
     [ -d "$STATE_MNT" ] && rmdir "$STATE_MNT"
     STATE_MNT=""
   fi
-  rm -f "$SRC_KERNEL"
-  rm -f "$SRC_ROOT"
-  rm -f "$DST_KERNEL"
-  rm -f "$DST_ROOT"
+  if [ -z "$FLAGS_src_kern_path" ]; then
+    rm -f "$SRC_KERNEL"
+  fi
+  if [ -z "$FLAGS_src_root_path" ]; then
+    rm -f "$SRC_ROOT"
+  fi
+  if [ -z "$FLAGS_kern_path" ]; then
+    rm -f "$DST_KERNEL"
+  fi
+  if [ -z "$FLAGS_root_path" ]; then
+    rm -f "$DST_ROOT"
+  fi
   [ -n "$1" ] || exit 1
 }
 
@@ -77,8 +85,8 @@
   local temp_file="$3"
   if [ -z "$temp_file" ]; then
     temp_file=$(mktemp /tmp/cros_generate_update_payload.XXXXXX)
-    echo "$temp_file"
   fi
+  echo "$temp_file"
 
   # Keep `local` decl split from assignment so return code is checked.
   local offset length
@@ -128,11 +136,11 @@
     die "missing root output filename"
   fi
 
-  extract_partition_to_temp_file "$bin_file" 2 "$kern_out"
+  extract_partition_to_temp_file "$bin_file" 2 "$kern_out" > /dev/null
   if [ "$FLAGS_patch_kernel" -eq "$FLAGS_TRUE" ]; then
     patch_kernel "$bin_file" "$kern_out"
   fi
-  extract_partition_to_temp_file "$bin_file" 3 "$root_out"
+  extract_partition_to_temp_file "$bin_file" 3 "$root_out" > /dev/null
 }
 
 DEFINE_string image "" "The image that should be sent to clients."
@@ -165,6 +173,14 @@
 DEFINE_string build_channel "" "Channel of the build of the target image."
 DEFINE_string build_version "" "Version of the build of the target image."
 
+# Specifying any of the following will cause it to not be cleaned up upon exit.
+DEFINE_string kern_path "" "File path for extracting the kernel partition."
+DEFINE_string root_path "" "File path for extracting the rootfs partition."
+DEFINE_string src_kern_path "" \
+  "File path for extracting the source kernel partition."
+DEFINE_string src_root_path "" \
+  "File path for extracting the source rootfs partition."
+
 
 # Parse command line
 FLAGS "$@" || exit 1
@@ -186,10 +202,14 @@
 
 if [ "$FLAGS_extract" -eq "$FLAGS_TRUE" ]; then
   if [ -n "$FLAGS_src_image" ]; then
-    extract_kern_root "$FLAGS_src_image" old_kern.dat old_root.dat
+    SRC_KERN_PATH="${FLAGS_src_kern_path:-old_kern.dat}"
+    SRC_ROOT_PATH="${FLAGS_src_root_path:-old_root.dat}"
+    extract_kern_root "$FLAGS_src_image" "$SRC_KERN_PATH" "$SRC_ROOT_PATH"
   fi
   if [ -n "$FLAGS_image" ]; then
-    extract_kern_root "$FLAGS_image" new_kern.dat new_root.dat
+    KERN_PATH="${FLAGS_kern_path:-new_kern.dat}"
+    ROOT_PATH="${FLAGS_root_path:-new_root.dat}"
+    extract_kern_root "$FLAGS_image" "$KERN_PATH" "$ROOT_PATH"
   fi
   echo Done extracting kernel/root
   exit 0
@@ -210,7 +230,8 @@
 trap cleanup INT TERM EXIT
 if [ "$DELTA" -eq "$FLAGS_TRUE" ]; then
   if [ "$FLAGS_full_kernel" -eq "$FLAGS_FALSE" ]; then
-    SRC_KERNEL=$(extract_partition_to_temp_file "$FLAGS_src_image" 2)
+    SRC_KERNEL=$(extract_partition_to_temp_file "$FLAGS_src_image" 2 \
+                 "$FLAGS_src_kern_path")
     if [ "$FLAGS_patch_kernel" -eq "$FLAGS_TRUE" ]; then
       patch_kernel "$FLAGS_src_image" "$SRC_KERNEL"
     fi
@@ -219,17 +240,20 @@
   else
     echo "Generating a full kernel update."
   fi
-  SRC_ROOT=$(extract_partition_to_temp_file "$FLAGS_src_image" 3)
+  SRC_ROOT=$(extract_partition_to_temp_file "$FLAGS_src_image" 3 \
+             "$FLAGS_src_root_path")
 
   echo md5sum of src root:
   md5sum "$SRC_ROOT"
 fi
 
-DST_KERNEL=$(extract_partition_to_temp_file "$FLAGS_image" 2)
+DST_KERNEL=$(extract_partition_to_temp_file "$FLAGS_image" 2 \
+             "$FLAGS_kern_path")
 if [ "$FLAGS_patch_kernel" -eq "$FLAGS_TRUE" ]; then
   patch_kernel "$FLAGS_image" "$DST_KERNEL"
 fi
-DST_ROOT=$(extract_partition_to_temp_file "$FLAGS_image" 3)
+DST_ROOT=$(extract_partition_to_temp_file "$FLAGS_image" 3 \
+           "$FLAGS_root_path")
 
 DST_MNT=$(mktemp -d /tmp/src_root.XXXXXX)
 sudo mount -o loop,ro "$DST_ROOT" "$DST_MNT"