#! /bin/bash

# This is a unit test program for kpartx, in particular for deleting partitions.
#
# The rationale is the following:
#
#  1) kpartx should delete all mappings it created beforehand.
#  2) kpartx should handle partitions on dm devices and other devices
#     (e.g. loop devices) equally well.
#  3) kpartx should only delete "partitions", which are single-target
#     linear mappings into a block device. Other maps should not be touched.
#  4) kpartx should only delete mappings it created itself beforehand.
#     In particular, it shouldn't delete LVM LVs, even if they are fully
#     contained in the block device at hand and thus look like partitions
#     in the first place. (For historical compatibility reasons, we allow
#     such mappings to be deleted with the -f/--force flag).
#  5) DM map names may be changed, thus kpartx shouldn't rely on them to
#     check whether a mapping is a partition of a particular device. It is
#     legal for a partition of /dev/loop0 to be named "loop0".

# Note: This program tries hard to clean up, but if tests fail,
# stale DM or loop devices may keep lurking around.

# Set WORKDIR in environment to existing dir to for persistence
# WARNING:  existing files will be truncated.
# If empty, test will be done in temporary dir
: ${WORKDIR:=}
# Set this environment variable to test an alternative kpartx executable
: ${KPARTX:=}
# Options to pass to kpartx always
: ${KPARTX_OPTS:=-s}
# Time to wait for device nodes to appear (microseconds)
# Waiting is only needed if "s" is not in $KPARTX_OPTS
: ${WAIT_US:=0}

# IMPORTANT: The ERR trap is essential for this program to work correctly!
#trap 'LINE=$LINENO; trap - ERR; echo "== error in $BASH_COMMAND on line $LINE ==" >&2; ls -la /dev/mapper; read a; exit 1' ERR
trap 'cleanup' 0

CLEANUP=:
cleanup() {
    trap - ERR
    trap - 0
    if [[ $OK ]]; then
        echo == all ${ALL_PASS} tests completed successfully == >&2
    else
        printf "== Final Results ==\n%s of %s failed!\n\n" "${ALL_FAIL}" "$(( ALL_PASS + ALL_FAIL ))" >&2
    fi
    eval "$CLEANUP" &>/dev/null
}

push_cleanup() {
    CLEANUP="$@;$CLEANUP"
}

pop_cleanup() {
    # CAUTION: simplistic
    CLEANUP=${CLEANUP#*;}
}

step() {
    STEP="$@"

    if [ $PASS -gt 0 -o $FAIL -gt 0 ]
    then
      printf "PASS: %s FAIL: %s\n\n" "${PASS}" "${FAIL}"
      PASS=0
      FAIL=0
    fi

    echo == Test step: $STEP == >&2
}

mk_partitions() {
    parted -s $1 mklabel msdos
    parted -s -- $1 mkpart prim ext2 1MiB -1s
}

wipe_ptable() {
    dd if=/dev/zero of=$1 bs=1b count=1
}

usleep() {
    if which usleep >/dev/null 2>&1
    then
        env usleep "${1}"
    else
        sleep $( echo "${1}" / 1000000.0 | bc )
    fi
}

invoke_kpartx() {
    $KPARTX $KPARTX_OPTS "${@}"
    usleep "${WAIT_US}"
}

assert_exists() {
    local result=0
    while [ "${1}" ]
    do
        if [ -b "$(readlink -f "${1}")" ]
        then
            PASS=$(( PASS + 1 ))
            ALL_PASS=$(( ALL_PASS + 1 ))
        else
            FAIL=$(( FAIL + 1 ))
            ALL_FAIL=$(( ALL_FAIL + 1 ))
            echo "assert_exists: Device does not exist when it should: ${1}"
            result=1
        fi
        shift
    done
    return $result
}

assert_unexists() {
    local result=0
    while [ "${1}" ]
    do
        if [ ! -b "$(readlink -f "${1}")" ]
        then
            PASS=$(( PASS + 1 ))
            ALL_PASS=$(( ALL_PASS + 1 ))
        else
            FAIL=$(( FAIL + 1 ))
            ALL_FAIL=$(( ALL_FAIL + 1 ))
            echo "assert_exists: Device exists when it should not: ${1}"
            result=1
        fi
        shift
    done
    return $result
}

PASS=0
ALL_PASS=0
FAIL=0
ALL_FAIL=0

step preparation

[[ $UID -eq 0 ]]
[[ $KPARTX ]] || {
    if [[ -x $PWD/kpartx/kpartx ]]; then
	KPARTX=$PWD/kpartx/kpartx
    else
	KPARTX=$(which kpartx)
    fi
}
[[ $KPARTX ]]

FILE1=kpartx1
FILE2=kpartx2
FILE3=kpartx3
FILE4=kpartx4

SIZE=$((1024*1024*1024))  # use bytes as units here
SECTSIZ=512
OFFS=32                # offset of linear mapping into dev, sectors
VG=kpvg  # volume group name
LV=kplv  # logical vol name
LVMCONF='devices { filter = [ "a|/dev/loop.*|", r".*" ] }'

OK=

[[ $WORKDIR ]] || {
    WORKDIR=$(mktemp -d /tmp/kpartx-XXXXXX)
    push_cleanup 'rm -rf $WORKDIR'
}

push_cleanup "cd $PWD"
cd "$WORKDIR"

step "create loop devices"
truncate -s $SIZE $FILE1
truncate -s $SIZE $FILE2
truncate -s $SIZE $FILE3
truncate -s $SIZE $FILE4

LO1=$(losetup -f $FILE1 --show)
push_cleanup 'losetup -d $LO1'
LO2=$(losetup -f $FILE2 --show)
push_cleanup 'losetup -d $LO2'
LO3=$(losetup -f $FILE3 --show)
push_cleanup 'losetup -d $LO3'
LO4=$(losetup -f $FILE4 --show)
push_cleanup 'losetup -d $LO4'

[[ $LO1 && $LO2 && $LO3 && $LO4 && -b $LO1 && -b $LO2 && -b $LO3 && -b $LO4 ]]
DEV1=$(stat -c "%t:%T" $LO1)
DEV2=$(stat -c "%t:%T" $LO2)
DEV3=$(stat -c "%t:%T" $LO3)

usleep $WAIT_US

step "create DM devices (spans)"
# Create two linear mappings spanning two loopdevs.
# One of them gets a pathological name colliding with
# the loop device name.
# These mappings must not be removed by kpartx.
# They also serve as DM devices to test partition removal on those.

TABLE="\
0 $((SIZE/SECTSIZ-OFFS)) linear $DEV1 $OFFS
$((SIZE/SECTSIZ-OFFS)) $((SIZE/SECTSIZ-OFFS)) linear $DEV2 $OFFS"

SPAN1=kpt
SPAN2=$(basename $LO2)
dmsetup create $SPAN1 <<<"$TABLE"
push_cleanup 'dmsetup remove -f $SPAN1'

dmsetup create $SPAN2 <<<"$TABLE"
push_cleanup 'dmsetup remove -f $SPAN2'

# This is a non-kpartx pseudo "partition" mapping
USER1=user1
push_cleanup 'dmsetup remove -f $USER1'
dmsetup create $USER1 <<EOF
0 $((SIZE/SECTSIZ-OFFS)) linear $DEV1 $OFFS
EOF

usleep $WAIT_US
assert_exists /dev/mapper/$SPAN1 /dev/mapper/$SPAN2 /dev/mapper/$USER1
[[ -b /dev/mapper/$SPAN1 ]]
[[ -b /dev/mapper/$SPAN2 ]]
[[ -b /dev/mapper/$USER1 ]]

step "create vg on $LO3"
# On the 3rd loop device, we create a VG and an LV
# The LV should not be removed by kpartx.
pvcreate --config "$LVMCONF" -f $LO3
vgcreate --config "$LVMCONF" $VG $LO3
push_cleanup 'vgremove --config "$LVMCONF" -f $VG'
lvcreate --config "$LVMCONF" -L $((SIZE/2))B -n $LV $VG
push_cleanup 'lvremove --config "$LVMCONF" -f $VG/$LV'
usleep $WAIT_US

assert_exists /dev/mapper/$VG-$LV
[[ -b /dev/mapper/$VG-$LV ]]

# dmsetup table /dev/mapper/$VG-$LV
# dmsetup info /dev/mapper/$VG-$LV

step "create partitions on loop devices"

mk_partitions $LO1
mk_partitions $LO2
mk_partitions $LO4

# Test invocation of kpartx with regular file here
LO2P1=/dev/mapper/$(basename $LO2)-foo1
invoke_kpartx -a -p -foo $FILE2
assert_exists $LO2P1
[[ -b $LO2P1 ]]
push_cleanup 'dmsetup remove -f $(basename $LO2P1)'

step "remove partitions with deleted ptable"
wipe_ptable $LO2
invoke_kpartx -d $LO2
assert_unexists $LO2P1

mk_partitions $LO2
invoke_kpartx -a -p -foo $FILE2
assert_exists $LO2P1

LO1P1=/dev/mapper/$(basename $LO1)-eggs1
invoke_kpartx -a -p -eggs $LO1
push_cleanup 'dmsetup remove -f $(basename $LO1P1)'

usleep $WAIT_US
assert_exists $LO1P1 $LO2P1

# dmsetup info $LO2P1

# Set pathological name for partition on $LO1 (same as loop device itself)
dmsetup rename $(basename $LO1P1) $(basename $LO1)
LO1P1=/dev/mapper/$(basename $LO1)
pop_cleanup
push_cleanup 'dmsetup remove -f $(basename $LO1P1)'

# dmsetup info $LO1P1

step "create partitions on DM devices"
mk_partitions /dev/mapper/$SPAN2

invoke_kpartx -a -p -bar /dev/mapper/$SPAN2
SPAN2P1=/dev/mapper/${SPAN2}-bar1

# udev rules may have created partition mappings without UUIDs
# which aren't removed by default (if system standard kpartx doesn't
# set the UUID). Remove them using -f
push_cleanup 'invoke_kpartx -f -d /dev/mapper/$SPAN2'
push_cleanup 'dmsetup remove -f $(basename $SPAN2P1)'

invoke_kpartx -a -p -spam /dev/mapper/$SPAN1
SPAN1P1=/dev/mapper/${SPAN1}-spam1
# see above
push_cleanup 'invoke_kpartx -f -d /dev/mapper/$SPAN1'
push_cleanup 'dmsetup remove -f $(basename $SPAN1P1)'

usleep $WAIT_US
assert_exists $SPAN2P1 $SPAN1P1

step "rename partitions on DM device to default"
invoke_kpartx -u /dev/mapper/$SPAN1
assert_unexists $SPAN1P1
# This assumes that $SPAN1 ends in a non-digit
assert_exists ${SPAN1P1//-spam/}

step "rename partitions on DM device back from default"
invoke_kpartx -u -p -spam /dev/mapper/$SPAN1
assert_exists $SPAN1P1
assert_unexists ${SPAN1P1//-spam/}

step "delete partitions on DM devices"
invoke_kpartx -d /dev/mapper/$SPAN1 >&2
usleep $WAIT_US

assert_exists $SPAN2P1 $LO1P1 $LO2P1
assert_unexists $SPAN1P1

invoke_kpartx -d /dev/mapper/$SPAN2
usleep $WAIT_US

assert_exists $LO1P1 $LO2P1
assert_unexists $SPAN2P1

step "rename partitions on loop device"
invoke_kpartx -u -p -spam $LO2
assert_unexists $LO2P1
assert_exists ${LO2P1//-foo/-spam}

step "rename partitions on loop device back"
invoke_kpartx -u -p -foo $LO2
assert_exists $LO2P1
assert_unexists ${LO2P1//-foo/-spam}

step "rename partitions on loop device to default"
invoke_kpartx -u $LO2
#read a
assert_unexists $LO2P1
assert_exists ${LO2P1//-foo/p}
# $LO1 ends in a digit

step "rename partitions on loop device back from default"
invoke_kpartx -u -p -foo $LO2
assert_exists $LO2P1
assert_unexists ${LO2P1//-foo/p}

step "rename partitions on loop devices"
invoke_kpartx -u -p spam $LO2

step "delete partitions on loop devices"

invoke_kpartx -d $LO3

# This will also delete the loop device
invoke_kpartx -d $FILE2
invoke_kpartx -d $LO1
usleep $WAIT_US

# ls -l /dev/mapper
assert_unexists $LO1P1
pop_cleanup
assert_unexists $LO2P1
pop_cleanup
# spans should not have been removed
# LVs neither
assert_exists /dev/mapper/$SPAN1 /dev/mapper/$SPAN2 /dev/mapper/$USER1 /dev/mapper/$VG-$LV

step "delete partitions on $LO3 with -f"

invoke_kpartx -f -d $LO3
# -d -f should delete the LV, too
assert_unexists /dev/mapper/$VG-$LV
assert_exists /dev/mapper/$SPAN1 /dev/mapper/$SPAN2

step "test kpartx creation/deletion on an image file with no existing loopdev"
losetup -d $LO4

OUTPUT=$(invoke_kpartx -v -a $FILE4 2>&1)
read loop dm < \
     <(sed -n  's/^add map \(loop[0-9]*\)p1 ([0-9]*:\([0-9]*\)).*$/\1 dm-\2/p' \
	   <<<$OUTPUT)
[[ $dm && $loop ]]
push_cleanup "dmsetup remove -f /dev/$dm"
push_cleanup "losetup -d /dev/$loop"

assert_exists /dev/mapper/${loop}p1
invoke_kpartx -d $FILE4
assert_unexists /dev/mapper/${loop}p1
# /dev/$loop is _not_ automatically deleted
assert_exists /dev/${loop}

[ $ALL_FAIL -eq 0 ] && OK=yes