Switch to GPT-format disk images.

This changes the disk image for both USB keys and hard-drive installs to use
the EFI GUID Partition Tables. This is a prequisite for booting with an EFI
BIOS. This change does not use the EFI Boot protocol; it still boots using
Legacy BIOS. The PMBR contains syslinux' gptmbr.bin, which searches the GPT
for a specified partition's GUID to boot.

I've tested it on my EeePC. The USB image works, chromeos-install works, and
the reimaged hard drive works. I have not yet tested the memento_updater.sh
script, but I wanted to start the review without waiting until it's all perfect.
I will also be refactoring build_gpt.sh and chromeos-install to share common
logic.

It's almost certain that all existing dogfood machines will need to be
reimaged from a USB key. We could probably figure out a way to upgrade
automatically, but not quickly or without risk.

In addition to the GPT formatting, the build_image script has changed to
emit a single usb.bin file. This can be copied directly onto a USB key and
booted. Installation of dev tools needs to happen with build_image, not
image_to_usb. I have not yet looked at the other image_to_* scripts.

Review URL: http://codereview.chromium.org/1100001

build_gpt.sh

@@ -0,0 +1,146 @@
+#!/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.
+#
+
+# Load common constants.  This should be the first executable line.
+# The path to common.sh should be relative to your script's location.
+. "$(dirname "$0")/common.sh"
+
+# Load functions and constants for chromeos-install
+. "$(dirname "$0")/chromeos-common.sh"
+
+# Script must be run inside the chroot.
+assert_inside_chroot
+
+get_default_board
+
+# Flags.
+DEFINE_string arch "" \
+  "The target architecture (\"arm\" or \"x86\")."
+DEFINE_string board "$DEFAULT_BOARD" \
+  "The board to build an image for."
+DEFINE_string board_root "" \
+  "The build directory, needed to find tools for ARM."
+
+# Usage.
+FLAGS_HELP=$(cat <<EOF
+ 
+Usage: $(basename $0) [flags] IMAGEDIR OUTDEV
+
+This takes the image components in IMAGEDIR and creates a bootable,
+GPT-formatted image in OUTDEV. OUTDEV can be a file or block device.
+ 
+EOF
+)
+
+# Parse command line.
+FLAGS "$@" || exit 1
+eval set -- "${FLAGS_ARGV}"
+
+if [[ -z "$FLAGS_board" ]] ; then
+  error "--board is required."
+  exit 1
+fi
+
+if [[ -z "$1" || -z "$2" ]] ; then
+  flags_help
+  exit 1
+fi
+IMAGEDIR="$1"
+OUTDEV="$2"
+
+if [[ -n "$FLAGS_arch" ]]; then
+  ARCH=${FLAGS_arch}
+else
+  # Figure out ARCH from the given toolchain.
+  # TODO: Move to common.sh as a function after scripts are switched over.
+  TC_ARCH=$(echo "$CHOST" | awk -F'-' '{ print $1 }')
+  case "$TC_ARCH" in
+    arm*)
+      ARCH="arm"
+      ;;
+    *86)
+      ARCH="x86"
+      ;;
+    *)
+      error "Unable to determine ARCH from toolchain: $CHOST"
+      exit 1
+  esac
+fi
+
+if [[ -z "$FLAGS_board_root" ]]; then
+  FLAGS_board_root="/build/${FLAGS_board}"
+fi
+
+# Only now can we die on error.  shflags functions leak non-zero error codes,
+# so will die prematurely if 'set -e' is specified before now.
+set -e
+# Die on uninitialized variables.
+set -u
+
+# Check for missing parts.
+ROOTFS_IMG="${IMAGEDIR}/rootfs.image"
+if [[ ! -s ${ROOTFS_IMG} ]]; then
+  error "Can't find ${ROOTFS_IMG}"
+  exit 1
+fi
+
+KERNEL_IMG="${IMAGEDIR}/vmlinuz.image"
+if [[ ! -s ${KERNEL_IMG} ]]; then
+  error "Can't find ${KERNEL_IMG}"
+  exit 1
+fi
+
+STATEFUL_IMG="${IMAGEDIR}/stateful_partition.image"
+if [[ ! -s ${STATEFUL_IMG} ]]; then
+  error "Can't find ${STATEFUL_IMG}"
+  exit 1
+fi
+
+# We'll need some code to put in the PMBR, for booting on legacy BIOS. Some ARM
+# systems will use a U-Boot script temporarily, but it goes in the same place.
+if [[ "$ARCH" = "arm" ]]; then
+  # U-Boot script copies the kernel into memory, then boots it.
+  KERNEL_OFFSET=$(printf "0x%08x" ${START_KERN_A})
+  KERNEL_SECS_HEX=$(printf "0x%08x" ${NUM_KERN_BLOCKS})
+  MBR_SCRIPT="${IMAGEDIR}/mbr_script"
+  echo -e "echo\necho ---- ChromeOS Boot ----\necho\n" \
+          "mmc read 1 C0008000 $KERNEL_OFFSET $KERNEL_SECS_HEX\n" \
+          "bootm C0008000" > ${MBR_SCRIPT}
+  MKIMAGE="${FLAGS_board_root}/u-boot/mkimage"
+  if [[ -f "$MKIMAGE".gz ]]; then
+    sudo gunzip "$MKIMAGE".gz
+  fi
+  if [[ -x "$MKIMAGE" ]]; then
+    MBR_SCRIPT_UIMG="${MBR_SCRIPT}.uimg"
+    "$MKIMAGE" -A "${ARCH}" -O linux  -T script -a 0 -e 0 -n "COS boot" \
+               -d ${MBR_SCRIPT} ${MBR_SCRIPT_UIMG}
+    MBR_IMG=${IMAGEDIR}/mbr.img
+    dd bs=1 count=`stat --printf="%s" ${MBR_SCRIPT_UIMG}` \
+       if="$MBR_SCRIPT_UIMG" of="$MBR_IMG" conv=notrunc
+    hexdump -v -C "$MBR_IMG"
+  else
+    echo "Error: u-boot mkimage not found or not executable."
+  fi
+  PMBRCODE=${MBR_IMG}
+else
+  PMBRCODE=$(readlink -f /usr/share/syslinux/gptmbr.bin)
+fi
+
+# Create the GPT. This has the side-effect of setting some global vars
+# describing the partition table entries (see the comments in the source).
+install_gpt $OUTDEV $ROOTFS_IMG $KERNEL_IMG $STATEFUL_IMG $PMBRCODE
+
+# Now populate the partitions.
+echo "Copying stateful partition..."
+dd if=${STATEFUL_IMG} of=${OUTDEV} conv=notrunc bs=512 seek=${START_STATEFUL}
+
+echo "Copying kernel..."
+dd if=${KERNEL_IMG} of=${OUTDEV} conv=notrunc bs=512 seek=${START_KERN_A}
+  
+echo "Copying rootfs..."
+dd if=${ROOTFS_IMG} of=${OUTDEV} conv=notrunc bs=512 seek=${START_ROOTFS_A}
+

build_image

@@ -70,7 +70,7 @@ OUTPUT_DIR="${FLAGS_output_root}/${FLAGS_board}/${IMAGE_SUBDIR}"
 ROOT_FS_DIR="${OUTPUT_DIR}/rootfs"
 ROOT_FS_IMG="${OUTPUT_DIR}/rootfs.image"
 MBR_IMG="${OUTPUT_DIR}/mbr.image"
-OUTPUT_IMG="${OUTPUT_DIR}/usb.img"
+OUTPUT_IMG="${OUTPUT_DIR}/chromiumos_image.bin"
 
 BOARD="${FLAGS_board}"
 BOARD_ROOT="${FLAGS_build_root}/${BOARD}"
@@ -320,6 +320,11 @@ if [[ $FLAGS_withdev -eq $FLAGS_TRUE ]]; then
   WITH_DEV="--withdev"
 fi
 
+# Extract the kernel from the root filesystem for use by the GPT image. Legacy
+# BIOS will use the kernel in the rootfs (via syslinux), ChromeOS BIOS will use
+# the kernel partition.
+sudo cp -f ${ROOT_FS_DIR}/boot/vmlinuz ${OUTPUT_DIR}/vmlinuz.image
+
 #TODO(sosa@chromium.org) - Does it make sense to leave /usr/local bound here?
 "${SCRIPTS_DIR}/customize_rootfs" \
   --root="$ROOT_FS_DIR" \
@@ -347,44 +352,11 @@ fi
 cleanup_stateful_fs_loop
 cleanup_rootfs_loop
 
-# Create a master boot record.
+# Create the GPT-formatted image
-# Start with the syslinux master boot record. We need to zero-pad to
+${SCRIPTS_DIR}/build_gpt.sh \
-# fill out a 512-byte sector size.
+   --arch=${ARCH} --board=${FLAGS_board} --board_root=${BOARD_ROOT} \
-SYSLINUX_MBR="/usr/share/syslinux/mbr.bin"
+   ${OUTPUT_DIR} ${OUTPUT_IMG}
-dd if="$SYSLINUX_MBR" of="$MBR_IMG" bs=512 count=1 conv=sync
+
-# Create a partition table in the MBR.
-NUM_SECTORS=$((`stat --format="%s" "$ROOT_FS_IMG"` / 512))
-KERNEL_SECTORS=8192
-sudo sfdisk -H64 -S32 -uS -f "$MBR_IMG" <<EOF
-,$NUM_SECTORS,L,-,
-,$NUM_SECTORS,S,-,
-,$NUM_SECTORS,L,*,
-,$KERNEL_SECTORS,L,-,
-;
-EOF
-if [[ "$ARCH" = "arm" ]]; then
-  # Write u-boot script into MBR code block that boots 4th partition kernel.
-  KERNEL_OFFSET=`printf "%08x" $(((3 * $NUM_SECTORS) + 1))`
-  KERNEL_SECS_HEX=`printf "%08x" $KERNEL_SECTORS`
-  MBR_SCRIPT="${OUTPUT_DIR}/mbr_script"
-  echo -e "echo\necho ---- ChromeOS Boot ----\necho\n" \
-          "mmc read 1 C0008000 0x$KERNEL_OFFSET 0x$KERNEL_SECS_HEX\n" \
-          "bootm C0008000" > ${MBR_SCRIPT}
-  MKIMAGE="${BOARD_ROOT}/u-boot/mkimage"
-  if [[ -f "$MKIMAGE".gz ]]; then
-    sudo gunzip "$MKIMAGE".gz
-  fi
-  if [[ -x "$MKIMAGE" ]]; then
-    MBR_SCRIPT_UIMG="${MBR_SCRIPT}.uimg"
-    "$MKIMAGE" -A "${ARCH}" -O linux  -T script -a 0 -e 0 -n "COS boot" \
-               -d ${MBR_SCRIPT} ${MBR_SCRIPT_UIMG}
-    dd bs=1 count=`stat --printf="%s" ${MBR_SCRIPT_UIMG}` \
-       if="$MBR_SCRIPT_UIMG" of="$MBR_IMG" conv=notrunc
-    hexdump -v -C "$MBR_IMG"
-  else
-    echo "Error: u-boot mkimage not found or not executable."
-  fi
-fi
 
 OUTSIDE_OUTPUT_DIR="../build/images/${FLAGS_board}/${IMAGE_SUBDIR}"
 echo "Done.  Image created in ${OUTPUT_DIR}"

chromeos-common.sh

@@ -0,0 +1 @@
+../platform/installer/chromeos-common.sh

image_to_usb.sh

@@ -10,6 +10,9 @@
 # The path to common.sh should be relative to your script's location.
 . "$(dirname "$0")/common.sh"
 
+# Load functions and constants for chromeos-install
+. "$(dirname "$0")/chromeos-common.sh"
+
 get_default_board
 
 # Flags
@@ -66,9 +69,10 @@ if [ -z "${FLAGS_to}" ]; then
   # Script can be run either inside or outside the chroot.
   if [ ${INSIDE_CHROOT} -eq 1 ]
   then
-    # Inside the chroot, so output to usb.img in the same dir as the other
+    # Inside the chroot, we'll only output to a file, and we probably already
-    # images.
+    # have a valid image. Make the user specify a filename.
-    FLAGS_to="${FLAGS_from}/usb.img"
+    echo "ERROR: Inside the chroot, you must specify a --to FILE to create."
+    exit 1
   else
     # Outside the chroot, so output to the default device for a usb key.
     FLAGS_to="/dev/sdb"
@@ -80,46 +84,23 @@ fi
 FLAGS_from=`eval readlink -f ${FLAGS_from}`
 FLAGS_to=`eval readlink -f ${FLAGS_to}`
 
-# Uses this rootfs image as the source image to copy
+# Use this image as the source image to copy
-ROOTFS_IMAGE="${FLAGS_from}/rootfs.image"
+SRC_IMAGE="${FLAGS_from}/chromiumos_image.bin"
-PART_SIZE=$(stat -c%s "${ROOTFS_IMAGE}")  # Bytes
 
-# Setup stateful partition variables
+# If we're asked to modify the image for test, then let's make a copy and
-STATEFUL_IMG="${FLAGS_from}/stateful_partition.image"
+# modify that instead.
-STATEFUL_DIR="${FLAGS_from}/stateful_partition"
-
-# TODO(sosa@chromium.org) - Remove legacy support.
-if [ ! -f "${STATEFUL_IMG}" ] ; then
-  echo "WARNING!  Stateful partition not found.  Creating clean stateful"
-  STATEFUL_LOOP_DEV=$(sudo losetup -f)
-  if [ -z "${STATEFUL_LOOP_DEV}" ] ; then
-    echo "No free loop device.  Free up a loop device or reboot.  exiting. "
-    exit 1
-  fi
-  set -x
-  dd if=/dev/zero of="${STATEFUL_IMG}" bs=1 count=1 \
-      seek=$(( (${PART_SIZE} - 1) ))
-  set +x
-  trap do_cleanup INT TERM EXIT
-  sudo losetup "$STATEFUL_LOOP_DEV" "$STATEFUL_IMG"
-  sudo mkfs.ext3 "$STATEFUL_LOOP_DEV"
-  sudo tune2fs -L "C-STATE" -c 0 -i 0 "$STATEFUL_LOOP_DEV"
-  sudo losetup -d "${STATEFUL_LOOP_DEV}"
-  trap - INT TERM EXIT
-fi
-
-# Modifies image for test if requested
 if [ ${FLAGS_test_image} -eq ${FLAGS_TRUE} ] ; then
-  if [ ! -f "${FLAGS_from}/rootfs_test.image" ] ; then
+  # Copy it.
-    echo "Test image not found, creating test image from original ... "
+  echo "Creating test image from original..."
-    cp "${FLAGS_from}/rootfs.image" "${FLAGS_from}/rootfs_test.image"
+  cp -f "${SRC_IMAGE}" "${FLAGS_from}/chromiumos_test_image.bin"
-    "${SCRIPTS_DIR}/mod_image_for_test.sh" \
+  # Use it.
-      --image "${FLAGS_from}/rootfs_test.image"
+  SRC_IMAGE="${FLAGS_from}/chromiumos_test_image.bin"
-  fi
+  # Modify it.
-  # Use the test image instead
+  "${SCRIPTS_DIR}/mod_image_for_test.sh" --image "${SRC_IMAGE}"
-  ROOTFS_IMAGE="${FLAGS_from}/rootfs_test.image"
 fi
 
+STATEFUL_DIR="${FLAGS_from}/stateful_partition"
+
 function do_cleanup {
   echo "Cleaning loopback devices: ${STATEFUL_LOOP_DEV}"
   if [ "${STATEFUL_LOOP_DEV}" != "" ]; then
@@ -129,26 +110,32 @@ function do_cleanup {
   fi
 }
 
-function install_autotest {
+if [ ${FLAGS_install_autotest} -eq ${FLAGS_TRUE} ] ; then
   echo "Detecting autotest at ${AUTOTEST_SRC}"
   if [ -d ${AUTOTEST_SRC} ]
   then
-    local stateful_loop_dev=$(sudo losetup -f)
+    # Figure out how to loop mount the stateful partition. It's always
-    local stateful_root="${STATEFUL_DIR}/dev_image"
+    # partition 1 on the disk image.
+    offset=$(partoffset "${SRC_IMAGE}" 1)
+
+    stateful_loop_dev=$(sudo losetup -f)
     if [ -z "${stateful_loop_dev}" ]
     then
       echo "No free loop device. Free up a loop device or reboot. exiting."
       exit 1
     fi
-    trap do_cleanup INT TERM EXIT
     STATEFUL_LOOP_DEV=$stateful_loop_dev
+    trap do_cleanup INT TERM EXIT
+
     echo "Mounting ${STATEFUL_DIR} loopback"
-    sudo losetup "${stateful_loop_dev}" "${STATEFUL_DIR}.image"
+    sudo losetup -o $(( $offset * 512 )) "${stateful_loop_dev}" "${SRC_IMAGE}"
     sudo mount "${stateful_loop_dev}" "${STATEFUL_DIR}"
-    
+    stateful_root="${STATEFUL_DIR}/dev_image"
-    echo -ne "Install autotest into stateful partition..."
+
-    local autotest_client="/home/autotest-client"
+    echo "Install autotest into stateful partition..."
-    sudo mkdir -p "${stateful_root}/${autotest_client}"
+    autotest_client="/home/autotest-client"
+    sudo mkdir -p "${stateful_root}${autotest_client}"
+
     sudo cp -fpru ${AUTOTEST_SRC}/client/* \
 	    "${stateful_root}/${autotest_client}"
     sudo chmod 755 "${stateful_root}/${autotest_client}"
@@ -162,17 +149,18 @@ function install_autotest {
     echo "Please call make_autotest.sh inside chroot first."
     exit -1
   fi
-}
+fi
 
-# Copy MBR and rootfs to output image
+
+# Let's do it.
 if [ -b "${FLAGS_to}" ]
 then
   # Output to a block device (i.e., a real USB key), so need sudo dd
-  echo "Copying USB image ${FLAGS_from} to device ${FLAGS_to}..."
+  echo "Copying USB image ${SRC_IMAGE} to device ${FLAGS_to}..."
 
   # Warn if it looks like they supplied a partition as the destination.
-  if echo ${FLAGS_to} | grep -q '[0-9]$'; then
+  if echo "${FLAGS_to}" | grep -q '[0-9]$'; then
-    local drive=$(echo ${FLAGS_to} | sed -re 's/[0-9]+$//')
+    local drive=$(echo "${FLAGS_to}" | sed -re 's/[0-9]+$//')
     if [ -b "${drive}" ]; then
       echo
       echo "NOTE: It looks like you may have supplied a partition as the "
@@ -185,8 +173,8 @@ then
   # Make sure this is really what the user wants, before nuking the device
   if [ ${FLAGS_yes} -ne ${FLAGS_TRUE} ]
   then
+    sudo fdisk -l "${FLAGS_to}" 2>/dev/null | grep Disk | head -1
     echo "This will erase all data on this device:"
-    sudo fdisk -l "${FLAGS_to}" | grep Disk | head -1
     read -p "Are you sure (y/N)? " SURE
     SURE="${SURE:0:1}" # Get just the first character
     if [ "${SURE}" != "y" ]
@@ -196,72 +184,30 @@ then
     fi
   fi
 
-  echo "attempting to unmount any mounts on the USB device"
+  echo "Attempting to unmount any mounts on the USB device..."
-  for i in "${FLAGS_to}"*
+  for i in $(mount | grep ^"${FLAGS_to}" | awk '{print $1}')
   do
-    ! sudo umount "$i"
+    sudo umount "$i"
   done
   sleep 3
 
-  if [ ${FLAGS_install_autotest} -eq ${FLAGS_TRUE} ] ; then
+  echo "Copying ${SRC_IMAGE} to ${FLAGS_to}..."
-    install_autotest
+  sudo dd if="${SRC_IMAGE}" of="${FLAGS_to}" bs=4M
-  fi
 
-  # Write stateful partition to first partition. 
-  echo "Copying stateful partition ..."
-  sudo "${SCRIPTS_DIR}"/file_copy.py \
-      if="${STATEFUL_IMG}" of="${FLAGS_to}" bs=4M \
-      seek_bytes=512
-
-  # Write root fs to third partition.
-  echo "Copying root fs partition ..."
-  sudo "${SCRIPTS_DIR}"/file_copy.py \
-      if="${ROOTFS_IMAGE}" of="${FLAGS_to}" bs=4M \
-      seek_bytes=$(( (${PART_SIZE} * 2) + 512 ))
-
-  trap - EXIT
-
-  if [ ${FLAGS_copy_kernel} -eq ${FLAGS_TRUE} ]
-  then
-    echo "Copying Kernel..."
-    "${SCRIPTS_DIR}"/kernel_fetcher.sh \
-      --from "${FLAGS_from}" \
-      --to "${FLAGS_to}" \
-      --offset "$(( (${PART_SIZE} * 3) + 512 ))"
-  fi
-
-  echo "Copying MBR..."
-  sudo "${SCRIPTS_DIR}"/file_copy.py \
-    if="${FLAGS_from}/mbr.image" of="${FLAGS_to}"
-  sync
   echo "Done."
 else
-  # Output to a file, so just cat the source images together
+  # Output to a file, so just make a copy.
-
+  echo "Copying ${SRC_IMAGE} to ${FLAGS_to}..."
-  PART_SIZE=$(stat -c%s "${ROOTFS_IMAGE}")
+  cp -f "${SRC_IMAGE}" "${FLAGS_to}"
-
-  if [ ${FLAGS_install_autotest} -eq ${FLAGS_TRUE} ] ; then
-    install_autotest
-  fi
-
-  # Create a sparse output file
-  dd if=/dev/zero of="${FLAGS_to}" bs=1 count=1 \
-      seek=$(( (${PART_SIZE} * 2) + 512 - 1))
-
-  echo "Copying USB image to file ${FLAGS_to}..."
-
-  dd if="${FLAGS_from}/mbr.image" of="${FLAGS_to}" conv=notrunc
-  dd if="${FLAGS_from}/stateful_partition.image" of="${FLAGS_to}" seek=1 bs=512 \
-      conv=notrunc
-  cat "${ROOTFS_IMAGE}" >> "${FLAGS_to}"
 
   echo "Done.  To copy to USB keyfob, outside the chroot, do something like:"
   echo "   sudo dd if=${FLAGS_to} of=/dev/sdb bs=4M"
   echo "where /dev/sdb is the entire keyfob."
   if [ ${INSIDE_CHROOT} -eq 1 ]
   then
+    example=$(basename "${FLAGS_to}")
     echo "NOTE: Since you are currently inside the chroot, and you'll need to"
     echo "run dd outside the chroot, the path to the USB image will be"
-    echo "different (ex: ~/chromeos/trunk/src/build/images/SOME_DIR/usb.img)."
+    echo "different (ex: ~/chromeos/trunk/src/build/images/SOME_DIR/$example)."
   fi
 fi

mod_image_for_test.sh

@@ -10,6 +10,9 @@
 # The path to common.sh should be relative to your script's location.
 . "$(dirname "$0")/common.sh"
 
+# Load functions and constants for chromeos-install
+. "$(dirname "$0")/chromeos-common.sh"
+
 get_default_board
 
 DEFINE_string board "$DEFAULT_BOARD" "Board for which the image was built"
@@ -24,7 +27,7 @@ FLAGS "$@" || exit 1
 eval set -- "${FLAGS_ARGV}"
 
 # No board, no default and no image set then we can't find the image
-if [ -z $FLAGS_IMAGE ] && [ -z $FLAGS_board ] ; then
+if [ -z $FLAGS_image ] && [ -z $FLAGS_board ] ; then
   setup_board_warning
   echo "*** mod_image_for_test failed.  No board set and no image set"
   exit 1
@@ -33,7 +36,8 @@ fi
 # 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}"
-  FLAGS_image="${IMAGES_DIR}/$(ls -t $IMAGES_DIR 2>&-| head -1)/rootfs.image"
+  FILENAME="chromiumos_test_image.bin"
+  FLAGS_image="${IMAGES_DIR}/$(ls -t $IMAGES_DIR 2>&-| head -1)/${FILENAME}"
 fi
 
 # Abort early if we can't find the image
@@ -51,7 +55,7 @@ cleanup_rootfs_mounts() {
   do
     local cmdline=`cat /proc/$pid/cmdline`
     echo "Killing process that has open file on our rootfs: $cmdline"
-    ! sudo kill $pid  # Preceded by ! to disable ERR trap.
+    sudo kill $pid || /bin/true
   done
 }
 
@@ -76,15 +80,6 @@ cleanup() {
 }
 
 # main process begins here.
-set -e
-trap cleanup EXIT
-
-ROOT_FS_DIR="`dirname ${FLAGS_image}`/rootfs"
-mkdir -p "${ROOT_FS_DIR}"
-
-LOOP_DEV=`sudo losetup -f`
-sudo losetup "${LOOP_DEV}" "${FLAGS_image}"
-sudo mount "${LOOP_DEV}" "${ROOT_FS_DIR}"
 
 # Make sure this is really what the user wants, before nuking the device
 if [ $FLAGS_yes -ne $FLAGS_TRUE ]; then
@@ -98,9 +93,27 @@ else
   echo "Modifying image ${FLAGS_image} for test..."
 fi
 
+set -e
+trap cleanup EXIT
+
+ROOT_FS_DIR=$(dirname "${FLAGS_image}")/rootfs
+mkdir -p "${ROOT_FS_DIR}"
+
+# Figure out how to loop mount the rootfs partition. It should be partition 3
+# on the disk image.
+offset=$(partoffset "${FLAGS_image}" 3)
+
+LOOP_DEV=$(sudo losetup -f)
+if [ -z "$LOOP_DEV" ]; then
+  echo "No free loop device"
+  exit 1
+fi
+sudo losetup -o $(( $offset * 512 )) "${LOOP_DEV}" "${FLAGS_image}"
+sudo mount "${LOOP_DEV}" "${ROOT_FS_DIR}"
+
 MOD_SCRIPTS_ROOT="${GCLIENT_ROOT}/src/scripts/mod_for_test_scripts"
 # Run test setup script inside chroot jail to modify the image
-sudo GCLIENT_ROOT=${GCLIENT_ROOT} ROOT_FS_DIR=${ROOT_FS_DIR} \
+sudo GCLIENT_ROOT="${GCLIENT_ROOT}" ROOT_FS_DIR="${ROOT_FS_DIR}" \
     "${MOD_SCRIPTS_ROOT}/test_setup.sh"
 
 # Run manufacturing test setup