#!/bin/bash
#chkconfig: 1235 80 30
#description: HCT config service script

VERSION=1.0.2-2023-0512-rc

#max supported number of processes
MAX_PROGRESS="128"

#CCP_NUM value can be: all | 1 | 2 | 3 | 4 | ... | max ccp num |
CCP_NUM="all"

#CCP_TYPE value can be: both | ntb | psp
CCP_TYPE="both"

#-------> system dirs and special files <---------
CURRENT_DIR=$(cd $(dirname $0)/../../; pwd)
LOCKFILE=/var/lock/subsys/hct.lock
HCT_INFO_LOG=/var/log/hct_info.log
MDEV_PATH="/sys/bus/mdev/devices/"
HCT_SCRIPT_DIR="${CURRENT_DIR}/hct/script"
ENGINE_DIR="${CURRENT_DIR}/lib64/engines-1.1/"
HCT_LIB_DIR="${CURRENT_DIR}/lib64/"
HCT_TEST_DIR="${CURRENT_DIR}/bin/"
HCT_SPEED_DIR="${CURRENT_DIR}/bin/"


#------->>> function definitions  <<<---------

#-------------------------------------
# get command opts
#-------------------------------------
_getopt() {
	ARGS=`getopt -o n:t:d:p: --long ccp_num:,ccp_type:,script_dir:,max_prog: -- "$@"`
	if [ $? != 0 ]; then
	echo "Args error terminating..."
	exit 1
	fi

	eval set -- "${ARGS}"

	while true
	do
	case "$1" in
		-n|--ccp_num)
			CCP_NUM=$2;
			shift 2
			;;
		-t|--ccp_type)
			CCP_TYPE=$2;
			shift 2
			;;
	-d|--script_dir)
			HCT_SCRIPT_DIR=$2;
			shift 2
			;;
	-p|--max_prog)
			MAX_PROGRESS=$2;
			shift 2
			;;
		--)
			shift
			break
			;;
		*)
			echo "Internal error!"
			exit 1
			;;
	esac
	done


	for arg in $@
	do
	echo "Invalid args $arg"
	echo "Terminating!"
	exit 1
	done

}

#-------------------------------------
# help messages for this service
#-------------------------------------
usage() {

	echo "
HYGON HCT CONFIG SERVICE

Version: $VERSION

Descript: manage hygon hct ccp binding & etc.

USAGE:
	#Start hct service with given parameters; \"ccp_type\" can be \"both\", \"ntb\", \"psp\"
	start [-n --ccp_num] [-t --ccp_type] [-p --max_prog]

	#Stop hct service and release ccp resources
	stop

	#Rebind CCP by given parameters;  \"ccp_type\" can be \"both\", \"ntb\", \"psp\"
	rebind [-n --ccp_num [-t --ccp_type] [-p --max_prog]

	#Show hct info
	status

	#Show this message
	usage
"

}

unbind_ccp() {

	MDEV_PATH="/sys/bus/mdev/devices/"
	MDEVS=$(ls $MDEV_PATH)
	for MDEV in ${MDEVS} ; do
		echo 1 > /sys/bus/mdev/devices/${MDEV}/remove
	done

	DEVS=(`${HCT_SCRIPT_DIR}/hct_ccp_bind.py -s | sed -n '/CCP/p' | sed -n '/drv=hct/p' | awk '{print $1}'`)
	for DEV in ${DEVS[*]} ; do
		${HCT_SCRIPT_DIR}/hct_ccp_bind.py --unbind ${DEV}
	done
}

bind_ccp() {

	NTB_DEVS=(`${HCT_SCRIPT_DIR}/hct_ccp_bind.py -s | sed -n '/CCP/p' | sed -n '/hct/p' | grep "NTB" | awk '{print $1}'`)
	PSP_DEVS=(`${HCT_SCRIPT_DIR}/hct_ccp_bind.py -s | sed -n '/CCP/p' | sed -n '/hct/p' | grep "PSP" | awk '{print $1}'`)
	CCP_DEVS=(`${HCT_SCRIPT_DIR}/hct_ccp_bind.py -s | sed -n '/CCP/p' | sed -n '/hct/p' | awk '{print $1}'`)

	#count system ntb ccp amount
	SYS_NTB_NUM=0
	for TPM_DEV in ${NTB_DEVS[*]} ; do
		let SYS_NTB_NUM++
	done

	#count system psp ccp amount
	SYS_PSP_NUM=0
	for TPM_DEV in ${PSP_DEVS[*]} ; do
		let SYS_PSP_NUM++
	done

	let SYS_PSP_NUM--

	case $CCP_TYPE in
	"both")
		if [ $CCP_NUM == "all" ]; then
			for TPM_DEV in ${CCP_DEVS[*]} ; do
				# fTPM issue
				if [ ${TPM_DEV} = "${PSP_DEVS[0]}" ]; then
						continue;
				fi
				${HCT_SCRIPT_DIR}/hct_ccp_bind.py --bind=hct ${TPM_DEV}
			done
		elif [ $(( $SYS_PSP_NUM + $SYS_NTB_NUM )) -ge $CCP_NUM ]; then
			local i=0
			for TPM_DEV in ${CCP_DEVS[*]} ; do
				if [ ${TPM_DEV} = "${PSP_DEVS[0]}" ]; then
						continue;
				fi
				let i++
				if [ $i -le $CCP_NUM ]; then
					${HCT_SCRIPT_DIR}/hct_ccp_bind.py --bind=hct ${TPM_DEV}
				else
					break
				fi
			done
		else
			echo "Error: Total available CCP amount for hct  is $(( $SYS_PSP_NUM + $SYS_NTB_NUM )) , ccp_num paramter:$CCP_NUM is too large or illegal!"
			exit -1
		fi
		;;
	"ntb")
		if [ $CCP_NUM == "all" ]; then
			for TPM_DEV in ${NTB_DEVS[*]} ; do
				${HCT_SCRIPT_DIR}/hct_ccp_bind.py --bind=hct ${TPM_DEV}
			done
		elif [ $SYS_NTB_NUM -ge $CCP_NUM ]; then
			local i=0
			for TPM_DEV in ${NTB_DEVS[*]} ; do
				let i++
				if [ $i -le $CCP_NUM ]; then
					${HCT_SCRIPT_DIR}/hct_ccp_bind.py --bind=hct ${TPM_DEV}
				else
					break
				fi
			done
		else
			echo "Error: Total NTB CCP amount is $SYS_NTB_NUM , ccp_num paramter:$CCP_NUM is too large or illegal!"
			exit -1
		fi
		;;
	"psp")
		if [ $CCP_NUM == "all" ]; then
			for TPM_DEV in ${PSP_DEVS[*]} ; do
				if [ ${TPM_DEV} = "${PSP_DEVS[0]}" ]; then
						continue;
				fi
				${HCT_SCRIPT_DIR}/hct_ccp_bind.py --bind=hct ${TPM_DEV}
			done
		elif [ $SYS_PSP_NUM -ge $CCP_NUM ]; then
			local i=0
			for TPM_DEV in ${PSP_DEVS[*]} ; do
				if [ ${TPM_DEV} = "${PSP_DEVS[0]}" ]; then
						continue;
				fi
				let i++
				if [ $i -le $CCP_NUM ]; then
					${HCT_SCRIPT_DIR}/hct_ccp_bind.py --bind=hct ${TPM_DEV}
				else
					break
				fi
			done
		else
			echo "Error: Total PSP CCP amount is $SYS_PSP_NUM , ccp_num paramter:$CCP_NUM is too large or illegal!"
			exit -1
		fi
		;;
	*)
	usage
		exit 7
		;;
	esac

	#create mdev per progress
	EFFECT_CCP_NUM=`${HCT_SCRIPT_DIR}/hct_ccp_bind.py -s | sed -n '/drv=hct/p' | grep "CCP" | awk 'END {print NR}'`
	echo "${EFFECT_CCP_NUM} CCP devices are bound with hct driver"

	MDEV_NUM=$((${EFFECT_CCP_NUM}*${MAX_PROGRESS}))
	for i in $(seq 1 ${MDEV_NUM}) ; do
				echo `uuidgen` > /sys/devices/virtual/hct/hct/mdev_supported_types/hct-1/create
	done
	echo "${MDEV_NUM} mdev devices are created to support max ${MAX_PROGRESS} progress"
}


#-------------------------------------
# service start function
#-------------------------------------
start() {
	_getopt "$@"

	modprobe vfio | modprobe vfio-pci | modprobe vfio_iommu_type1 | modprobe mdev | modprobe vfio_mdev | modprobe hct
	if [ $? -ne 0 ]; then
		echo "lack of necessary kernel modules!"
		echo "
Tips: run below commands to check
	  $ sudo modprobe vfio
	  $ sudo modprobe vfio-pci
	  $ sudo modprobe vfio_iommu_type1
	  $ sudo modprobe mdev
	  $ sudo modprobe vfio_mdev
	  $ sudo modprobe hct

"
		exit 8
	fi

	# unbind ccp first
	unbind_ccp

	# bind ccp now
	bind_ccp
}

#-------------------------------------
# service stop function
#-------------------------------------
stop() {
	echo "Stopping HCT service and release CCPs..."
	unbind_ccp
	rmmod hct
}

#-------------------------------------
# service restart function
#-------------------------------------
restart() {
	stop
	start $*
}


#-------------------------------------
# service status function
#-------------------------------------
status() {
	if [ -d `echo /sys/devices/virtual/hct/hct/mdev_supported_types/hct-1/devices` ];then
		MDEV_NUM=`ls -l /sys/devices/virtual/hct/hct/mdev_supported_types/hct-1/devices | grep -v "total" | wc -l`
		EFFECT_CCP_NUM=`${HCT_SCRIPT_DIR}/hct_ccp_bind.py -s | sed -n '/drv=hct/p' | grep "CCP" | awk 'END {print NR}'`
		MAX_PROGRESS=`expr ${MDEV_NUM} / ${EFFECT_CCP_NUM}`
	else
		MDEV_NUM=0
		MAX_PROGRESS=0
	fi
	echo "------------>>>  H C T	I N F O	D U M P  <<<--------------"
	echo ""
	printf "HCT Version:			%s\n"  $VERSION
	printf "HCT engine Dir:			%s\n"  $ENGINE_DIR
	printf "HCT lib Dir:			%s\n"  $HCT_LIB_DIR
	printf "HCT script Dir:			%s\n"  $HCT_SCRIPT_DIR
	printf "HCT function test tool Dir:	%s\n"  $HCT_TEST_DIR
	printf "HCT performance tool Dir:	%s\n"  $HCT_SPEED_DIR
	printf "HCT created %s mdev devices to support max %s progress \n"  $MDEV_NUM $MAX_PROGRESS
	echo ""
	printf "HCT CCP effective Binding info :\n"
	${HCT_SCRIPT_DIR}/hct_ccp_bind.py -s | sed -n '/CCP/p'
	echo ""
	echo "-----------------------------------------------------------------------------"
	echo ""


	printf "\n\n"
}

#-------------------------------------
# service rebind function
#-------------------------------------
rebind() {
	start $*
}


#==========================
# service entry
#==========================
process() {
	case $1 in
	start)
		shift 1
		start $*
		;;
	stop)
		stop
		;;
	restart)
		shift 1
		restart $*
		;;
	rebind)
		shift 1
		rebind $*
		;;
	status)
		status
		;;
	*)
		usage
		exit 7
		;;
	esac
}

main() {
	if [ $# -lt 1 ]; then
		start ;
		exit 1
	fi

	process "$@"
}

main $*
