mod_image_for_recovery.sh: preserve whitelisted files on stateful.

We want the recovery image to contain CRX files that will be used to preload
the CRX cache on the stateful partition. This change adds a whitelist
of files on the stateful partition to preserve when creating a recovery
image with --minimize_image (the normal case).

To do this, we have to increase the size of the recovery image's stateful
partition. This CL dynamically discovers the approximate smallest size that
will hold the files to preserve.

We will operate correctly if these CRX files are not on the recovery image,
but more bandwidth is used at first boot/first login. We chose to make
the recovery image larger because:

1) Testers use recovery to setup machines for testing. It will be much better
if they test images as close to customer experience as possible.

2) You are going to pay the bandwidth during recovery one way or
another (getting the image, or after boot). If you recover more than
one machine, then the total bandwidth usage is reduced. This might be
of particular importance to an institution that processes many
machines at once.

3) Recovery is sometimes used in the factory process. In particular,
it's sometimes run end of the factory process if there was something
odd with the machine in question. We want these files to always be
there when leaving factory.

BUG=chromium:288255
TEST=Built images, and inspected stateful.

Change-Id: I39090377ea4b1f68bcc16a127a88a79228680535
Reviewed-on: https://chromium-review.googlesource.com/169061
Reviewed-by: Will Drewry <wad@chromium.org>
Commit-Queue: Don Garrett <dgarrett@chromium.org>
Tested-by: Don Garrett <dgarrett@chromium.org>
diff --git a/mod_image_for_recovery.sh b/mod_image_for_recovery.sh
index ecf862c..cb0badc 100755
--- a/mod_image_for_recovery.sh
+++ b/mod_image_for_recovery.sh
@@ -72,6 +72,9 @@
 . "${BUILD_LIBRARY_DIR}/board_options.sh" || exit 1
 EMERGE_BOARD_CMD="emerge-$BOARD"
 
+# Files to preserve from original stateful, if minimize_image is true.
+# If minimize_image is false, everything is always preserved.
+WHITELIST="vmlinuz_hd.vblock unencrypted/import_extensions"
 
 get_install_vblock() {
   # If it exists, we need to copy the vblock over to stateful
@@ -229,29 +232,81 @@
   return 0
 }
 
+find_sectors_needed() {
+  # Find the minimum disk sectors needed for a file system to hold a list of
+  # files or directories.
+  local base_dir file_list
+
+  base_dir="$1"
+  file_list="$2"
+
+  local sectors_needed=1
+  local size name
+  for name in ${file_list}; do
+    if [ -e "${base_dir}/${name}" ]; then
+      size=$(du -B512 -s "${base_dir}/${name}" | awk '{ print $1 }')
+      sectors_needed=$((${sectors_needed} + ${size}))
+    fi
+  done
+  # Add 5% overhead for the FS, rounded down.
+  sectors_needed=$((${sectors_needed} + (${sectors_needed} / 20)))
+
+  echo "${sectors_needed}"
+}
+
 maybe_resize_stateful() {
   # If we're not minimizing, then just copy and go.
   if [ $FLAGS_minimize_image -eq $FLAGS_FALSE ]; then
     return 0
   fi
 
-  # Rebuild the image with a 1 sector stateful partition
-  local err=0
-  local small_stateful=$(mktemp)
+  local old_stateful_offset old_stateful_mnt sectors_needed
+  local small_stateful new_stateful_mnt
+
+  # Mount the old stateful partition so we can copy selected values
+  # off of it.
+  old_stateful_offset=$(partoffset "$FLAGS_image" 1)
+  old_stateful_mnt=$(mktemp -d)
+
+  sudo mount -o ro,loop,offset=$((old_stateful_offset * 512)) \
+    "$FLAGS_image" $old_stateful_mnt
+
+  # Add 5% overhead for the FS, rounded down.
+  sectors_needed=$(find_sectors_needed "${old_stateful_mnt}" "${WHITELIST}")
+
+  if [ ${FLAGS_statefulfs_sectors} -gt ${sectors_needed} ]; then
+    sectors_needed="${FLAGS_statefulfs_sectors}"
+  fi
+
+  # Rebuild the image with stateful partition sized by sectors_needed.
+  small_stateful=$(mktemp)
   dd if=/dev/zero of="$small_stateful" bs=512 \
-    count=${FLAGS_statefulfs_sectors} 1>&2
+    count="${sectors_needed}" 1>&2
   trap "rm $small_stateful" RETURN
   # Don't bother with ext3 for such a small image.
   /sbin/mkfs.ext2 -F -b 4096 "$small_stateful" 1>&2
 
   # If it exists, we need to copy the vblock over to stateful
   # This is the real vblock and not the recovery vblock.
-  local new_stateful_mnt=$(mktemp -d)
+  new_stateful_mnt=$(mktemp -d)
 
-  set +e
   sudo mount -o loop $small_stateful $new_stateful_mnt
-  sudo cp "$INSTALL_VBLOCK" "$new_stateful_mnt/vmlinuz_hd.vblock"
+
+  # Create the directories that are going to be needed below. With correct
+  # permissions.
+  mkdir --mode=766 "${new_stateful_mnt}/unencrypted"
+
+  # Copy over any files that need to be preserved.
+  for name in ${WHITELIST}; do
+    if [ -e "${old_stateful_mnt}/${name}" ]; then
+      sudo cp -a "${old_stateful_mnt}/${name}" "${new_stateful_mnt}/${name}"
+    fi
+  done
+
+  # Cleanup everything.
+  safe_umount "$old_stateful_mnt"
   safe_umount "$new_stateful_mnt"
+  rmdir "$old_stateful_mnt"
   rmdir "$new_stateful_mnt"
   switch_to_strict_mode
 
@@ -259,9 +314,9 @@
   # TODO(wad) Make the developer script case create a custom GPT with
   # just the kernel image and stateful.
   update_partition_table "${FLAGS_image}" "$small_stateful" \
-                         ${FLAGS_statefulfs_sectors} \
+                         "${sectors_needed}" \
                          "${RECOVERY_IMAGE}" 1>&2
-  return $err
+  return 0
 }
 
 cleanup() {