#!/bin/bash

# SPDX-FileCopyrightText: Copyright (c) 2019-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

# This is a script to resize partition and filesystem on the root partition
# This will consume all un-allocated sapce on SD card after boot.

set -e

function usage()
{
	if [ -n "${1}" ]; then
		echo "${1}"
	fi

	echo "Usage:"
	echo "${script_name} [options]"
	echo ""
	echo "Available options are:"
	echo ""
	echo "	-h | --help"
	echo "		Show this usage"
	echo ""
	echo "	-c | --check"
	echo "		Check whether ${script_name} can be used on current platform"
	echo ""
	echo "	-g | --get"
	echo "		Show APP partition size in MB"
	echo ""
	echo "	-m | --max"
	echo "		Show APP partition maximum available size in MB"
	echo ""
	echo "	-s <app_size> | --size <app_size>"
	echo "		Set APP partition size in MB"
	echo ""
	echo "Without any option means consume all unallocated"
	echo "space on SD card."
	echo "This script extends the filesystem to the maximum"
	echo "size by default."
	echo ""
	echo "Example:"
	echo "${script_name}"
	echo "${script_name} -s 16384"
}

function check_pre_req()
{
	# Return directly if rootfs A/B is enabled
	if ! nvbootctrl -t rootfs is-rootfs-ab-enabled; then
		echo "WARNING: expanding rootfs is not supported when rootfs A/B is enabled" >&2
		return
	fi

	root_dev="$(sed -ne 's/.*\broot=\([^ ]*\)\b.*/\1/p' < /proc/cmdline)"

	if [[ "${root_dev}" == *UUID* ]]; then
		root_dev="$(/sbin/findfs "${root_dev}")"
		# Return if root partition is LUKS formatted
		root_type="$(blkid -s TYPE -o value "${root_dev}")"
		if [ "${root_type}" == "crypto_LUKS" ]; then
			return
		fi
	fi

	if [[ "${root_dev}" == /dev/mmcblk* ]] || [[ "${root_dev}" == /dev/sd* ]] \
		|| [[ "${root_dev}" == /dev/nvme* ]]; then
		block_dev="$(/bin/lsblk -no pkname "${root_dev}")"
		root_dev="$(echo "${root_dev}" | sed 's/^.....//g')"
	fi

	if [ "${block_dev}" != "" ]; then
		# Check whether APP partition is located at the end of the disk
		if ! cat /sys/block/"${block_dev}"/"${root_dev}"/start > /dev/null 2>&1 ; then
			return
		fi
		root_start_sector="$(cat /sys/block/"${block_dev}"/"${root_dev}"/start)"
		all_start_sectors="$(cat /sys/block/"${block_dev}"/*/start)"
		hw_sector_size="$(cat /sys/block/"${block_dev}"/queue/hw_sector_size)"

		is_last="true"
		for start_sector in ${all_start_sectors}; do
			if [[ "${start_sector}" -gt "${root_start_sector}" ]]; then
				is_last="false"
				break
			fi
		done

		if [ "${is_last}" == "true" ]; then
			get_app_size
			if max_available_app_size; then
				support_resizefs="true"
			fi
		fi

		# Extend the filesystem to the maximum size if the
		# APP reaches the maximum size
		if [ "${is_last}" == "true" ] && [ "${support_resizefs}" == "false" ]; then
			resize2fs "/dev/${root_dev}" > /dev/null 2>&1
		fi
	fi
}

function get_app_size()
{
	partition_size="$(cat /sys/block/"${block_dev}"/"${root_dev}"/size)"
	cur_app_size_sector="${partition_size}"
}

function max_available_app_size()
{
	# Move backup GPT header to end of disk
	echo -e "Fix\n" | parted ---pretend-input-tty /dev/"${block_dev}" print free >/dev/null 2>&1

	app_start_sector="$(cat /sys/block/"${block_dev}"/"${root_dev}"/start)"
	last_usable_sector="$(sgdisk -p /dev/"${block_dev}" | \
				grep "last usable sector" | \
				awk '{print $10}')"
	# Get maximum APP size in sector
	max_app_size_sector="$((last_usable_sector - app_start_sector + 1))"
	# Align the size to 4KiB in sector
	max_app_size_sector="$((max_app_size_sector * hw_sector_size / 4096 * 4096 / hw_sector_size))"
	# Return 1 if APP has been extended to its maximum size
	if [ "${cur_app_size_sector}" -ge "${max_app_size_sector}" ]; then
		return 1
	fi
	return 0
}

function process_specified_size()
{
	# Extend root partition to maximum size if the "--size" is not specified.
	if [ "${user_app_size}" == "" ]; then
		app_size_sector="${max_app_size_sector}"
		return 0
	fi

	# If specified size is invalid, report error and exit.
	# If specified size is larger than the maximum size, report warning and
	# extend it to the maximum size.
	# If sepcified size is 0, report warning and exit
	local size_MB=
	size_MB="$(echo "${user_app_size}" | sed 's/^[+]//' | sed 's/M$//')"
	# Confirm that the specified size is valid
	size_MB="$(echo "${size_MB}" | sed -n "/^[0-9]\+$/p")"
	if [ "${size_MB}" == "" ]; then
		echo "ERROR: invalid size ${user_app_size}" >&2
		exit 1
	fi
	size_MB="$((size_MB))"
	if [ "${size_MB}" == "0" ]; then
		echo "WARNING: specified size(${user_app_size}) is zero, exit" >&2
		exit 0
	fi
	# Convert the specified APP size (MB) into sectors
	app_size_sector="$((size_MB * 1048576 / hw_sector_size))"
	if [ "${app_size_sector}" -eq 0 ]; then
		echo "WARNING: specified size (${user_app_size}) is smaller than one sector, exit" >&2
		exit 0
	fi
	if [ "${app_size_sector}" -gt "${max_app_size_sector}" ]; then
		echo "WARNING: ${user_app_size} is beyond the size of the disk, just extend it to the maximum size ($((max_app_size_sector * hw_sector_size / 1048576)))" >&2
		app_size_sector="${max_app_size_sector}"
	fi
	# Make sure APP size is alinged to 4KiB in sector
	app_size_sector="$((app_size_sector * hw_sector_size / 4096 * 4096 / hw_sector_size))"
}

function parse_args()
{
	while [ -n "${1}" ]; do
		case "${1}" in
		-h | --help)
			usage
			exit 0
			;;
		-c | --check)
			echo "${support_resizefs}"
			exit 0
			;;
		-g | --get)
			if [ "${support_resizefs}" = "false" ]; then
				echo "ERROR: ${script_name} doesn't support this platform." >&2
				exit 1
			fi
			cur_app_size_MB="$((cur_app_size_sector * hw_sector_size / 1048576))"
			echo "${cur_app_size_MB}"
			exit 0
			;;
		-m | --max)
			if [ "${support_resizefs}" = "false" ]; then
				echo "ERROR: ${script_name} doesn't support this platform." >&2
				exit 1
			fi
			max_app_size_MB="$((max_app_size_sector * hw_sector_size / 1048576))"
			echo "${max_app_size_MB}"
			exit 0
			;;
		-s | --size)
			[ -n "${2}" ] || usage "Not enough parameters"
			user_app_size="+${2}M"
			shift 2
			;;
		*)
			usage "Unknown option: ${1}"
			exit 1
			;;
		esac
	done
}

script_name="$(basename "${0}")"
support_resizefs="false"
user_app_size=
cur_app_size_sector="0"
max_app_size_sector="0"
root_dev=""
block_dev=""
hw_sector_size=
app_size_sector=

check_pre_req
parse_args "${@}"

if [ "${support_resizefs}" = "false" ]; then
	echo "ERROR: ${script_name} doesn't support this platform." >&2
	exit 1
fi

# Get partition number of the root partition
partition_num="$(cat /sys/block/"${block_dev}"/"${root_dev}"/partition)"

# Get start sector of the root partition
start_sector="$(cat /sys/block/"${block_dev}"/"${root_dev}"/start)"

# Process the "user_app_size" variant
process_specified_size

# Get the end sector of extended root partition
end_sector="$((start_sector + app_size_sector - 1))"
echo -e "Yes\n${end_sector}s\n" | parted ---pretend-input-tty /dev/"${block_dev}" unit "s" resizepart "${partition_num}" "${end_sector}s"

# Inform kernel and OS about change in partition table and root
# partition size
partprobe "/dev/${block_dev}"

# Resize filesystem on root partition to consume all un-allocated
# space on disk
resize2fs "/dev/${root_dev}"