Merge "pbuilder: Update native to USERSPACE_ARCH"
diff --git a/flash.sh b/flash.sh
index 9996207..b6b30c9 100755
--- a/flash.sh
+++ b/flash.sh
@@ -14,6 +14,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+ROOTDIR=$(dirname $0)/..
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+FASTBOOT_CMD="$(which fastboot)"
+
 function partition_table_image {
     if [[ $1 -gt 7000000000 ]] && [[ $1 -lt 8000000000 ]]; then
         echo "partition-table-8gb.img"
@@ -30,77 +34,256 @@
 }
 
 function try {
-    $@ || die "Failed to execute $@"
+    $@ || die "Failed to execute '$@'"
 }
 
-ROOTDIR=$(dirname $0)/..
-SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+function ensure_flash_files_present
+{
+    local found_all_files=true
+    local files_dir="$1"; shift
+    local arch="$1"; shift
+    local files=(
+        "u-boot.imx"
+        "home.img"
+        "partition-table-8gb.img"
+        "partition-table-16gb.img"
+        "partition-table-64gb.img"
+        "boot_${arch}.img"
+        "rootfs_${arch}.img"
+    )
 
-FASTBOOT_CMD="$(which fastboot)"
-if [[ ! -x ${FASTBOOT_CMD} ]]; then
-    die "Couldn't find fastboot on your PATH -- did you install it?"
-fi
+    for file in "${files[@]}"; do
+        if [[ ! -f "${files_dir}/${file}" ]]; then
+            echo "${file} is missing"
+            found_all_files=""
+        fi
+    done
 
-if [[ -n ${SERIAL} ]]; then
-    FASTBOOT_CMD="fastboot -s ${SERIAL}"
-fi
-
-# Set USERSPACE_ARCH to arm64, if not set in the environment
-USERSPACE_ARCH=${USERSPACE_ARCH:=arm64}
-
-# If there's a boot.img in the same folder as flash.sh,
-# find artifacts in the directory containing the script.
-if [ -e ${SCRIPT_DIR}/boot_${USERSPACE_ARCH}.img ]; then
-    PRODUCT_OUT=${SCRIPT_DIR}
-else
-    PRODUCT_OUT=${PRODUCT_OUT:=${ROOTDIR}/out/target/product/imx8m_phanbell}
-fi
-
-# Figure out which partition map we need based upon fastboot vars
-MMC_SIZE=$(${FASTBOOT_CMD} getvar mmc_size 2>&1 | awk '/mmc_size:/ { print $2 }')
-PART_IMAGE=$(partition_table_image ${MMC_SIZE})
-
-if [[ -z ${PART_IMAGE} ]]; then
-    die "No partition map available for an emmc of size ${MMC_SIZE}"
-fi
-
-for i in u-boot.imx boot_${USERSPACE_ARCH}.img rootfs_${USERSPACE_ARCH}.img ${PART_IMAGE}; do
-    [[ ! -f ${PRODUCT_OUT}/$i ]] && die "${PRODUCT_OUT}/$i is missing."
-done
-
-# Flash bootloader
-try ${FASTBOOT_CMD} flash bootloader0 ${PRODUCT_OUT}/u-boot.imx
-try ${FASTBOOT_CMD} reboot-bootloader
-sleep 3
-
-# Check to see if HOME existed before we clobber the GPT
-HOME_MISSING=$(${FASTBOOT_CMD} getvar partition-type:home 2>&1 |grep -q FAILED && echo true)
-if [[ "${HOME_MISSING}" == "true" ]]; then
-    echo "*** home partition is missing -- will flash home partition"
-else
-    echo -n "*** home partition present; "
-
-    if [[ "$1" == "--clobber-home" ]]; then
-        echo "--clobber-home specified. WILL flash home partition"
-    else
-        echo "WILL NOT flash home partition"
+    if [[ -z "${found_all_files}" ]]; then
+        die "Required image files are missing. Cannot continue."
     fi
-fi
+}
 
-# Flash partition table
-try ${FASTBOOT_CMD} flash gpt ${PRODUCT_OUT}/${PART_IMAGE}
-try ${FASTBOOT_CMD} reboot-bootloader
-sleep 3
+function flash
+{
+    local partition="$1"; shift
+    local file="$1"; shift
 
-# Flash filesystems
-try ${FASTBOOT_CMD} erase misc
+    try ${FASTBOOT_CMD} flash "${partition}" "${file}"
+}
 
-# If home was present before we stuffed in the new GPT, only clobber
-# if the user asked for it.
-if [[ "${HOME_MISSING}" == "true" ]] || [[ "$1" == "--clobber-home" ]] ; then
-    try ${FASTBOOT_CMD} flash home ${PRODUCT_OUT}/home.img
-fi
+function erase
+{
+    local partition="$1"; shift
+    try ${FASTBOOT_CMD} erase "${partition}"
+}
 
-try ${FASTBOOT_CMD} flash boot ${PRODUCT_OUT}/boot_${USERSPACE_ARCH}.img
-try ${FASTBOOT_CMD} flash rootfs ${PRODUCT_OUT}/rootfs_${USERSPACE_ARCH}.img
-try ${FASTBOOT_CMD} reboot
+function reboot_bootloader
+{
+    try ${FASTBOOT_CMD} reboot-bootloader
+}
+
+function get_emmc_size
+{
+    local output=$(${FASTBOOT_CMD} getvar mmc_size 2>&1)
+    if [[ "$?" != "0" ]]; then
+        die "Couldn't get emmc size (nonzero exit, output was [${output}])"
+    fi
+
+    local emmc_size=$(echo $output |awk '/mmc_size:/ { print $2 }')
+
+    if [[ -z "${emmc_size}" ]]; then
+        die "Can't get emmc size (output was [${emmc_size}])"
+    fi
+
+    echo "${emmc_size}"
+}
+
+function partition_present
+{
+    local partition="$1"; shift
+    local output=$(${FASTBOOT_CMD} getvar "partition-type:${partition}" 2>&1)
+    if [[ "$?" != "0" ]]; then
+        die "Couldn't communicate with fastboot (nonzero exit, output was [${output}])"
+    fi
+
+    if [[ "${output}" =~ "FAILED" ]]; then
+        return 1
+    else
+        return 0
+    fi
+}
+
+function flash_partitions
+{
+    local files_dir="$1"; shift
+    local arch="$1"; shift
+    local overwrite_home="$1"; shift
+    local will_flash_home=true
+    local mmc_size=$(get_emmc_size)
+    local part_image=$(partition_table_image "${mmc_size}")
+
+    if [[ ! -f "${files_dir}/${part_image}" ]]; then
+        die "Can't find partition image ${part_image} to flash"
+    fi
+
+    if partition_present home && [[ -z "${overwrite_home}" ]]; then
+        will_flash_home=""
+        echo "Existing home partition detected, will not flash home (override with -H)"
+    fi
+
+    flash bootloader0 "${files_dir}/u-boot.imx"
+    reboot_bootloader
+    sleep 3
+
+    flash gpt "${files_dir}/${part_image}"
+    reboot_bootloader
+    sleep 3
+
+    erase misc
+
+    if [[ ! -z "${will_flash_home}" ]]; then
+        flash home "${files_dir}/home.img"
+    fi
+
+    flash boot "${files_dir}/boot_${arch}.img"
+    flash rootfs "${files_dir}/rootfs_${arch}.img"
+}
+
+function main {
+    local usage=$(cat <<EOF
+Usage: flash.sh [-n] [-H] [-d <files_dir>] [-s <serial>] [-r <detect_retries>]
+
+  -n                  prevents the automatic reboot of the connected board
+  -H                  overwrites the home partition if one is detected
+  -d <files_dir>      flashes files from <files_dir> (defaults to current dir)
+  -s <serial>         only flashes the board with the given fastboot serial number
+  -r <detect_retries> number of times to retry waiting for a device (defaults to 0)
+  -u <userspace_arch> sets the name of the userspace to flash (defaults to arm64)
+EOF
+)
+
+    local dont_reboot                # -n
+    local overwrite_home             # -H
+    local serial_number              # -s <serial>
+    local arch="arm64"               # -a <userspace_arch>
+    local files_dir                  # -d <files_dir>
+    local detect_retries=0           # -r <detect_retries>
+    local args=$(getopt hnHs:d:r: $*)
+
+    set -- $args
+
+    for i; do
+        case "$1" in
+            -n)  # don't rebooot
+                dont_reboot=true
+                shift 1
+                ;;
+
+            -H)  # overwrite home if it exists
+                overwrite_home=true
+                shift 1
+                ;;
+
+            -s)  # serial <serial_no>
+                serial_number="$2"
+                shift 2
+                ;;
+
+            -a)  # userspace architecture <arch>
+                arch="$2"
+                shift 2
+                ;;
+
+            -d)  # files dir <files_dir>
+                files_dir="$2"
+                shift 2
+                ;;
+
+            -r)  # detection retries <detect_retries>
+                detect_retries="$2"
+                if ! [[ "${detect_retries}" -gt "0" ]]; then
+                    die "Retry count \"${detect_retries}\" is not a positive integer"
+                    return 1
+                fi
+                shift 2
+                ;;
+
+            --)  # end of args
+                shift
+                break
+                ;;
+
+            *)   # help
+                die "${usage}"
+                ;;
+        esac
+    done
+
+    if [[ ! -x "${FASTBOOT_CMD}" ]]; then
+        die "Couldn't find fastboot on your PATH -- did you install it?"
+    fi
+
+    if [[ -n "${serial_number}" ]]; then
+        FASTBOOT_CMD="fastboot -s ${serial_number}"
+    fi
+
+    if [[ -z "${files_dir}" ]]; then
+        if [[ -n "${PRODUCT_OUT}" ]]; then
+            files_dir="${PRODUCT_OUT}"
+        else
+            files_dir="."
+        fi
+    fi
+
+    # Start the flashing process
+
+    if [[ -n "${serial_number}" ]]; then
+        detect_device_or_die "${detect_retries}" "${serial_number}"
+    fi
+
+    ensure_flash_files_present "${files_dir}" "${arch}"
+    flash_partitions "${files_dir}" "${arch}" "${overwrite_home}"
+
+    if [[ -z "${dont_reboot}" ]]; then
+        try ${FASTBOOT_CMD} reboot
+    else
+        echo "Skipping reboot as requested -- manual reboot will be required"
+    fi
+
+    echo "Flash completed."
+}
+
+function detect_device_or_die
+{
+    local retries="$1"; shift
+    local serial="$1"; shift
+    local count=0
+    local output
+    local found
+
+    echo -n "Waiting up to ${retries} seconds for device ${serial}"
+    for ((count=0; count<"${retries}"; count++)); do
+        output=$(${FASTBOOT_CMD} devices)
+        if [[ "$?" != "0" ]]; then
+            die "Unable to communicate with fastboot (nonzero exit, output was [${output}])"
+        fi
+
+        if [[ "${output}" =~ "${serial}" ]]; then
+            found=true
+            break
+        fi
+
+        echo -n '.'
+        sleep 1
+    done
+    echo
+
+    if [[ -z "${found}" ]]; then
+        die "Couldn't find device with serial ${serial} in fastboot output"
+    else
+        echo "Found device ${serial}"
+    fi
+}
+
+main "$@"
diff --git a/flashcard.mk b/flashcard.mk
index e5ab1b4..18a1699 100644
--- a/flashcard.mk
+++ b/flashcard.mk
@@ -31,7 +31,7 @@
 flashcard: $(FLASHCARD_PATH)
 
 $(FLASHCARD_PATH): $(ROOTDIR)/board/flashcard/init $(ROOTDIR)/board/flashcard/boot.txt multistrap | \
-                   busybox partition-table $(PRODUCT_OUT)/u-boot.imx $(HOST_OUT)/bin/mkimage
+                   busybox partition-table $(PRODUCT_OUT)/u-boot.imx $(HOST_OUT)/bin/mkimage home_raw
 	# TODO(math the size)
 	-rm -rf $(FLASHCARD_WORK_DIR)
 	mkdir -p $(FLASHCARD_WORK_DIR)
@@ -90,6 +90,7 @@
 	cp $(PRODUCT_OUT)/u-boot.imx $(FLASHCARD_MOUNT_PATH)
 	cp $(PRODUCT_OUT)/multistrap/boot_$(USERSPACE_ARCH).img $(FLASHCARD_MOUNT_PATH)/boot.img
 	cp $(PRODUCT_OUT)/multistrap/rootfs_$(USERSPACE_ARCH).img.wip $(FLASHCARD_MOUNT_PATH)/rootfs.img
+	cp $(PRODUCT_OUT)/obj/HOME/home.raw.img $(FLASHCARD_MOUNT_PATH)/home.img
 
 	sudo umount $(FLASHCARD_BOOT_MOUNT_PATH)
 	sudo umount $(FLASHCARD_MOUNT_PATH)
diff --git a/flashcard/init b/flashcard/init
index 9af18d3..7e58a9f 100755
--- a/flashcard/init
+++ b/flashcard/init
@@ -44,6 +44,9 @@
 log "* Writing boot filesystem..."
 dd if=/tmp/sdcard/boot.img of=/dev/mmcblk0p1
 
+log "* Writing home filesystem..."
+dd if=/tmp/sdcard/home.img of=/dev/mmcblk0p3
+
 log "* Writing root filesystem..."
 dd if=/tmp/sdcard/rootfs.img of=/dev/mmcblk0p4