From 28a753b22cb0e0515ab4494e11ea74f031a210d4 Mon Sep 17 00:00:00 2001
From: Hung-Te Lin <hungte@chromium.org>
Date: Mon, 15 Aug 2011 15:27:12 +0800
Subject: [PATCH] crosutils: refine factory and imaging scripts

In order to support building arbitrary image, the partition copying scripts has
been changed to support "copying partitions in same size" and "overwriting
partitions in different size", and "copying partition from external file".

We need these APIs to create disk/usb image with release images that is using
partition with different size (ex, recovery images).

Image copying buffer selection and disk image creation time are also improved.

BUG=chromium-os:15050
TEST=./make_factory_package.sh ... --diskimg preimage.bin
     ./make_factory_package.sh ... --usbimg rma.bin
     ./make_factory_package.sh ...  # omaha mode

Change-Id: I6a4c820abf59e780985c95dc35f9340b347bd952
Reviewed-on: http://gerrit.chromium.org/gerrit/5981
Reviewed-by: Nick Sanders <nsanders@chromium.org>
Tested-by: Hung-Te Lin <hungte@chromium.org>
---
 lib/cros_image_common.sh | 104 +++++++++++++++++++++++++++------------
 make_factory_package.sh  |  19 +++----
 2 files changed, 79 insertions(+), 44 deletions(-)

diff --git a/lib/cros_image_common.sh b/lib/cros_image_common.sh
index 7c9befcc17..c4732c221c 100644
--- a/lib/cros_image_common.sh
+++ b/lib/cros_image_common.sh
@@ -92,15 +92,13 @@ image_dump_partial_file() {
   local sectors="$3"
   local bs=512
 
-  # Try to use larger buffer if offset/size can be re-aligned.
-  # 2M / 512 = 4096
-  local buffer_ratio=4096
-  if [ $((offset % buffer_ratio)) -eq 0 -a \
-       $((sectors % buffer_ratio)) -eq 0 ]; then
-    offset=$((offset / buffer_ratio))
-    sectors=$((sectors / buffer_ratio))
-    bs=$((bs * buffer_ratio))
-  fi
+  # Increase buffer size as much as possible until 8M
+  while [ $((bs < (8 * 1024 * 1024) && sectors > 0 &&
+             offset % 2 == 0 && sectors % 2 == 0)) = "1" ]; do
+    bs=$((bs * 2))
+    offset=$((offset / 2))
+    sectors=$((sectors / 2))
+  done
 
   if image_has_command pv; then
     dd if="$file" bs=$bs skip="$offset" count="$sectors" \
@@ -124,6 +122,43 @@ image_dump_partition() {
   image_dump_partial_file "$file" "$offset" "$size"
 }
 
+# Updates a file (from stdin) by given offset and size (in sectors)
+image_update_partial_file() {
+  local file="$1"
+  local offset="$2"
+  local sectors="$3"
+  local bs=512
+
+  # Increase buffer size as much as possible until 8M
+  while [ $((bs < (8 * 1024 * 1024) && sectors > 0 &&
+             offset % 2 == 0 && sectors % 2 == 0)) = "1" ]; do
+    bs=$((bs * 2))
+    offset=$((offset / 2))
+    sectors=$((sectors / 2))
+  done
+
+  if image_has_command pv; then
+    pv -ptreb -B $bs -s $((sectors * bs)) |
+      dd of="$file" bs=$bs seek="$offset" count="$sectors" \
+        iflag=fullblock oflag=dsync conv=notrunc status=noxfer 2>/dev/null
+  else
+    dd of="$file" bs=$bs seek="$offset" count="$sectors" \
+      iflag=fullblock oflag=dsync conv=notrunc status=noxfer 2>/dev/null
+  fi
+}
+
+# Updates a specific partition in given image file (from stdin)
+image_update_partition() {
+  local file="$1"
+  local part_num="$2"
+  local offset="$(image_part_offset "$file" "$part_num")" ||
+    image_die "failed to find partition #$part_num from: $file"
+  local size="$(image_part_size "$file" "$part_num")" ||
+    image_die "failed to find partition #$part_num from: $file"
+
+  image_update_partial_file "$file" "$offset" "$size"
+}
+
 # Maps a specific partition from given image file to a loop device
 image_map_partition() {
   local file="$1"
@@ -173,34 +208,39 @@ image_umount_partition() {
   sudo umount -d "$mount_point"
 }
 
-# Copy a partition from one image to another.
+# Copy a partition from one image to another (size must be equal)
 image_partition_copy() {
-  local src="$1"
-  local srcpart="$2"
-  local dst="$3"
-  local dstpart="$4"
-
-  local srcoffset=$(image_part_offset "${src}" "${srcpart}")
-  local dstoffset=$(image_part_offset "${dst}" "${dstpart}")
-  local length=$(image_part_size "${src}" "${srcpart}")
-  local dstlength=$(image_part_size "${dst}" "${dstpart}")
-
-  if [ "${length}" -gt  "${dstlength}" ]; then
-    exit 1
+  local src="$1" src_part="$2" dst="$3" dst_part="$4"
+  local size1="$(image_part_size "$src" "$src_part")"
+  local size2="$(image_part_size "$dst" "$dst_part")"
+  if [ "$size1" != "$size2" ]; then
+    die "Partition size different: ($size1 != $size2)"
   fi
+  image_dump_partition "$src" "$src_part" 2>/dev/null |
+    image_update_partition "$dst" "$dst_part"
+}
 
-  # Try to use larger buffer if offset/size can be re-aligned.
-  # 2M / 512 = 4096
-  local buffer_ratio=4096
-  local bs=512
-  if [ $((dstoffset % buffer_ratio)) -eq 0 -a \
-       $((length % buffer_ratio)) -eq 0 ]; then
-    dstoffset=$((dstoffset / buffer_ratio))
-    bs=$((bs * buffer_ratio))
+# Copy a partition from one image to another (source <= dest)
+image_partition_overwrite() {
+  local src="$1" src_part="$2" dst="$3" dst_part="$4"
+  local size1="$(image_part_size "$src" "$src_part")"
+  local size2="$(image_part_size "$dst" "$dst_part")"
+  if [ "$size1" -gt "$size2" ]; then
+    die "Destination is too small: ($size1 > $size2)"
   fi
+  image_dump_partition "$src" "$src_part" 2>/dev/null |
+    image_update_partition "$dst" "$dst_part"
+}
 
-  image_dump_partition "${src}" "${srcpart}" |
-      dd of="${dst}" bs="${bs}" seek="${dstoffset}" conv=notrunc oflag=dsync
+# Copy a partition image from file to a full disk image.
+image_partition_copy_from_file() {
+  local src="$1" dst="$2" dst_part="$3"
+  local size1="$(($(stat -c"%s" "$src") / 512))"
+  local size2="$(image_part_size "$dst" "$dst_part")"
+  if [ "$size1" != "$size2" ]; then
+    die "Partition size different: ($size1 != $size2)"
+  fi
+  cat "$src" | image_update_partition "$dst" "$dst_part"
 }
 
 # Temporary object management
diff --git a/make_factory_package.sh b/make_factory_package.sh
index dfeeb5c7d8..88a2e930ec 100755
--- a/make_factory_package.sh
+++ b/make_factory_package.sh
@@ -230,8 +230,8 @@ prepare_img() {
         "$(stat -c %s ${outdev})" != "$(( ${sectors} * 512 ))"  -o \
         "$FLAGS_preserve" = "$FLAGS_FALSE" ]; then
     echo "Generating empty image file"
-    image_dump_partial_file /dev/zero 0 "${sectors}" |
-        dd of="${outdev}" bs=8M
+    truncate -s "0" "$outdev"
+    truncate -s "$((sectors * 512))" "$outdev"
   else
     echo "Reusing $outdev"
   fi
@@ -381,28 +381,23 @@ generate_img() {
   local release_image="${RELEASE_DIR}/${RELEASE_IMAGE}"
   echo "Release Kernel"
   if [ -n "$RELEASE_KERNEL" ]; then
-    local newkernel="$(image_map_partition "${outdev}" "4")"
-    local failed=""
-    sudo dd if="$RELEASE_KERNEL" of="$newkernel" bs=512 ||
-      failed="TRUE"
-    image_unmap_partition "$newkernel"
-    [ -z "$failed" ] || die "Failed to build release kernel."
+    image_partition_copy_from_file "${RELEASE_KERNEL}" "${outdev}" 4
   else
     image_partition_copy "${release_image}" 2 "${outdev}" 4
   fi
   echo "Release Rootfs"
-  image_partition_copy "${release_image}" 3 "${outdev}" 5
+  image_partition_overwrite "${release_image}" 3 "${outdev}" 5
   echo "OEM parition"
-  image_partition_copy "${release_image}" 8 "${outdev}" 8
+  image_partition_overwrite "${release_image}" 8 "${outdev}" 8
 
   # Go to retrieve the factory test image.
   local factory_image="${FACTORY_DIR}/${FACTORY_IMAGE}"
   echo "Factory Kernel"
   image_partition_copy "${factory_image}" 2 "${outdev}" 2
   echo "Factory Rootfs"
-  image_partition_copy "${factory_image}" 3 "${outdev}" 3
+  image_partition_overwrite "${factory_image}" 3 "${outdev}" 3
   echo "Factory Stateful"
-  image_partition_copy "${factory_image}" 1 "${outdev}" 1
+  image_partition_overwrite "${factory_image}" 1 "${outdev}" 1
   echo "EFI Partition"
   image_partition_copy "${factory_image}" 12 "${outdev}" 12
   apply_hwid_updater "${hwid_updater}" "${outdev}"