#!/bin/sh
# vim: set ts=4 sw=4:
set -eu

# Prints size (in megabytes) of the specified device.
dev_size_mb() {
	local dev_path="$1"

	local bytes="$(blockdev --getsize64 "$dev_path")"
	expr $bytes / 1024 / 1024
}

is_partition() {
	local dev_path="$1"
	test -e "/sys/class/block/${dev_path##*/}/partition"
}

is_last_partition() {
	local part_path="$1"
	local disk_path="$2"

	local part_uuid="$(partx -sg -o UUID "$part_path" "$disk_path")"
	local last_uuid="$(partx -sg -o UUID "$disk_path" | tail -n 1)"

	[ "$part_uuid" = "$last_uuid" ]
}

part_resize() {
	local disk_path="$1"
	local partno="$2"

	# Note: sfdisk prints "Resource busy" error if --no-tell-kernel is not used
	#  and we must use 'partx --update' anyway.
	echo ', +' | sfdisk --no-reread --no-tell-kernel --color=never --partno "$partno" "$disk_path"

	# Tell kernel that the partition table has been modified.
	partx --update "$disk_path"
}

part_resize_if_needed() {
	local part_path="$1"

	local part_name="${dev_path##*/}"
	local disk_path="$(readlink -f /sys/class/block/$part_name/..)"
	disk_path="/dev/${disk_path##*/}"

	# sfdisk supports only DOS and GPT partitions.
	local scheme="$(partx -sg -o SCHEME "$part_path" "$disk_path")"
	case "$scheme" in
		dos | gpt) ;;  # continue
		*) echo "WARN: Found unsupported partition table on $disk_path: $scheme" >&2; return 0;;
	esac

	# Only the last partition can be expanded.
	if ! is_last_partition "$part_path" "$disk_path"; then
		return 0
	fi

	local partno="$(cat /sys/class/block/$part_name/partition)"
	# Size of the block device in sectors.
	local disk_size="$(blockdev --getsz "$disk_path")"
	# The end of the last partition in sectors.
	local last_part_end="$(partx -sg --sector-size 512 -o END "$part_path")"

	# Ignore difference that is less than ~4 MiB.
	if [ "$(( $disk_size - $last_part_end ))" -lt 8192 ]; then
		return 0
	fi

	echo "Resizing partition no. $partno on disk $disk_path" >&2
	part_resize "$disk_path" "$partno"
}

# Prints size (in megabytes) of the specified filesystem.
fs_size_mb() {
	local fs_type="$1"
	local dev_path="$2"

	case "$fs_type" in
		btrfs)
			btrfs filesystem show --mbytes "$dev_path" \
				| sed -En "s|.* size ([0-9]+).*path $dev_path$|\1|p"
		;;
		ext*)
			# Note: df doesn't count reserved blocks, that's why we use dumpe2fs.
			dumpe2fs -h "$dev_path" 2>/dev/null \
				| sed -En 's/Block (count|size):\s*([0-9]+)/\2/p' \
				| xargs \
				| awk '{ print int($1 * $2 / 1024 / 1024) }'
		;;
		*) return 1;;
	esac
}

# Prints the first mount point of the specified device.
fs_mountpoint() {
	local dev_path="$1"

	mount | grep "^$dev_path " | cut -d' ' -f3 | head -n1
}

fs_resize() {
	local fs_type="$1"
	local dev_path="$2"

	case "$fs_type" in
		btrfs) btrfs filesystem resize max "$(fs_mountpoint "$dev_path")";;
		ext*) resize2fs "$dev_path";;
		*) return 1;;
	esac
}

fs_resize_if_needed() {
	local fs_type="$1"
	local dev_path="$2"

	local fs_size="$(fs_size_mb "$fs_type" "$dev_path")"
	local dev_size="$(dev_size_mb "$dev_path")"

	if [ -z "$fs_size" ] || [ -z "$dev_size" ]; then
		echo "WARN: Failed to get size of $fs_type FS or device on $dev_path" >&2

	elif [ $fs_size -lt $dev_size ]; then
		echo "Resizing $fs_type on $dev_path from $fs_size MiB to device size $dev_size MiB" >&2
		fs_resize "$fs_type" "$dev_path"
	fi
}


#-------------------------------- Main -------------------------------

. "$(dirname "$(readlink -f "$0")")/utils.sh"

if yesno "${GROWFS_DISABLE:-}"; then
	exit 0
fi

mount | cut -d' ' -f1,5 | sort | uniq | while read dev_path fs_type; do
	case "$fs_type" in
		# Note: ext2 can't be resized online.
		ext3 | ext4 | btrfs)
			is_partition "$dev_path" && part_resize_if_needed "$dev_path"
			fs_resize_if_needed "$fs_type" "$dev_path"
		;;
	esac
done
