#!/bin/bash

[ -z "$CRYPTSETUP_PATH" ] && CRYPTSETUP_PATH=".."
CRYPTSETUP=$CRYPTSETUP_PATH/cryptsetup
DEV=""
DEV_STACKED="luks0xbabe"
DEV_NAME="dummyalign"
MNT_DIR="./mnt_luks"
PWD1="93R4P4pIqAH8"
PWD2="mymJeD8ivEhE"
FAST_PBKDF="--pbkdf pbkdf2 --pbkdf-force-iterations 1000"

cleanup() {
	udevadm settle >/dev/null 2>&1
	if [ -d "$MNT_DIR" ] ; then
	    umount -f $MNT_DIR 2>/dev/null
	    rmdir $MNT_DIR 2>/dev/null
	fi
	[ -b /dev/mapper/$DEV_STACKED ] && dmsetup remove --retry $DEV_STACKED >/dev/null 2>&1
	[ -b /dev/mapper/$DEV_NAME ] && dmsetup remove --retry $DEV_NAME >/dev/null 2>&1
	# FIXME scsi_debug sometimes in-use here
	sleep 1
	rmmod scsi_debug 2>/dev/null
	sleep 1
}

fail()
{
	if [ -n "$1" ] ; then echo "FAIL $1" ; fi
	echo "FAILED backtrace:"
	while caller $frame; do ((frame++)); done
	cleanup
	exit 100
}

skip()
{
	echo "TEST SKIPPED: $1"
	cleanup
	exit 0
}

function dm_crypt_features()
{
	VER_STR=$(dmsetup targets | grep crypt | cut -f2 -dv)
	[ -z "$VER_STR" ] && fail "Failed to parse dm-crypt version."

	VER_MAJ=$(echo $VER_STR | cut -f 1 -d.)
	VER_MIN=$(echo $VER_STR | cut -f 2 -d.)
	VER_PTC=$(echo $VER_STR | cut -f 3 -d.)

	[ $VER_MAJ -lt 1 ] && return
	[ $VER_MAJ -gt 1 ] && {
		DM_PERF_CPU=1
		DM_SECTOR_SIZE=1
		return
	}

	[ $VER_MIN -lt 14 ] && return
	DM_PERF_CPU=1
	if [ $VER_MIN -ge 17 -o \( $VER_MIN -eq 14 -a $VER_PTC -ge 5 \) ]; then
		DM_SECTOR_SIZE=1
	fi
}

add_device() {
	modprobe scsi_debug $@ delay=0
	if [ $? -ne 0 ] ; then
		echo "This kernel seems to not support proper scsi_debug module, test skipped."
		exit 77
	fi

	sleep 2
	DEV=$(grep -l -e scsi_debug /sys/block/*/device/model | cut -f4 -d /)

	if [ ! -e /sys/block/$DEV/alignment_offset ] ; then
		echo "This kernel seems to not support topology info, test skipped."
		cleanup
		exit 77
	fi

	DEV="/dev/$DEV"
	[ -b $DEV ] || fail "Cannot find $DEV."
}

format() # expected [forced] [encryption_sector_size]
{
	local _sec_size=512

	local _exp=$1

	if [ "${2:0:1}" = "s" ]; then
		_sec_size=${2:1}
		shift
	fi

	test "${3:0:1}" = "s" && _sec_size=${3:1}

	test $_sec_size -eq 512 || local _smsg=" (encryption sector size $_sec_size)"

	if [ -z "$2" ] ; then
		echo -n "Formatting using topology info$_smsg..."
		echo $PWD1 | $CRYPTSETUP luksFormat $FAST_PBKDF --type luks2 $DEV -q -c aes-cbc-essiv:sha256 --sector-size $_sec_size >/dev/null || fail
	else
		echo -n "Formatting using forced sector alignment $2$_smsg..."
		echo $PWD1 | $CRYPTSETUP luksFormat $FAST_PBKDF --type luks2 $DEV -q -c aes-cbc-essiv:sha256 --align-payload=$2 --sector-size $_sec_size >/dev/null || fail
	fi

	# check the device can be activated
	if [ -n "$DM_SECTOR_SIZE" ] ; then
		echo $PWD1 | $CRYPTSETUP luksOpen $DEV $DEV_NAME || fail
		$CRYPTSETUP close $DEV_NAME || fail
	fi

	ALIGN=$($CRYPTSETUP luksDump $DEV | tee  /tmp/last_dump | grep -A1 "0: crypt" | grep "offset:" | cut -d ' ' -f2)
	# echo "ALIGN = $ALIGN"

	[ -z "$ALIGN" ] && fail
	ALIGN=$((ALIGN/512))
	[ $ALIGN -ne $_exp ] && fail "Expected alignment differs: expected $_exp != detected $ALIGN"

	# test some operation, just in case
	echo -e "$PWD1\n$PWD2" | $CRYPTSETUP luksAddKey $DEV $FAST_PBKDF --key-slot 1
	[ $? -ne 0 ] && fail "Keyslot add failed."

	$CRYPTSETUP -q luksKillSlot $DEV 1
	[ $? -ne 0 ] && fail "Keyslot removal failed."

	echo "PASSED"
}

format_fail() # expected [forced] [encryption_sector_size]
{
	local _sec_size=512

	local _exp=$1

	if [ "${2:0:1}" = "s" ]; then
		_sec_size=${2:1}
		shift
	fi

	test "${3:0:1}" = "s" && _sec_size=${3:1}

	test $_sec_size -eq 512 || local _smsg=" (encryption sector size $_sec_size)"

	if [ -z "$2" ] ; then
		echo -n "Formatting using topology info$_smsg (must fail)..."
		echo $PWD1 | $CRYPTSETUP luksFormat $FAST_PBKDF --type luks2 $DEV -q -c aes-cbc-essiv:sha256 --sector-size $_sec_size >/dev/null 2>&1 && fail
	else
		echo -n "Formatting using forced sector alignment $2$_smsg (must fail)..."
		echo $PWD1 | $CRYPTSETUP luksFormat $FAST_PBKDF --type luks2 $DEV -q -c aes-cbc-essiv:sha256 --align-payload=$2 --sector-size $_sec_size >/dev/null 2>&1 && fail
	fi

	echo "PASSED"
}

if [ $(id -u) != 0 ]; then
	echo "WARNING: You must be root to run this test, test skipped."
	exit 77
fi

dm_crypt_features
modprobe --dry-run scsi_debug || exit 77
cleanup

add_device dev_size_mb=32
echo $PWD1 | $CRYPTSETUP luksFormat $FAST_PBKDF --type luks2 $DEV -q  >/dev/null || fail
EXPCT=$($CRYPTSETUP luksDump $DEV | grep "offset: " | cut -f 2 -d ' ')
test "$EXPCT" -gt 512 || fail
EXPCT=$((EXPCT/512))
echo "Default alignment detected: $EXPCT sectors"
cleanup

echo "# Create desktop-class 4K drive"
echo "# (logical_block_size=512, physical_block_size=4096, alignment_offset=0)"
add_device dev_size_mb=32 sector_size=512 physblk_exp=3 num_tgts=1
format $EXPCT
format $EXPCT s1024
format $EXPCT s2048
format $EXPCT s4096
format $EXPCT 1
format $EXPCT 1 s1024
format $EXPCT 1 s2048
format $EXPCT 1 s4096
format $EXPCT 8
format $EXPCT 8 s1024
format $EXPCT 8 s2048
format $EXPCT 8 s4096
format $((EXPCT+1)) $((EXPCT+1))
format_fail $((EXPCT+1)) $((EXPCT+1)) s1024
format_fail $((EXPCT+1)) $((EXPCT+1)) s2048
format_fail $((EXPCT+1)) $((EXPCT+1)) s4096
format $EXPCT $EXPCT
format $EXPCT $EXPCT s1024
format $EXPCT $EXPCT s2048
format $EXPCT $EXPCT s4096
cleanup

echo "# Create desktop-class 4K drive with misaligned opt-io (some bad USB enclosures)"
echo "# (logical_block_size=512, physical_block_size=4096, alignment_offset=0, opt-io=1025)"
add_device dev_size_mb=32 sector_size=512 physblk_exp=3 num_tgts=1 opt_blks=1025
format $EXPCT
format $EXPCT s1024
format $EXPCT s2048
format $EXPCT s4096
format $EXPCT 1
format $EXPCT 1 s1024
format $EXPCT 1 s2048
format $EXPCT 1 s4096
format $EXPCT 8
format $EXPCT 8 s1024
format $EXPCT 8 s2048
format $EXPCT 8 s4096
format $((EXPCT+1)) $((EXPCT+1))
format_fail $((EXPCT+1)) $((EXPCT+1)) s1024
format_fail $((EXPCT+1)) $((EXPCT+1)) s2048
format_fail $((EXPCT+1)) $((EXPCT+1)) s4096
format $EXPCT $EXPCT
format $EXPCT $EXPCT s1024
format $EXPCT $EXPCT s2048
format $EXPCT $EXPCT s4096
cleanup

echo "# Create desktop-class 4K drive w/ 1-sector shift (original bug report)"
echo "# (logical_block_size=512, physical_block_size=4096, alignment_offset=512)"
add_device dev_size_mb=32 sector_size=512 physblk_exp=3 lowest_aligned=1 num_tgts=1
format $((EXPCT+1))
format_fail $((EXPCT+1)) s1024
format_fail $((EXPCT+1)) s2048
format_fail $((EXPCT+1)) s4096
format $EXPCT 1
format $EXPCT 1 s1024
format $EXPCT 1 s2048
format $EXPCT 1 s4096
format $EXPCT 8
format $EXPCT 8 s1024
format $EXPCT 8 s2048
format $EXPCT 8 s4096
format $((EXPCT+1)) $((EXPCT+1))
format_fail $((EXPCT+1)) $((EXPCT+1)) s1024
format_fail $((EXPCT+1)) $((EXPCT+1)) s2048
format_fail $((EXPCT+1)) $((EXPCT+1)) s4096
format $EXPCT $EXPCT
format $EXPCT $EXPCT s1024
format $EXPCT $EXPCT s2048
format $EXPCT $EXPCT s4096
cleanup

echo "# Create desktop-class 4K drive w/ 63-sector DOS partition compensation"
echo "# (logical_block_size=512, physical_block_size=4096, alignment_offset=3584)"
add_device dev_size_mb=32 sector_size=512 physblk_exp=3 lowest_aligned=7 num_tgts=1
format $((EXPCT+7))
format_fail $((EXPCT+7)) s1024
format_fail $((EXPCT+7)) s2048
format_fail $((EXPCT+7)) s4096
format $EXPCT 1
format $EXPCT 1 s1024
format $EXPCT 1 s2048
format $EXPCT 1 s4096
format $EXPCT 8
format $EXPCT 8 s1024
format $EXPCT 8 s2048
format $EXPCT 8 s4096
format $((EXPCT+1)) $((EXPCT+1))
format_fail $((EXPCT+1)) $((EXPCT+1)) s1024
format_fail $((EXPCT+1)) $((EXPCT+1)) s2048
format_fail $((EXPCT+1)) $((EXPCT+1)) s4096
format $EXPCT $EXPCT
format $EXPCT $EXPCT s1024
format $EXPCT $EXPCT s2048
format $EXPCT $EXPCT s4096
cleanup

echo "# Create enterprise-class 4K drive"
echo "# (logical_block_size=4096, physical_block_size=4096, alignment_offset=0)"
add_device dev_size_mb=32 sector_size=4096 num_tgts=1 opt_blks=64
format $EXPCT
format $EXPCT s1024
format $EXPCT s2048
format $EXPCT s4096
format $EXPCT 1
format $EXPCT 1 s1024
format $EXPCT 1 s2048
format $EXPCT 1 s4096
format $EXPCT 8
format $EXPCT 8 s1024
format $EXPCT 8 s2048
format $EXPCT 8 s4096
#FIXME: kernel limits issue?
##format $((EXPCT+1)) $((EXPCT+1))
format_fail $((EXPCT+1)) $((EXPCT+1)) s1024
format_fail $((EXPCT+1)) $((EXPCT+1)) s2048
format_fail $((EXPCT+1)) $((EXPCT+1)) s4096
format $EXPCT $EXPCT
format $EXPCT $EXPCT s1024
format $EXPCT $EXPCT s2048
format $EXPCT $EXPCT s4096
cleanup

echo "# Create classic 512B drive and stack dm-linear"
echo "# (logical_block_size=512, physical_block_size=512, alignment_offset=0)"
add_device dev_size_mb=32 sector_size=512 num_tgts=1
DEV2=$DEV
DEV=/dev/mapper/$DEV_STACKED
dmsetup create $DEV_STACKED --table "0 65536 linear $DEV2 0"
format $EXPCT
format $EXPCT s1024
format $EXPCT s2048
format $EXPCT s4096
format $EXPCT 1
format $EXPCT 1 s1024
format $EXPCT 1 s2048
format $EXPCT 1 s4096
format $EXPCT 8
format $EXPCT 8 s1024
format $EXPCT 8 s2048
format $EXPCT 8 s4096
format $((EXPCT+1)) $((EXPCT+1))
format_fail $((EXPCT+1)) $((EXPCT+1)) s1024
format_fail $((EXPCT+1)) $((EXPCT+1)) s2048
format_fail $((EXPCT+1)) $((EXPCT+1)) s4096
format $EXPCT $EXPCT
format $EXPCT $EXPCT s1024
format $EXPCT $EXPCT s2048
format $EXPCT $EXPCT s4096
cleanup

echo "# Create enterprise-class 4K drive with fs and LUKS images."
# loop device here presents 512 block but images have 4k block
# cryptsetup should properly use 4k block on direct-io
add_device dev_size_mb=32 sector_size=4096 physblk_exp=0 num_tgts=1 opt_blks=64
for file in $(ls img_fs_*.img.xz) ; do
    echo "Format using fs image $file."
    xz -d -c $file | dd of=$DEV bs=1M 2>/dev/null || fail "bad image"
    [ ! -d $MNT_DIR ] && mkdir $MNT_DIR
    mount $DEV $MNT_DIR || skip "Mounting image is not available."
    echo $PWD1 | $CRYPTSETUP luksFormat $FAST_PBKDF --type luks2 $MNT_DIR/luks.img --offset 8192 || fail
    echo $PWD2 | $CRYPTSETUP luksFormat $FAST_PBKDF --type luks2 $MNT_DIR/luks.img --header $MNT_DIR/luks_header.img || fail
    umount $MNT_DIR
done
cleanup