cros_generate_update_payload: respect TMPDIR and --work_dir

Makes it far more usable for cases where the hard-coded /tmp won't do.
It is named --work_dir for consistency with other tools (e.g. paygen).

BUG=chromium:253622
TEST=Works with both TMPDIR and --work_dir.

Change-Id: Iac9141018cd626a706570858652a468dff427901
Reviewed-on: https://chromium-review.googlesource.com/181222
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 54b829d..c2f2593 100755
--- a/host/cros_generate_update_payload
+++ b/host/cros_generate_update_payload
@@ -84,7 +84,8 @@
   local partition="$2"
   local temp_file="$3"
   if [ -z "$temp_file" ]; then
-    temp_file=$(mktemp /tmp/cros_generate_update_payload.XXXXXX)
+    temp_file=$(mktemp --tmpdir="${FLAGS_work_dir}" \
+                cros_generate_update_payload.XXXXXX)
   fi
   echo "$temp_file"
 
@@ -116,7 +117,7 @@
   local offset
   offset=$(partoffset "${IMAGE}" 1)
   : $(( offset *= 512 ))
-  STATE_MNT=$(mktemp -d /tmp/state.XXXXXX)
+  STATE_MNT=$(mktemp -d --tmpdir="${FLAGS_work_dir}" state.XXXXXX)
   sudo mount -o ro,loop,offset=$offset "$IMAGE" "$STATE_MNT"
   dd if="$STATE_MNT"/vmlinuz_hd.vblock of="$KERN_FILE" conv=notrunc 2>/dev/null
   sudo umount "$STATE_MNT"
@@ -181,6 +182,8 @@
 DEFINE_string src_root_path "" \
   "File path for extracting the source rootfs partition."
 
+DEFINE_string work_dir "" "Where to dump temporary files."
+
 
 # Parse command line
 FLAGS "$@" || exit 1
@@ -255,7 +258,7 @@
 DST_ROOT=$(extract_partition_to_temp_file "$FLAGS_image" 3 \
            "$FLAGS_root_path")
 
-DST_MNT=$(mktemp -d /tmp/src_root.XXXXXX)
+DST_MNT=$(mktemp -d --tmpdir="${FLAGS_work_dir}" src_root.XXXXXX)
 sudo mount -o loop,ro "$DST_ROOT" "$DST_MNT"
 
 source "$DST_MNT"/usr/sbin/write_gpt.sh || warn "Could not read write_gpt.sh"
@@ -269,7 +272,7 @@
 fi
 
 if [ "$DELTA" -eq "$FLAGS_TRUE" ]; then
-  SRC_MNT=$(mktemp -d /tmp/src_root.XXXXXX)
+  SRC_MNT=$(mktemp -d --tmpdir="${FLAGS_work_dir}" src_root.XXXXXX)
   sudo mount -o loop,ro "$SRC_ROOT" "$SRC_MNT"
 
   sudo LD_LIBRARY_PATH=${LD_LIBRARY_PATH}  PATH=${PATH} "$GENERATOR" \