path: root/scripts/disk_image

#!/bin/bash

# This helps to see what is going on
set -e

PROG_REQ=( kpartx mkfs.ext4 losetup sgdisk readlink )

function check_depends()
{
    for ((i=0; i<${#PROG_REQ[@]}; i++))
    do
        if ! command -v ${PROG_REQ[i]} > /dev/null
        then
            echo "Please install the needed program: ${PROG_REQ[i]}."
            exit 1
        fi
    done
}

function add_partition()
{
    local rootfs="$1"
    _part_size=`stat -L --printf="%s" $rootfs`
    _part_size=$(( $_part_size + $offset - 1))
    _part_size=$(( $_part_size & ~($offset - 1) ))
    # account for gzip compression
    _part_size=$(( $_part_size * 2 ))
    # add some slack
    _part_size=$(( 128*1024*1024 + $_part_size ))
    _part_size=$(( $_part_size / 1024 / 1024 ))
    echo PART size: "$_part_size"MB
    
    prev=$(( $_npart - 1 ))
    _sector_start[$_npart]=$(( ${_sector_end[$prev]} + 1 ))
    _sector_end[$_npart]=$(( $_part_size * 1024 * 1024 / $_sector_size + ${_sector_start[$_npart]} - 1))

    _tot_size=$(( $_tot_size + $_part_size ))
    _npart=$(( $_npart + 1 ))
}

function write_rootfs()
{
    local j=$1
    local rootfs=$2

    # create mount point and mount diskn
    mkdir -p ${DESTDIR}/part/disk$j
    mount -t ext4 /dev/mapper/diskimage$j $DESTDIR/part/disk$j
    
    # Extract rootfs cpio archive into `.../part/vos_$j`
    cd ${DESTDIR}/part/disk$j

    if [[ $rootfs = *.cpio.gz ]]
    then
        cat "${UBOOT_OUT_ABS}/$rootfs" | gunzip | cpio -id
    elif [[ $rootfs = *.tar.gz ]]
    then
        tar xvzf "${UBOOT_OUT_ABS}/$rootfs"
    else
        echo "Ignoring $rootfs: unsupported file format. Use cpio.gz or tar.gz."
    fi

    cd -
    # umount
    sync
    umount $DESTDIR/part/disk$j
}

function print_help
{
    echo "usage:"
    echo "	$0 -c CONFIG_FILE -d UBOOT_DIRECTORY <-w WORK_DIRECTORY> -o IMG_FILE"
    echo "	$0 -h"
    echo "where:"
    echo "	-c CONFIG_FILE - configuration file"
    echo "	-d UBOOT_DIRECTORY - root directory for the paths specified in CONFIG_FILE"
    echo "	-w WORK_DIRECTORY - work directory used when building the image"
    echo "	-o IMG_FILE - the output img file "
    echo "Example:"
    echo "	$0 -c ../config -d ./build42 -w tmp -o disk.img"
}

# before anything else, check if we have root privilege
if ! [ $(id -u) = 0 ]
then
    echo "This script needs root privilege to run, exiting."
    exit 1
fi

while getopts ":w:d:c:o:h" opt
do
    case ${opt} in
    w )
        DESTDIR=$OPTARG
        ;;
    d )
        UBOOT_OUT=$OPTARG
        ;;
    c )
        CFG_FILE=$OPTARG
        ;;
    o )
        IMG=$OPTARG
        ;;
    h )
        print_help
        exit 0
        ;;
    * )
        echo "Unknown option, see \"$0 -h\""
        exit 1
        ;;
    esac
done
shift $((OPTIND -1))

if [ -z "$UBOOT_OUT" ] || [ -z "$CFG_FILE" ] || [ -z "$IMG" ]
then
    echo "Undefined arguments, see \"$0 -h\""
    exit 1
fi

# if the user hasn't specified a working directing, create it
if [ -z "$DESTDIR" ]
then
    DESTDIR="$(mktemp -d /tmp/imagebuilder.XXXXXX)"
    DESTDIR_DEL=true
else
    DESTDIR_DEL=false
fi

UBOOT_OUT_ABS="$(readlink -f $UBOOT_OUT)"
DESTDIR_ABS="$(readlink -f $DESTDIR)"

check_depends

source "$CFG_FILE"

offset=$((2*1024*1024))
_part1_size=`stat -L --printf="%s" $UBOOT_OUT/$XEN`
_part1_size=$(( $_part1_size + `stat -L --printf="%s" $UBOOT_OUT/$DOM0_KERNEL` ))
_part1_size=$(( $_part1_size + `stat -L --printf="%s" $UBOOT_OUT/$DEVICE_TREE` ))
_part1_size=$(( $_part1_size + `stat -L --printf="%s" $UBOOT_OUT/$UBOOT_SOURCE` ))
_part1_size=$(( $_part1_size + `stat -L --printf="%s" $UBOOT_OUT/$UBOOT_SCRIPT` ))
i=0
while test $i -lt $NUM_DOMUS
do
    _part1_size=$(( $_part1_size + `stat -L --printf="%s" $UBOOT_OUT/${DOMU_KERNEL[$i]}` ))
    if test "${DOMU_RAMDISK[$i]}"
    then
        _part1_size=$(( $_part1_size + `stat -L --printf="%s" $UBOOT_OUT/${DOMU_RAMDISK[$i]}` ))
    fi
    if test "${DOMU_PASSTHROUGH_DTB[$i]}"
    then
        _part1_size=$(( $_part1_size + `stat -L --printf="%s" $UBOOT_OUT/${DOMU_PASSTHROUGH_DTB[$i]}` ))
    fi
    i=$(( $i + 1 ))
done

# add 16 MB slack
_part1_size=$(( $_part1_size + 16777216 ))
_part1_size=$(( $_part1_size + $offset - 1))
_part1_size=$(( $_part1_size & ~($offset - 1) ))
_part1_size=$(( $_part1_size / 1024 / 1024 ))
echo PART1 size: "$_part1_size"MB

_sector_size=512
# _sector_start[0] needs to be aligned to 2048
_sector_start[0]=2048
_sector_end[0]=$(( $_part1_size * 1024 * 1024 / $_sector_size + ${_sector_start[0]} - 1))
_tot_size=$(( $_part1_size ))
_npart=1

if test "$DOM0_ROOTFS"
then
    add_partition "$UBOOT_OUT/$DOM0_ROOTFS"
fi

i=0
while test $i -lt $NUM_DOMUS
do
    if test "${DOMU_ROOTFS[$i]}"
    then
        add_partition "$UBOOT_OUT/${DOMU_ROOTFS[$i]}"
    fi
    i=$(( $i + 1 ))
done
_tot_size=$(( $_tot_size + 16 ))

# NOTE: Increase vos_a to 256 to accomodate rootfs
# 528 MiB (256 + 256 + 16)
truncate $IMG -s "$_tot_size"M

# create GPT partition table
sgdisk -og $IMG

i=0
j=1
while test $i -lt $_npart
do
    sgdisk -n $j:${_sector_start[$i]}:${_sector_end[$i]} -c $j:"Linux""$j" -t $j:8300 $IMG
    i=$(( $i + 1 ))
    j=$(( $j + 1 ))
done

# find the first available loop device
_loop_dev=$(losetup -f)

# attach loopback device to $IMG
losetup $_loop_dev $IMG

_disksize=$(blockdev --getsize $_loop_dev)

dmsetup create diskimage --table "0 $_disksize linear $_loop_dev 0"

# ensure that /dev/mapper/diskimage exists
while [ ! -b /dev/mapper/diskimage ]
do
    sleep 2
done

kpartx -a /dev/mapper/diskimage

# ensure that /dev/mapperdiskimage1 exists
while [ ! -b /dev/mapper/diskimage1 ]
do
    sleep 2
done

i=0
j=1
while test $i -lt $_npart
do
    mkfs.ext4 -L vos_$j -F /dev/mapper/diskimage$j
    i=$(( $i + 1 ))
    j=$(( $j + 1 ))
done

# create mount point and mount disk1
mkdir -p ${DESTDIR}/part/disk1
mount -t ext4 /dev/mapper/diskimage1 $DESTDIR/part/disk1

# only copy over files that were counted for the partition size
cd "$UBOOT_OUT"
cp --parents "$XEN" "${DESTDIR_ABS}/part/disk1/"
cp --parents "$DOM0_KERNEL" "${DESTDIR_ABS}/part/disk1/"
cp --parents "$DEVICE_TREE" "${DESTDIR_ABS}/part/disk1/"
cp --parents "$UBOOT_SOURCE" "${DESTDIR_ABS}/part/disk1/"
cp --parents "$UBOOT_SCRIPT" "${DESTDIR_ABS}/part/disk1/"
i=0
while test $i -lt $NUM_DOMUS
do
    cp --parents "${DOMU_KERNEL[$i]}" "${DESTDIR_ABS}/part/disk1/"
    if test "${DOMU_RAMDISK[$i]}"
    then
        cp --parents "${DOMU_RAMDISK[$i]}" "${DESTDIR_ABS}/part/disk1/"
    fi
    if test "${DOMU_PASSTHROUGH_DTB[$i]}"
    then
        cp --parents "${DOMU_PASSTHROUGH_DTB[$i]}" "${DESTDIR_ABS}/part/disk1/"
    fi
    i=$(( $i + 1 ))
done
cd -

# unmount
sync
# This fails for some reason. It could work now because we are not using qemu-user
# fstrim $DESTDIR/part/disk1
umount $DESTDIR/part/disk1

j=2
if test "$DOM0_ROOTFS"
then
    write_rootfs 2 "$DOM0_ROOTFS"
    j=$(( $j + 1 ))
fi

i=0
while test $i -lt $NUM_DOMUS
do
    if test "${DOMU_ROOTFS[$i]}"
    then
        write_rootfs $j "${DOMU_ROOTFS[$i]}"
        j=$(( $j + 1 ))
    fi
    i=$(( $i + 1 ))
done


kpartx -d /dev/mapper/diskimage
dmsetup remove diskimage
losetup -d $_loop_dev

if [ "$DESTDIR_DEL" = true ]
then
    rm -rf "$DESTDIR"
fi