From e1e306566ac972e6ebd4093ca43948be11bc152c Mon Sep 17 00:00:00 2001
From: Tan Gao <tgao@chromium.org>
Date: Fri, 8 Oct 2010 15:22:51 -0700
Subject: [PATCH] Issue 6759: reduce size of stateful partition for
 non-developer recovery image

Change-Id: I636f3156e549508230515174edffaff650c9d46e

BUG=chromium-os:6759
TEST=(1)manually built a non-developer recovery image and verified recovery install
completes on an agz device and the device is able to boot from HD upon reboot
(2) repeat for a developer recovery image and verified dev image from dev_payload is installed onto target and the target is able to boot from HD

Review URL: http://codereview.chromium.org/3391018
---
 mod_image_for_dev_recovery.sh | 114 +++++----------------------------
 mod_image_for_recovery.sh     | 117 +++++++++++++++++-----------------
 resize_stateful_partition.sh  |  96 ++++++++++++++++++++++++++++
 3 files changed, 169 insertions(+), 158 deletions(-)
 create mode 100755 resize_stateful_partition.sh

diff --git a/mod_image_for_dev_recovery.sh b/mod_image_for_dev_recovery.sh
index e94f687b6d..9c68b1e323 100755
--- a/mod_image_for_dev_recovery.sh
+++ b/mod_image_for_dev_recovery.sh
@@ -7,11 +7,12 @@
 # Script to create a Chrome OS dev recovery image using a dev install shim
 
 # Source constants and utility functions
-. "$(dirname "$0")/common.sh"
-. "$(dirname "$0")/chromeos-common.sh"
+. "$(dirname "$0")/resize_stateful_partition.sh"
 
 get_default_board
-locate_gpt
+
+# Constants
+TEMP_IMG=$(mktemp "/tmp/temp_img.XXXXXX")
 
 DEFINE_string board "$DEFAULT_BOARD" "Board for which the image was built \
 Default: ${DEFAULT_BOARD}"
@@ -73,94 +74,6 @@ fi
 INSTALL_SHIM_DIR="$(dirname "$FLAGS_dev_install_shim")"
 DEV_RECOVERY_IMAGE="dev_recovery_image.bin"
 
-umount_from_loop_dev() {
-  local mnt_pt=$1
-  mount | grep -q " on ${mnt_pt}" && sudo umount ${mnt_pt}
-  rmdir ${mnt_pt}
-}
-
-cleanup_loop_dev() {
-  sudo losetup -d ${1} || /bin/true
-}
-
-get_loop_dev() {
-  local loop_dev=$(sudo losetup -f)
-  if [ -z "${loop_dev}" ]; then
-    die "No free loop device. Free up a loop device or reboot. Exiting."
-  fi
-  echo ${loop_dev}
-}
-
-# Resize stateful partition of install shim to hold payload content
-# Due to this resize, we need to create a new partition table and a new image.
-# (see update_partition_table() for details)
-resize_partition() {
-  local source_part=$1  # source partition
-  local add_num_sectors=$2  # number of 512-byte sectors to be added
-
-  local source_sectors=$(roundup $(numsectors ${source_part}))
-  info "source partition has ${source_sectors} 512-byte sectors."
-  local resized_sectors=$(roundup $(expr $source_sectors + $add_num_sectors))
-  info "resized partition has ${resized_sectors} 512-byte sectors."
-
-  local loop_dev=$(get_loop_dev)
-  trap "cleanup_loop_dev ${loop_dev}" EXIT
-
-  # Extend the source file size to the new size.
-  dd if=/dev/zero of="${source_part}" bs=1 count=1 \
-      seek=$((512 * ${resized_sectors} - 1)) &>/dev/null
-
-  # Resize the partition.
-  sudo losetup "${loop_dev}" "${source_part}"
-  sudo e2fsck -fp "${loop_dev}" &> /dev/null
-  sudo resize2fs "${loop_dev}" &> /dev/null
-  # trap handler will clean up the loop device
-  echo "${resized_sectors}"
-}
-
-# Update partition table with resized stateful partition and create the final
-# dev recovery image
-update_partition_table() {
-  local temp_state=$1       # stateful partition image
-  local resized_sectors=$2  # number of sectors in resized stateful partition
-  local temp_img=$(mktemp "/tmp/temp_img.XXXXXX")
-
-  local kernel_offset=$(partoffset ${FLAGS_dev_install_shim} 2)
-  local kernel_count=$(partsize ${FLAGS_dev_install_shim} 2)
-  local rootfs_offset=$(partoffset ${FLAGS_dev_install_shim} 3)
-  local rootfs_count=$(partsize ${FLAGS_dev_install_shim} 3)
-  local oem_offset=$(partoffset ${FLAGS_dev_install_shim} 8)
-  local oem_count=$(partsize ${FLAGS_dev_install_shim} 8)
-  local esp_offset=$(partoffset ${FLAGS_dev_install_shim} 12)
-  local esp_count=$(partsize ${FLAGS_dev_install_shim} 12)
-
-  local temp_pmbr=$(mktemp "/tmp/pmbr.XXXXXX")
-  dd if="${FLAGS_dev_install_shim}" of="${temp_pmbr}" bs=512 count=1 &>/dev/null
-
-  # Set up a new partition table
-  install_gpt "${temp_img}" "${rootfs_count}" "${resized_sectors}" \
-    "${temp_pmbr}" "${esp_count}" false $(roundup ${rootfs_count}) &>/dev/null
-
-  rm -rf "${temp_pmbr}"
-
-  # Copy into the partition parts of the file
-  dd if="${FLAGS_dev_install_shim}" of="${temp_img}" conv=notrunc bs=512 \
-    seek="${START_ROOTFS_A}" skip=${rootfs_offset} count=${rootfs_count} \
-    &>/dev/null
-  dd if="${temp_state}" of="${temp_img}" conv=notrunc bs=512 \
-    seek="${START_STATEFUL}" &>/dev/null
-  # Copy the full kernel (i.e. with vboot sections)
-  dd if="${FLAGS_dev_install_shim}" of="${temp_img}" conv=notrunc bs=512 \
-    seek="${START_KERN_A}" skip=${kernel_offset} count=${kernel_count} \
-    &>/dev/null
-  dd if="${FLAGS_dev_install_shim}" of="${temp_img}" conv=notrunc bs=512 \
-    seek="${START_OEM}" skip=${oem_offset} count=${oem_count} &>/dev/null
-  dd if="${FLAGS_dev_install_shim}" of="${temp_img}" conv=notrunc bs=512 \
-    seek="${START_ESP}" skip=${esp_offset} count=${esp_count} &>/dev/null
-
-  echo ${temp_img}
-}
-
 # Creates a dev recovery image using an existing dev install shim
 # If successful, content of --payload_dir is copied to a directory named
 # "dev_payload" under the root of stateful partition.
@@ -171,14 +84,15 @@ create_dev_recovery_image() {
   dd if="${FLAGS_dev_install_shim}" of="${temp_state}" conv=notrunc bs=512 \
     skip=${stateful_offset} count=${stateful_count} &>/dev/null
 
-  local resized_sectors=$(resize_partition $temp_state $PAYLOAD_DIR_SIZE)
+  local resized_sectors=$(enlarge_partition_image $temp_state $PAYLOAD_DIR_SIZE)
 
   # Mount resized stateful FS and copy payload content to its root directory
   local temp_mnt=$(mktemp -d "/tmp/temp_mnt.XXXXXX")
   local loop_dev=$(get_loop_dev)
-  trap "umount_from_loop_dev ${temp_mnt} && rm -f \"${temp_state}\"" EXIT
+  trap "cleanup_loop_dev ${loop_dev}" EXIT
   mkdir -p "${temp_mnt}"
   sudo mount -o loop=${loop_dev} "${temp_state}" "${temp_mnt}"
+  trap "umount_from_loop_dev ${temp_mnt} && rm -f \"${temp_state}\"" EXIT
   sudo cp -R "${FLAGS_payload_dir}" "${temp_mnt}"
   sudo mv "${temp_mnt}/$(basename ${FLAGS_payload_dir})" \
     "${temp_mnt}/dev_payload"
@@ -187,17 +101,21 @@ create_dev_recovery_image() {
   sudo touch "${temp_mnt}/.dev_recovery"
 
   # TODO(tgao): handle install script (for default and custom cases)
-  local temp_img=$(update_partition_table $temp_state $resized_sectors)
+  (update_partition_table $FLAGS_dev_install_shim $temp_state \
+       $resized_sectors $TEMP_IMG)
 
   # trap handler will clean up loop device and temp mount point
-  echo ${temp_img}
 }
 
 # Main
 DST_PATH="${INSTALL_SHIM_DIR}/${DEV_RECOVERY_IMAGE}"
 info "Attempting to create dev recovery image using dev install shim \
 ${FLAGS_dev_install_shim}"
-TEMP_IMG=$(create_dev_recovery_image)
+(create_dev_recovery_image)
 
-mv -f $TEMP_IMG $DST_PATH
-info "Dev recovery image created at ${DST_PATH}"
+if [ -n ${TEMP_IMG} ] && [ -f ${TEMP_IMG} ]; then
+  mv -f $TEMP_IMG $DST_PATH
+  info "Dev recovery image created at ${DST_PATH}"
+else
+  info "Failed to create developer recovery image"
+fi
diff --git a/mod_image_for_recovery.sh b/mod_image_for_recovery.sh
index 6acf46daa8..7d59095753 100755
--- a/mod_image_for_recovery.sh
+++ b/mod_image_for_recovery.sh
@@ -6,16 +6,15 @@
 
 # Script to modify a pristine/dev Chrome OS image to be used for recovery
 
-. "$(dirname "$0")/common.sh"
-
-# Load functions and constants for chromeos-install
-. "$(dirname "$0")/chromeos-common.sh"
+. "$(dirname "$0")/resize_stateful_partition.sh"
 
 # Script must be run inside the chroot.
 restart_in_chroot_if_needed $*
 
 get_default_board
 
+# Constants
+TEMP_IMG=$(mktemp "/tmp/temp_img.XXXXXX")
 RECOVERY_IMAGE="recovery_image.bin"
 
 DEFINE_string board "$DEFAULT_BOARD" "Board for which the image was built"
@@ -33,7 +32,7 @@ if [ -z $FLAGS_image ] && [ -z $FLAGS_board ] ; then
   die "mod_image_for_recovery failed.  No board set and no image set"
 fi
 
-# We have a board name but no image set.  Use image at default location
+# We have a board name but no image set. Use image at default location
 if [ -z $FLAGS_image ] ; then
   IMAGES_DIR="${DEFAULT_BUILD_ROOT}/images/${FLAGS_board}"
   FILENAME="chromiumos_image.bin"
@@ -49,75 +48,73 @@ if [ ! -f $FLAGS_image ] ; then
   exit 1
 fi
 
+set -u
 set -e
 
 # Constants
 IMAGE_DIR="$(dirname "$FLAGS_image")"
-IMAGE_NAME="$(basename "$FLAGS_image")"
 
-# loop device utility methods mostly duplicated from
-#   src/platform/installer/chromeos-install
-# TODO(tgao): minimize duplication by refactoring these methods into a separate
-# library script which both scripts can reference
+# Creates a dev recovery image using an existing dev install shim
+# If successful, content of --payload_dir is copied to a directory named
+# "dev_payload" under the root of stateful partition.
+create_recovery_image() {
+  local src_img=$1  # base image
+  local src_state=$(mktemp "/tmp/src_state.XXXXXX")
+  local stateful_offset=$(partoffset ${src_img} 1)
+  local stateful_count=$(partsize ${src_img} 1)
 
-# Set up loop device for an image file at specified offset
-loop_offset_setup() {
-  local filename=$1
-  local offset=$2  # 512-byte sectors
+  dd if="${src_img}" of="${src_state}" conv=notrunc bs=512 \
+    skip=${stateful_offset} count=${stateful_count}
 
-  LOOP_DEV=$(sudo losetup -f)
-  if [ -z "$LOOP_DEV" ]
-  then
-    echo "No free loop device. Free up a loop device or reboot. Exiting."
-    exit 1
-  fi
+  # Mount original stateful partition to figure out its actual size
+  local src_loop_dev=$(get_loop_dev)
+  trap "cleanup_loop_dev ${src_loop_dev}" EXIT
 
-  sudo losetup -o $(($offset * 512)) ${LOOP_DEV} ${filename}
-}
+  # Setup loop dev
+  sudo losetup $src_loop_dev $src_state
+  local block_size=$(sudo /sbin/dumpe2fs $src_loop_dev | grep "Block size:" \
+                     | tr -d ' ' | cut -f2 -d:)
+  echo "block_size = $block_size"
+  local min_size=$(sudo /sbin/resize2fs -P $src_loop_dev | tr -d ' ' \
+                   | cut -f2 -d:)
+  echo "min_size = $min_size $block_size blocks"
 
-loop_offset_cleanup() {
-  sudo losetup -d ${LOOP_DEV} || /bin/true
-}
+  # Add 20%, convert to 512-byte sectors and round up to 2Mb boundary
+  local min_sectors=$(roundup $(((min_size * block_size * 120) / (512 * 100))))
+  echo "min_sectors = ${min_sectors} 512-byte blocks"
+  sudo e2fsck -fp "${src_loop_dev}"
+  # Resize using 512-byte sectors
+  sudo /sbin/resize2fs $src_loop_dev ${min_sectors}s
 
-mount_on_loop_dev() {
-  TMPMNT=$(mktemp -d)
-  sudo mount ${LOOP_DEV} ${TMPMNT}
-}
-
-# Unmount loop-mounted device.
-umount_from_loop_dev() {
-  mount | grep -q " on ${TMPMNT} " && sudo umount ${TMPMNT}
-}
-
-# Undo both mount and loop.
-my_cleanup() {
-  umount_from_loop_dev
-  loop_offset_cleanup
-}
-
-# Modifies an existing image for recovery use
-update_recovery_packages() {
-  local image_name=$1
-
-  echo "Modifying image ${image_name} for recovery use"
-
-  locate_gpt  # set $GPT env var
-  loop_offset_setup ${image_name} $(partoffset "${image_name}" 1)
-  trap loop_offset_cleanup EXIT
-  mount_on_loop_dev "readwrite"
-  trap my_cleanup EXIT
-  sudo touch ${TMPMNT}/.recovery
-  umount_from_loop_dev
-  trap loop_offset_cleanup EXIT
-  loop_offset_cleanup
+  # Delete the loop
   trap - EXIT
+  cleanup_loop_dev ${src_loop_dev}
+
+  # Truncate the image at the new size
+  dd if=/dev/zero of=$src_state bs=512 seek=$min_sectors count=0
+
+  # Mount and touch .recovery  # Soon not to be needed :/
+  local new_mnt=$(mktemp -d "/tmp/src_mnt.XXXXXX")
+  mkdir -p "${new_mnt}"
+  local new_loop_dev=$(get_loop_dev)
+  trap "cleanup_loop_dev ${new_loop_dev} && rmdir ${new_mnt} && \
+        rm -f ${src_state}" EXIT
+  sudo mount -o loop=${new_loop_dev} "${src_state}" "${new_mnt}"
+  trap "umount_from_loop_dev ${new_mnt} && rm -f ${src_state}" EXIT
+  sudo touch "${new_mnt}/.recovery"
+
+  (update_partition_table $src_img $src_state $min_sectors $TEMP_IMG)
+  # trap handler will handle unmount and clean up of loop device and temp files
 }
 
 # Main
-
 DST_PATH="${IMAGE_DIR}/${FLAGS_output}"
 echo "Making a copy of original image ${FLAGS_image}"
-cp $FLAGS_image $DST_PATH
-update_recovery_packages $DST_PATH
+(create_recovery_image $FLAGS_image)
 
-echo "Recovery image created at ${DST_PATH}"
+if [ -n ${TEMP_IMG} ] && [ -f ${TEMP_IMG} ]; then
+  mv -f $TEMP_IMG $DST_PATH
+  echo "Recovery image created at ${DST_PATH}"
+else
+  echo "Failed to create recovery image"
+fi
diff --git a/resize_stateful_partition.sh b/resize_stateful_partition.sh
new file mode 100755
index 0000000000..d83e0b8f42
--- /dev/null
+++ b/resize_stateful_partition.sh
@@ -0,0 +1,96 @@
+#!/bin/bash
+
+# Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+# Utility methods used to resize a stateful partition and update the GPT table
+
+# Source constants and utility functions
+. "$(dirname "$0")/common.sh"
+. "$(dirname "$0")/chromeos-common.sh"
+
+locate_gpt
+
+umount_from_loop_dev() {
+  local mnt_pt=$1
+  mount | grep -q " on ${mnt_pt}" && sudo umount -d ${mnt_pt}
+  rmdir ${mnt_pt}
+}
+
+cleanup_loop_dev() {
+  sudo losetup -d ${1} || /bin/true
+}
+
+get_loop_dev() {
+  local loop_dev=$(sudo losetup -f)
+  if [ -z "${loop_dev}" ]; then
+    die "No free loop device. Free up a loop device or reboot. Exiting."
+  fi
+  echo ${loop_dev}
+}
+
+# Resize stateful partition of install shim to hold payload content
+# Due to this resize, we need to create a new partition table and a new image.
+# (see update_partition_table() for details)
+enlarge_partition_image() {
+  local source_part=$1  # source partition
+  local add_num_sectors=$2  # number of 512-byte sectors to be added
+
+  local source_sectors=$(roundup $(numsectors ${source_part}))
+  info "source partition has ${source_sectors} 512-byte sectors."
+  local resized_sectors=$(roundup $(expr $source_sectors + $add_num_sectors))
+  info "resized partition has ${resized_sectors} 512-byte sectors."
+
+  local loop_dev=$(get_loop_dev)
+  trap "cleanup_loop_dev ${loop_dev}" EXIT
+
+  # Extend the source file size to the new size.
+  dd if=/dev/zero of="${source_part}" bs=1 count=1 \
+      seek=$((512 * ${resized_sectors} - 1)) &>/dev/null
+
+  # Resize the partition.
+  sudo losetup "${loop_dev}" "${source_part}"
+  sudo e2fsck -fp "${loop_dev}" &> /dev/null
+  sudo resize2fs "${loop_dev}" &> /dev/null
+  # trap handler will clean up the loop device
+  echo "${resized_sectors}"
+}
+
+# Update partition table with resized stateful partition
+update_partition_table() {
+  local src_img=$1          # source image
+  local temp_state=$2       # stateful partition image
+  local resized_sectors=$3  # number of sectors in resized stateful partition
+  local temp_img=$4
+
+  local kernel_offset=$(partoffset ${src_img} 2)
+  local kernel_count=$(partsize ${src_img} 2)
+  local rootfs_offset=$(partoffset ${src_img} 3)
+  local rootfs_count=$(partsize ${src_img} 3)
+  local oem_offset=$(partoffset ${src_img} 8)
+  local oem_count=$(partsize ${src_img} 8)
+  local esp_offset=$(partoffset ${src_img} 12)
+  local esp_count=$(partsize ${src_img} 12)
+
+  local temp_pmbr=$(mktemp "/tmp/pmbr.XXXXXX")
+  dd if="${src_img}" of="${temp_pmbr}" bs=512 count=1 &>/dev/null
+
+  trap "rm -rf \"${temp_pmbr}\"" EXIT
+  # Set up a new partition table
+  install_gpt "${temp_img}" "${rootfs_count}" "${resized_sectors}" \
+    "${temp_pmbr}" "${esp_count}" false $(roundup ${rootfs_count})
+
+  # Copy into the partition parts of the file
+  dd if="${src_img}" of="${temp_img}" conv=notrunc bs=512 \
+    seek="${START_ROOTFS_A}" skip=${rootfs_offset} count=${rootfs_count}
+  dd if="${temp_state}" of="${temp_img}" conv=notrunc bs=512 \
+    seek="${START_STATEFUL}"
+  # Copy the full kernel (i.e. with vboot sections)
+  dd if="${src_img}" of="${temp_img}" conv=notrunc bs=512 \
+    seek="${START_KERN_A}" skip=${kernel_offset} count=${kernel_count}
+  dd if="${src_img}" of="${temp_img}" conv=notrunc bs=512 \
+    seek="${START_OEM}" skip=${oem_offset} count=${oem_count}
+  dd if="${src_img}" of="${temp_img}" conv=notrunc bs=512 \
+    seek="${START_ESP}" skip=${esp_offset} count=${esp_count}
+}