#!/bin/sh

udevd_timeout=10
dev=$1
osize=$2

# Check arguments count
if [ "$#" -lt 1 -o "$#" -gt 2 -o -z "$1" ]; then
	echo "usage: $0 <device> [<xGB>|<xTB>|<sectors>]"
	exit 1
fi

# Make sure provided disk is valid
if ! grep -q " ${dev}$" /proc/partitions; then
	echo "The disk is specified incorrectly, aborting."
	exit 1
fi

# Sanitize/convert size argument.
size=${osize}
if [ -n "${osize}" ]; then
	size=`echo "${osize}" | awk '/^[0-9]*$/'`
	gsize=`echo "${osize}" | awk -F '[GT]' '/^[0-9]*TB?$/ { print $1*1000 } /^[0-9]*GB?$/ { print $1 }'`
	if [ -n "${gsize}" ]; then
		# Supplied size in gigabytes or terabytes convert to bytes
		if [ ${gsize} -lt 8000 ]; then
			# For small disks fancy math is rounded to 4KB.
			bsize=$((((${gsize} * 1000194048 + 10838016) + 4095) & ~4095))
		else
			# For large disks GB size rounded to next TiB.
			bsize=$(((${gsize} * 1000000000 + ((1 << 30) - 1)) & ~((1 << 30) - 1)))
		fi
		# Convert bytes into sectors.
		sector=`lsblk -rndo log-sec /dev/${dev}`
		size=$((${bsize} / ${sector}))
	elif [ -n "${size}" ]; then
		# Supplied size in sectors use as-is.
		:
	else
		# Size not in sectors or gigabytes.  Fail.
		echo "Incorrect size '${osize}'.  Specify in TB, GB or sectors."
		exit 1
	fi
	echo "Resizing ${dev} to ${size} sectors"
else
	echo "Resizing ${dev} to full capacity"
fi

# ATA, SCSI and NVMe have completely different resize commands.
# SAT we handle as ATA via pass-through, but with some complications.
devtype=`lsblk -rndo tran /dev/${dev}`
if [ "${devtype}" = "sas" -a -e "/sys/block/${dev}/device/vpd_pg89" ]; then
	devtype="sata"
fi
case "${devtype}" in
ata|sata)
	# For SATL disks we first need to try to enable descriptor format
	# sense data.  Without it some SATL devices return fixed format,
	# that provides not enough response data for some ATA commands.
	# All ATA/SATA disks go through some SATL in Linux, be it libata
	# or some other hardware/driver.
	echo "Enabling descriptor sense"
	field=`sginfo -C /dev/${dev} | awk '/---/ { start = NR } /D_SENSE/ { print NR - start }'`
	sginfo -CX /dev/${dev} | awk "\$$field = 1" | xargs sginfo -CXNR /dev/${dev}

	# Fetch the ATA disk capabilities, since there are many options.
	ident=`hdparm -I /dev/${dev}`
	ama=`echo "${ident}" | grep -Fc "unknown 119[8]"`
	hpa=`echo "${ident}" | grep -Fc "Host Protected Area feature set"`
	if [ ${ama} -eq 0 -a ${hpa} -eq 0 ]; then
		echo "ATA device supports neither HPA nor AMA, can't resize"
		exit 1
	fi
	ssd=`echo "${ident}" | grep -Fc "Nominal Media Rotation Rate: Solid State Device"`
	block=`echo "${ident}" | grep -Fc "*	BLOCK_ERASE_EXT command"`
	crypto=`echo "${ident}" | grep -Fc "*	CRYPTO_SCRAMBLE_EXT command"`
	security=`echo "${ident}" | grep -Fc " for SECURITY ERASE UNIT"`
	if [ ${security} -ne 0 ]; then
		snotfrozen=`echo "${ident}" | grep -Fc "not	frozen"`
		if [ ${snotfrozen} -eq 0 ]; then
			security="0"
		fi
	fi
	trim=`echo "${ident}" | grep -Fc "*	Data Set Management TRIM supported"`
	maxsect=`hdparm -N /dev/${dev} | awk -F '[/,]' '/max sectors/ { print $2 }'`

	# When resizing SSD to specified size (down?) deallocate the flash.
	if [ ${ssd} -ne 0 -a -n "${size}" ]; then
		err=1
		if [ ${block} -ne 0 -a ${err} -ne 0 ]; then
			echo "Doing block erase sanitize."
			hdparm --yes-i-know-what-i-am-doing --sanitize-block-erase /dev/${dev}
			err=$?
		fi
		if [ ${crypto} -ne 0 -a ${err} -ne 0 ]; then
			echo "Doing cryptograhic erase sanitize."
			hdparm --yes-i-know-what-i-am-doing --sanitize-crypto-scramble /dev/${dev}
			err=$?
		fi
		if [ ${err} -eq 0 ]; then
			echo "Waiting for sanitize to complete."
			until hdparm --sanitize-status /dev/${dev} | grep -Fq "SD0 Sanitize Idle"
			do
				sleep 1
			done
		fi
		if [ ${security} -ne 0 -a ${err} -ne 0 ]; then
			echo "Doing security erase."
			hdparm --yes-i-know-what-i-am-doing --security-erase MyPass /dev/${dev}
			err=$?
		fi
		if [ ${trim} -ne 0 -a ${err} -ne 0 ]; then
			echo "Doing TRIM."
			blkdiscard /dev/${dev}
			err=$?
		fi
		if [ ${err} -ne 0 ]; then
			echo "No method found to deallocate the flash."
		fi
	fi

	if [ -z "${size}" ]; then
		size=${maxsect}
	fi
	echo "Waiting up to ${udevd_timeout} seconds for udevd events to settle"
	udevadm settle -t ${udevd_timeout}
	if [ $? -ne 0 ]; then
		echo "Resize may fail. Udevd events failed to settle within ${udevd_timeout} seconds."
	fi
	echo "Setting Max Sectors to ${size}"
	hdparm --yes-i-know-what-i-am-doing -N p${size} /dev/${dev}	# NB: the p is for persistence
	if [ $? -eq 0 ]; then
		echo "Resize completed successfully.  Reboot may be needed."
		echo "Note that resize can be done only once per power cycle."
	else
		echo "Resize failed."
		exit 1
	fi
	;;
sas|scsi)
	echo "Changing number of LBAs in block descriptor"
	if [ -z "${size}" ]; then
		lsize="-1"
	else
		lsize="${size}"
	fi
	sg_format --resize -c ${lsize} /dev/${dev}
	err=$?

	# When resizing SSD to specified size (down?) format to deallocate the flash.
	ssd=`sg_vpd -p bdc /dev/${dev} 2>/dev/null | grep -Fc "Non-rotating medium"`
	if [ ${err} -eq 0 -a ${ssd} -ne 0 -a -n "${size}" ]; then
		echo "Formatting device."
		sg_format --format --quick /dev/${dev}
	fi

	# Tell the kernel to refresh the device properties.
	echo 1 > /sys/block/${dev}/device/rescan

	if [ ${err} -eq 0 ]; then
		echo "Resize completed successfully."
	else
		echo "Resize failed."
		exit 1
	fi
	;;
nvme)
	# Get controller and current namespace ID.
	ctrlr=/dev/${dev%n*}
	nsid=`nvme get-ns-id /dev/${dev} | awk -F 'namespace-id:' '{ print $2 }'`

	# if nvme device provided is a "fabric" device
	# (i.e. nvme over tcp/rdma) then gracefully exit
	transport="$(cat /sys/class/nvme/${dev%n*}/transport)" 2>/dev/null
	if [ -z ${transport} ]; then
		echo "Unable to determine transport protocol for /dev/${dev}"
		exit 1
	elif [ ${transport} != "pcie" ]; then
		echo "Transport protocol (${transport}) unsupported at this time"
		exit 1
	fi

	# Identify controller properties.
	idctrl=`nvme id-ctrl ${ctrlr}`
	ctrls=`nvme list-ctrl ${ctrlr} | awk -F: '/^\[\s*[0-9]+\]/ { print $2 }' | xargs | tr ' ' ,`

	# Check Namespace Management is supported by the controller.
	oacs=`echo "${idctrl}" | awk -F ' +: ' '$1 == "oacs" { print $2 }'`
	nsm_bit=$((1 << 3))
	nsm=$(((oacs & nsm_bit) == nsm_bit))
	if [ ${nsm} -eq 0 ]; then
		echo "Namespace management not supported, can't resize"
		exit 1
	fi

	# Identify namespace properties.
	idns=`nvme id-ns -n -1 ${ctrlr}`

	# Get current LBA format of the namespace.
	lbafds=`echo "${idns}" | sed -En '/(in use)/s/^lbaf +([[:digit:]]+) .* lbads:([[:digit:]]+) .*$/\1 \2/p'`
	lbaf=${lbafds%% *}
	sector=$((1 << ${lbafds##* }))
	if [ -z "${lbaf}" -o -z "${sector}" ]; then
		echo "Can't get current LBA format"
		exit 1
	fi

	# Make sure we have enough capacity to not fail after delete.
	if [ -n "${size}" ]; then
		ucap=`echo "${idctrl}" | awk -F ' +: ' '$1 == "unvmcap" { print $2 }'`
		ocap=`echo "${idctrl}" | awk -F ' +: ' '$1 == "tnvmcap" { print $2 }'`
		if [ "${size}" -gt $(((ucap + ocap) / sector)) ]; then
			echo "Not enough capacity."
			exit 1
		fi
	fi

	echo "Detaching old namespace."
	nvme detach-ns -n ${nsid} -c ${ctrls} ${ctrlr}
	if [ $? -ne 0 ]; then
		echo "Can't detach old namespace, but continuing anyway"
	fi

	echo "Deleting old namespace."
	nvme delete-ns -n ${nsid} ${ctrlr}
	if [ $? -ne 0 ]; then
		echo "Can't delete old namespace, but continuing anyway"
	fi

	echo "Creating new namespace."
	if [ -z "${size}" ]; then
		idctrl=`nvme id-ctrl ${ctrlr}`
		ucap=`echo "${idctrl}" | awk -F ' +: ' '$1 == "unvmcap" { print $2 }'`
		size=$((ucap / sector))
	fi
	mic=`echo "${idctrl}" | awk -F ' +: ' '$1 == "cmic" { print $2 }'`
	if [ "${mic}" != "0" ]; then
		mic=1
	fi
	created=`nvme create-ns -s ${size} -c ${size} -f ${lbaf} -m ${mic} -d 0 ${ctrlr}`
	echo "${created}"
	nsid=`echo "${created}" | awk -F nsid: '{ print $2 }'`
	if [ -z "${nsid}" ]; then
		echo "Namespace creation failed"
		exit 1
	fi

	echo "Attaching new namespace."
	nvme attach-ns -n ${nsid} -c ${ctrls} ${ctrlr}
	if [ $? -ne 0 ]; then
		echo "Namespace attach failed"
		exit 1
	fi

	echo "Verifying namespace."
	nsize=$((`nvme id-ns -n ${nsid} ${ctrlr} | awk -F ' +: ' '$1 == "nsze" { print $2}'`))
	if [ ${nsize} -eq ${size} ]; then
		echo "Resize completed successfully."
	else
		echo "Size \"${nsize}\" does not match requested.  Resize failed."
		exit 1
	fi
	;;
*)
	echo "Unknown device type"
	exit 1
	;;
esac

exit 0
