#!/usr/bin/bash
###############################################################################
#
# Copyright 2024 NVIDIA Corporation
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of
# this software and associated documentation files (the "Software"), to deal in
# the Software without restriction, including without limitation the rights to
# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
# the Software, and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
#
###############################################################################

rlog()
{
    msg=$(echo "$*" | sed 's/INFO://;s/ERROR:/ERR/;s/WARNING:/WARN/')
    bfrshlog "$msg"
}

ilog()
{
    msg="[$(date +%H:%M:%S)] $*"
    echo "$msg" >> $LOG
    echo "$msg" > /dev/hvc0
}

log()
{
    ilog "$*"
    rlog "$*"
}

bind_partitions()
{
    local wdir=$1
    mount --bind /proc ${wdir}/proc
    mount --bind /dev ${wdir}/dev
    mount --bind /dev/pts ${wdir}/dev/pts
    mount --bind /sys ${wdir}/sys
    mount -t efivarfs none ${wdir}/sys/firmware/efi/efivars
}

unmount_partition()
{
    retries=0
    while (grep -wq $1 /proc/mounts) && [ $retries -lt 5 ]; do
        ilog "Unmounting $1"
        umount $1 || umount -l $1
        sync
        let retries++
    done
}

unmount_partitions()
{
    local wdir=$1
    ilog "Unmount partitions"
    for part in \
        ${wdir}/dev/pts \
        ${wdir}/dev \
        ${wdir}/sys/fs/fuse/connections \
        ${wdir}/sys/firmware/efi/efivars \
        ${wdir}/sys \
        ${wdir}/proc
    do
        unmount_partition $part
    done
}

RC=0
ACTVATION_STATUS=0
# Maximum idle time in seconds to decide no more data from BOOTFIFO.
IDLE_MAX=5
INITIAL_DATA_RECEIVE_TIMEOUT=20

logfile=dpu.installation.log
LOG=/var/log/$logfile

# flow: upgrade, activate, cancel, compatibility
# compatibility: compatibility flow - to support old infrastructure
COMPATIBILITY_FLOW=0
UPGRADE_FLOW=0
ACTIVATE_FLOW=0
CANCEL_FLOW=0

SCRIPTS_DIR=$(dirname $0)

. ${SCRIPTS_DIR}/bf-upgrade.env/rshim

# Source a configuration if exists.
[ -f /etc/bf-upgrade.conf ] && source /etc/bf-upgrade.conf
[ -f /etc/bf.cfg ] && source /etc/bf.cfg

[ ! -e /proc/acpi/button/lid/LID/state ] && exit 0

BOOTFIFO="/sys/bus/platform/devices/MLNXBF04:00/driver/bootfifo"
if [ ! -e "$BOOTFIFO" ]; then
  BOOTFIFO="/sys/devices/platform/MLNXBF04:00/bootfifo"
  [ ! -e "$BOOTFIFO" ] && exit 1
fi

# Only handles upgrade in closed state.
state=$(cat /proc/acpi/button/lid/LID/state | awk '{print $2}' 2>/dev/null)
[ "$state" != "closed" ] && exit 0

TMPDIR=${TMPDIR:-"/tmp"}
WDIR=${WDIR:-"$(mktemp -d ${TMPDIR}/bfb.XXXXXX)"}
BFB_DIR=${TMPDIR}/bfb
UPDATE_IN_PROGRESS=0
UPDATE_IN_PROGRESS_FILE=${TMPDIR}/update_in_progress
if [ -e $UPDATE_IN_PROGRESS_FILE ]; then
	ilog "Update in progress. Awaiting for activation."
	UPDATE_IN_PROGRESS=1
	/bin/rm -rf $UPDATE_IN_PROGRESS_FILE
else
	/bin/rm -rf ${BFB_DIR}/* ${BFB_DIR}
	/bin/ln -sf ${WDIR} ${BFB_DIR}
fi

# Read 8 bytes starting from the 5th byte (offset 4)
RSHIM_SP2=$(read_SP2)
ilog "UINT64 value starting from the 5th byte (little-endian): $RSHIM_SP2"

set_upgrade_progress 0 0

if is_mode_upgrade $RSHIM_SP2; then
	ilog "RSHIM SP2 is $RSHIM_SP2: Running Upgrade flow"
	touch $UPDATE_IN_PROGRESS_FILE
	if [ $UPDATE_IN_PROGRESS -eq 1 ]; then
		ilog "Previous update in progress. Awaiting for activation."
		exit 1
	fi
	UPGRADE_FLOW=1
fi
if is_mode_activate $RSHIM_SP2; then
	ilog "RSHIM SP2 is $RSHIM_SP2: Running ActivateFirmware flow"
	ACTIVATE_FLOW=1
fi
if is_mode_cancel $RSHIM_SP2; then
	ilog "RSHIM SP2 is $RSHIM_SP2: Running Cancel flow"
	CANCEL_FLOW=1
fi
if [ $UPGRADE_FLOW -eq 0 ] && [ $ACTIVATE_FLOW -eq 0 ] && [ $CANCEL_FLOW -eq 0 ]; then
	ilog "RSHIM SP2 is $RSHIM_SP2"
	COMPATIBILITY_FLOW=1
fi

# IS_BUNDLE is set to 1 if the BFB is a bf-bundle or bf-fwbundle (non-flat) BFB
IS_BUNDLE=0
export UPDATE_ATF_UEFI=${UPDATE_ATF_UEFI:-"yes"}
export UPDATE_DPU_OS=${UPDATE_DPU_OS:-"no"}
export UPDATE_BMC_FW=${UPDATE_BMC_FW:-"yes"}
export UPDATE_CEC_FW=${UPDATE_CEC_FW:-"yes"}
export UPDATE_DPU_GOLDEN_IMAGE=${UPDATE_DPU_GOLDEN_IMAGE:-"yes"}
export UPDATE_NIC_FW_GOLDEN_IMAGE=${UPDATE_NIC_FW_GOLDEN_IMAGE:-"yes"}
export WITH_NIC_FW_UPDATE=${WITH_NIC_FW_UPDATE:-"yes"}
export RUNTIME_UPGRADE="yes"
export USE_BFB_INSTALL=${USE_BFB_INSTALL:-"no"}
export COMPATIBILITY_ACTIVATION=${COMPATIBILITY_ACTIVATION:-"no"}

if [[ $COMPATIBILITY_FLOW -eq 1 || $UPGRADE_FLOW -eq 1 ]]; then
	# Add a delay for the boot-fifo to be filled.
	sleep $IDLE_MAX

	# Example code to fetch the upgrade bfb.
	# !!!Use eMMC or NVME to avoid running out of memory in NIC mode.!!!
	UPGRADE_IMAGE=${WDIR}/upgrade.bfb
	rm -f ${UPGRADE_IMAGE}* 2>/dev/null

	idle_cnt=0
	data_receive_timeout=$INITIAL_DATA_RECEIVE_TIMEOUT
	while [ $idle_cnt -lt $data_receive_timeout ]; do
		cat "$BOOTFIFO" > ${UPGRADE_IMAGE}.tmp
		[ -e ${UPGRADE_IMAGE}.tmp ] && filesize=$(du -b ${UPGRADE_IMAGE}.tmp | awk '{print $1}')
		if [[ -z "$filesize" || ."$filesize" == ."0" ]]; then
			# Done if no more data in 5 seconds.
			idle_cnt=$((idle_cnt + 1))
			sleep 1
		else
			idle_cnt=0
			data_receive_timeout=$IDLE_MAX
			cat ${UPGRADE_IMAGE}.tmp >> ${UPGRADE_IMAGE}
		fi
	done

	/bin/rm -f ${UPGRADE_IMAGE}.tmp
	cd ${WDIR}
	/bin/rm -f dump-*

	if [ ! -s "${UPGRADE_IMAGE}" ]; then
		log "ERR Failed to receive a BFB"
		/bin/rm -rf ${WDIR} ${BFB_DIR} $UPDATE_IN_PROGRESS_FILE
		set_upgrade_progress 100 1
		exit 1
	fi

	log "Extracting BFB ${UPGRADE_IMAGE}"
	mlx-mkbfb -x ${UPGRADE_IMAGE}
	rm -f ${UPGRADE_IMAGE}

	# PLDM BFB format:
	# dump-boot-args-v0 - bf.cfg
	# dump-capsule-v0 - BSP upgrade capsule
	# dump-image-v0 - BMC firmware image
	# dump-upgrade-image-v0 - CEC upgrade image
	# dump-ramdisk-v0 - NIC firmware Golden Image
	# dump-dpu-gi-v0 - DPU Golden Image
	# dump-nicfw-v0 - NIC firmware

	# BUNDLE BFB format:
	# dump-boot-args-v0 - bf.cfg
	# dump-initramfs-v0 - BFB installation environment that contains all the components (BMC, CEC, NIC FW, etc.)
	# dump-image-v0 - Linux kernel
fi # compatibility flow or upgrade flow - finish reading BFB

cd ${BFB_DIR}
if [ -e dump-boot-args-v0 ]; then
	if ( bash -n dump-boot-args-v0 ); then
		if (grep -qw "initrd=initramfs" dump-boot-args-v0); then
			ilog "Detected non-flat bundle BFB"
			IS_BUNDLE=1
			if [ "$COMPATIBILITY_ACTIVATION" == "yes" ]; then
				ACTIVATE_FLOW=1
			fi
		else
			if [ "$(grep -c 'rootwait' dump-boot-args-v0)" -eq 1 ] && [ "$(wc -l < dump-boot-args-v0)" -le 1 ]; then
				ilog "dump-boot-args-v0: Only one line and it contains rootwait"
			else
				ilog "Found bf.cfg"
				cat dump-boot-args-v0 > /etc/bf.cfg
				. dump-boot-args-v0
			fi
			if [ "$LFWP" == "yes" ]; then
				ilog "LFWP upgrade flow including NIC Firmware activation"
				ACTIVATE_FLOW=1
			fi
		fi
	else
		log "ERR Invalid bf.cfg"
		set_upgrade_progress 100 1
		exit 1
	fi
fi

if [ $COMPATIBILITY_FLOW -eq 1 ]; then
	ilog "Running Compatibility flow"
fi

if [ -e dump-initramfs-v0 ]; then
	log "Detected Bundle BFB with initramfs"
	IS_BUNDLE=1
else
	if [ -e dump-nicfw-v0 ]; then
		log "Detected Flat fwbundle BFB"
	else
		log "Detected PLDM BFB format"
	fi
	IS_BUNDLE=0
fi

# Include scripts that provide upgrade infrastructure
if (bash -n ${SCRIPTS_DIR}/bf-upgrade.env/common 2>/dev/null); then
	. ${SCRIPTS_DIR}/bf-upgrade.env/common
fi

if (bash -n ${SCRIPTS_DIR}/bf-upgrade.env/atf-uefi 2>/dev/null); then
	. ${SCRIPTS_DIR}/bf-upgrade.env/atf-uefi
fi

if (bash -n ${SCRIPTS_DIR}/bf-upgrade.env/nic-fw 2>/dev/null); then
	. ${SCRIPTS_DIR}/bf-upgrade.env/nic-fw
fi

if (bash -n ${SCRIPTS_DIR}/bf-upgrade.env/bmc 2>/dev/null); then
	. ${SCRIPTS_DIR}/bf-upgrade.env/bmc
fi

if [[ -n "$BMC_USER" && -n "$BMC_PASSWORD" ]]; then
	if [ -n "$BMC_IP" ]; then
		if ! (ping -c 3 $BMC_IP > /dev/null 2>&1); then
			if ! create_vlan; then
				log "ERR Failed to create VLAN"
				skip_bmc 0
				RC=$(( RC + 1 ))
			fi
		fi
	else
		log "WARN BMC_IP is not set"
		skip_bmc 0
	fi
else
	log "WARN BMC credentials are not set"
	skip_bmc 0
fi

BMC_LINK_UP="no"
if [ -n "$BMC_IP" ]; then
	if (ping -c 3 $BMC_IP > /dev/null 2>&1); then
		BMC_LINK_UP="yes"
	fi
fi

cd ${BFB_DIR}
if [ $IS_BUNDLE -eq 0 ]; then
	# PLDM BFB or Flat fwbundle BFB
	if [[ $ACTIVATE_FLOW -eq 1 || $COMPATIBILITY_FLOW -eq 1 ]]; then
		if [ "$UPDATE_ATF_UEFI" == "yes" ]; then
			if [ -e dump-capsule-v0 ]; then
				if ! function_exists update_atf_uefi; then
					log  "ERR BSP upgrade function does not exist"
					UPDATE_ATF_UEFI="no"
				fi
			else
				log "ERR BSP upgrade capsule dump-capsule-v0 was not found"
				UPDATE_ATF_UEFI="no"
			fi
		fi
	fi

	if [[ $UPGRADE_FLOW -eq 1 || $COMPATIBILITY_FLOW -eq 1 ]]; then
		if ! function_exists bmc_components_update; then
			log  "ERR BMC upgrade function does not exist"
			skip_bmc 0
		fi

		if [ "$UPDATE_BMC_FW" == "yes" ]; then
			if [ -e dump-image-v0 ]; then
				BMC_IMAGE=$(readlink -f dump-image-v0)
			else
				log "ERR BMC firmware dump-image-v0 was not found"
				UPDATE_BMC_FW="no"
			fi
		fi

		if [ "$UPDATE_CEC_FW" == "yes" ]; then
			if [ -e dump-upgrade-image-v0 ]; then
				CEC_IMAGE=$(readlink -f dump-upgrade-image-v0)
			else
				log "ERR CEC firmware dump-upgrade-image-v0 was not found"
				UPDATE_CEC_FW="no"
			fi
		fi

		if [ "$UPDATE_DPU_GOLDEN_IMAGE" == "yes" ]; then
			if [ -e dump-dpu-gi-v0 ]; then
				DPU_GOLDEN_IMAGE=$(readlink -f dump-dpu-gi-v0)
			else
				log "DPU Golden Image dump-dpu-gi-v0 was not found. Skipping..."
				UPDATE_DPU_GOLDEN_IMAGE="no"
			fi
		fi

		if [ "$UPDATE_NIC_FW_GOLDEN_IMAGE" == "yes" ]; then
			if [ -e dump-ramdisk-v0 ]; then
				NIC_FW_GOLDEN_IMAGE=$(readlink -f dump-ramdisk-v0)
			else
				log "NIC firmware golden image dump-upgrade-image-v0 was not found. Skipping..."
				UPDATE_NIC_FW_GOLDEN_IMAGE="no"
			fi
		fi

		if [ "$WITH_NIC_FW_UPDATE" == "yes" ]; then
			if [ -e dump-nicfw-v0 ]; then
				if function_exists update_nic_firmware; then
					NIC_FW_BIN=$(readlink -f dump-nicfw-v0)
				else
					log  "ERR NIC Firmware upgrade function does not exist"
					WITH_NIC_FW_UPDATE="no"
				fi
			else
				# log "ERR NIC firmware binary dump-nicfw-v0 was not found"
				WITH_NIC_FW_UPDATE="no"
			fi
		fi

		calculate_total_weight
		update_progress install_setup 0
	fi # compatibility flow or upgrade flow - finish updating components

	if [ "$UPDATE_ATF_UEFI" == "yes" ]; then
		update_atf_uefi_status=0
		if [[ $ACTIVATE_FLOW -eq 1 || $COMPATIBILITY_FLOW -eq 1 ]]; then
			# Update components
			update_atf_uefi $(readlink -f dump-capsule-v0)
			update_atf_uefi_status=$?
			ACTVATION_STATUS=$(( ACTVATION_STATUS + update_atf_uefi_status ))
		fi
		if [[ $UPGRADE_FLOW -eq 1 || $COMPATIBILITY_FLOW -eq 1 ]]; then
			update_progress atf_uefi $update_atf_uefi_status
		fi
	fi

	if [[ $UPGRADE_FLOW -eq 1 || $COMPATIBILITY_FLOW -eq 1 ]]; then
		bmc_components_update

		if [ "$WITH_NIC_FW_UPDATE" == "yes" ]; then
			nic_firmware_update
		fi
	fi

	if [[ $ACTIVATE_FLOW -eq 1 || $COMPATIBILITY_FLOW -eq 1 ]]; then
		if [ -e ${BFB_DIR}/bmc_firmware_staged ]; then
			ilog "Activating BMC firmware"
			bmc_activate
			ACTVATION_STATUS=$(( ACTVATION_STATUS + $? ))
			/bin/rm -f ${BFB_DIR}/bmc_firmware_staged
		else
			ilog "BMC firmware was not staged. Skipping BMC firmware activation."
		fi
	fi
else
	#BUNDLE BFB that contains Linux kernel and initramfs
	if [ ! -e dump-initramfs-v0 ]; then
		log "ERR BFB's initramfs dump-initramfs-v0 was not found"
		set_upgrade_progress 100 1
		exit 1
	fi

	log "Extracting BFB's initramfs"
	mkdir -p initramfs
	cd initramfs
	gzip -d < ../dump-initramfs-v0 | cpio -id
	cat ../dump-boot-args-v0 > ./etc/bf.cfg
	echo "UPDATE_DPU_OS=no" >> ./etc/bf.cfg
	echo "RUNTIME_UPGRADE=$RUNTIME_UPGRADE" >> ./etc/bf.cfg
	if (ping -c 3 $BMC_IP > /dev/null 2>&1); then
	    BMC_LINK_UP="yes" >> ./etc/bf.cfg
	fi
	if [ "$USE_BFB_INSTALL" == "yes" ]; then
		if [ ! -e ubuntu/install.sh ]; then
			log "ERR ubuntu/install.sh was not found. Only bf-fwbundle is supported."
			set_upgrade_progress 100 1
			exit 1
		fi
		log "Running ubuntu/install.sh"

		bind_partitions ${WDIR}/initramfs

		echo "[$(date +%H:%M:%S)] Starting upgrade from bundle" > /var/log/bf-upgrade.log
		chroot ${WDIR}/initramfs bash -c  /ubuntu/install.sh
		cat ${WDIR}/initramfs/tmp/*log >> /var/log/bf-upgrade.log 2> /dev/null || true
		unmount_partitions ${WDIR}/initramfs
	else
		CHROOT_DIR="$PWD"
		. ./etc/bf.cfg
		calculate_total_weight
		update_progress install_setup 0

		# Update components
		if [[ $ACTIVATE_FLOW -eq 1 || $COMPATIBILITY_FLOW -eq 1 ]]; then
			if [ "$UPDATE_ATF_UEFI" == "yes" ]; then
				if [ -e ./lib/firmware/mellanox/boot/capsule/boot_update2.cap ]; then
					update_atf_uefi $(readlink -f ./lib/firmware/mellanox/boot/capsule/boot_update2.cap)
				fi
			fi
		fi

		if [[ $UPGRADE_FLOW -eq 1 || $COMPATIBILITY_FLOW -eq 1 ]]; then
			update_progress atf_uefi 0
		fi
		if [[ $UPGRADE_FLOW -eq 1 || $COMPATIBILITY_FLOW -eq 1 ]]; then
			if [ -d ./lib/firmware/mellanox/bmc/ ]; then
				BMC_PATH=$(readlink -f ./lib/firmware/mellanox/bmc/)
			fi
			BMC_IMAGE=$(/bin/ls ${BMC_PATH}/bf3-bmc*.fwpkg 2> /dev/null)
			if [ -d ./lib/firmware/mellanox/cec/ ]; then
				CEC_PATH=$(readlink -f ./lib/firmware/mellanox/cec/)
			fi
			CEC_IMAGE=$(/bin/ls ${CEC_PATH}/bf3-cec*.fwpkg 2> /dev/null)
			if [ -d ./lib/firmware/mellanox/bmc/ ]; then
				DPU_GI_PATH=$(readlink -f ./lib/firmware/mellanox/bmc/)
				NIC_FW_GI_PATH=$(readlink -f ./lib/firmware/mellanox/bmc/)
			fi
			DPU_GOLDEN_IMAGE=$(/bin/ls ${DPU_GI_PATH}/bf3*preboot*.bfb 2> /dev/null || /bin/ls ${DPU_GI_PATH}/bf3*preboot*.pldm 2> /dev/null)
			NIC_FW_GOLDEN_IMAGE=$(/bin/ls ${NIC_FW_GI_PATH}/*${dpu_part_number}*bfb 2> /dev/null || /bin/ls ${NIC_FW_GI_PATH}/*${dpu_part_number}*pldm 2> /dev/null)

			bmc_components_update

			if [ "$WITH_NIC_FW_UPDATE" == "yes" ]; then
				if [[ $UPGRADE_FLOW -eq 1 || $COMPATIBILITY_FLOW -eq 1 ]]; then
					nic_firmware_update
				fi
			fi
		fi
	fi
	cd ..
fi

if [ -e ${BFB_DIR}/nic_firmware_updated ]; then
	if [[ $ACTIVATE_FLOW -eq 1 || $COMPATIBILITY_FLOW -eq 1 ]]; then
		ilog "Resetting NIC firmware"
		NIC_FW_UPDATE_DONE=1
		NIC_FW_UPDATE_PASSED=1
		reset_nic_firmware
		reset_nic_firmware_status=$?
		ACTVATION_STATUS=$(( ACTVATION_STATUS + reset_nic_firmware_status ))
		/bin/rm -f ${BFB_DIR}/nic_firmware_updated
	fi
elif [ -e ${BFB_DIR}/nic_firmware_update_failed ]; then
	if [[ $ACTIVATE_FLOW -eq 1 || $COMPATIBILITY_FLOW -eq 1 ]]; then
		ilog "ERR NIC firmware update failed. Skipping activation."
		ACTIVATE_FLOW=0
		ACTVATION_STATUS=$(( ACTVATION_STATUS + 1 ))
		RC=$(( RC + 1 ))
		/bin/rm -f ${BFB_DIR}/nic_firmware_update_failed
	fi
fi

# Keep the WDIR for activate
# During the activate flow, the WDIR is used to run the ATF/UEFI upgrade
if [[ $ACTIVATE_FLOW -eq 1 || $CANCEL_FLOW -eq 1 || $COMPATIBILITY_FLOW -eq 1 || $RC -ne 0 ]]; then
	if [ $RC -ne 0 ]; then
		log "ERR $0 flow failed with RC: $RC"
		ilog "Removing BFB directory: $BFB_DIR"
	fi
	cd ${WDIR}
	/bin/rm -rf dump-*
	cd ..

	if [ ! -d ${WDIR}/initramfs/dev/pts ]; then
		/bin/rm -rf ${WDIR} ${BFB_DIR} $UPDATE_IN_PROGRESS_FILE
	else
		ilog "Failed to unmount ${WDIR}/initramfs"
	fi
fi

if [ $UPGRADE_FLOW -eq 1 ]; then
	if [ $RC -eq 0 ]; then
		if [ $ACTIVATE_FLOW -eq 1 ]; then
			log "Runtime upgrade finished successfully"
		else
			log "Runtime upgrade finished. Activate firmware to complete the update."
		fi
	else
		log "ERR Runtime upgrade finished with errors"
	fi
	update_progress finished $RC
fi

if [ $ACTIVATE_FLOW -eq 1 ]; then
	if [ $ACTVATION_STATUS -eq 0 ]; then
		log "Activation finished successfully"
	else
		log "ERR Activation finished with errors"
	fi
	update_progress activate $((ACTVATION_STATUS + RC))
fi

sleep 3
exit $((RC + ACTVATION_STATUS))
