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

VERSION=2.0.0-2024-0430-release

#max supported number of processes
MAX_PROGRESS="128"

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

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

#DRIVER_MOD value can be: hct | vfio-pci
DRIVER_MOD="hct"

#hard file limit for process. "0" won't make any change.
HARD_FILE_LIMIT=0

#-------> 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/"

#------->>> NTBCCP or/and PSPCCP deviceid <<<---------
NTBCCP="1468\|14b8"
PSPCCP="1456\|1486\|14a6\|1537"
CCPALL=${NTBCCP}"\|"${PSPCCP}

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

#-------------------------------------
# get command opts
#-------------------------------------
_getopt() {
	ARGS=`getopt -o n:t:d:p:m:f: --long ccp_num:,ccp_type:,script_dir:,max_prog:,driver_mod:,max_file_limit: -- "$@"`
	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
			;;
		-m|--driver_mod)
			DRIVER_MOD=$2;
			shift 2
			;;
		-f|--max_file_limit)
			HARD_FILE_LIMIT=$2;
			shift 2
			;;
		--)
			shift
			break
			;;
		*)
			echo "Parameter error, $1"
			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;
	start [-n --ccp_num] [-t --ccp_type] [-p --max_prog] [-f --max_file_limit]
		\"ccp_num\" can be set to a number which no more than the number of CCP, or \"all\", \"default\". If set \"ccp_num\" to \"all\", please disable fTPM in BIOS at first.
		\"ccp_type\" can be \"both\", \"ntb\", \"psp\"

	#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
"

}

# check and update the hard file limit in limits.conf
update_fd_limit() {
	# backup old config
	current_time=$(date "+%Y%m%d-%H%M%S")
	limit_file="/etc/security/limits.conf"
	backup_file="${limit_file}.bak-$current_time"
	cp "$limit_file" "$backup_file"

	local root_pattern="^root\s+hard\s+nofile\s+([0-9]+)"
	local general_pattern="^\s*hard\s+nofile\s+([0-9]+)$"

	local new_limit_file="${limit_file}.new"
	: > "$new_limit_file"

	root_configured=false
	general_configured=false

	while IFS= read -r line; do
		# check file limit for root
		if [[ $line =~ $root_pattern ]]; then
			local current_limit="${BASH_REMATCH[1]}"
			if [[ "$current_limit" =~ ^[0-9]+$ ]] && [ "$current_limit" -lt $HARD_FILE_LIMIT ]; then
				echo -e "#$line" >> "$new_limit_file"
				echo -e "root hard nofile $HARD_FILE_LIMIT" >> "$new_limit_file"
			else
				echo -e "$line" >> "$new_limit_file"
			fi
			root_configured=true
		# check file limit for general user
		elif [[ $line =~ $general_pattern ]]; then
			local current_limit="${BASH_REMATCH[1]}"
			if [[ "$current_limit" =~ ^[0-9]+$ ]] && [ "$current_limit" -lt $HARD_FILE_LIMIT ]; then
				echo -e "#$line" >> "$new_limit_file"
				echo -e "hard nofile $HARD_FILE_LIMIT" >> "$new_limit_file"
			else
				echo -e "$line" >> "$new_limit_file"
			fi
			general_configured=true
		else
			echo -e "$line" >> "$new_limit_file"
		fi
	done < "$limit_file"

	if ! $root_configured; then
		echo -e "root hard nofile $HARD_FILE_LIMIT" >> "$new_limit_file"
	fi
	if ! $general_configured; then
		echo -e "hard nofile $HARD_FILE_LIMIT" >> "$new_limit_file"
	fi
	mv "$new_limit_file" "$limit_file"
	echo "The backup of the original limits.conf has been created at $backup_file. Relogin shell or execute 'ulimit -Hn' to set file limit"
}

unbind_ccp() {

	MDEV_PATH="/sys/bus/mdev/devices/"

	if [ -d $MDEV_PATH ]; then
		MDEVS=$(ls $MDEV_PATH)
		for MDEV in ${MDEVS} ; do
			echo 1 > /sys/bus/mdev/devices/${MDEV}/remove
		done
	fi

	DEVS=(`${HCT_SCRIPT_DIR}/hct_ccp_bind.py -s | grep ${CCPALL} | grep -E "drv=hct|drv=vfio-pci" | 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 | grep ${NTBCCP} | awk '{print $1}'`)
	PSP_DEVS=(`${HCT_SCRIPT_DIR}/hct_ccp_bind.py -s | grep ${PSPCCP} | awk '{print $1}'`)
	CCP_DEVS=(`${HCT_SCRIPT_DIR}/hct_ccp_bind.py -s | grep ${CCPALL} | 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

	if [ "$CCP_NUM" != "all" ]; then
		if [ "$CCP_TYPE" != "psp" ] && ([ "$CCP_NUM" == "default" ] || [ "$CCP_NUM" -lt $(( $SYS_PSP_NUM + $SYS_NTB_NUM )) ]); then
			let SYS_PSP_NUM--
		fi
		if [ "$CCP_TYPE" == "psp" ] && ([ "$CCP_NUM" == "default" ] || [ "$CCP_NUM" -lt $(( $SYS_PSP_NUM )) ]); then
			let SYS_PSP_NUM--
		fi
	fi

	case $CCP_TYPE in
	"both")
		if [ $CCP_NUM == "all" ]; then
			for TPM_DEV in ${CCP_DEVS[*]} ; do
				${HCT_SCRIPT_DIR}/hct_ccp_bind.py --bind=$DRIVER_MOD ${TPM_DEV}
			done
		elif [ $CCP_NUM == "default" ]; 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=$DRIVER_MOD ${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 [ $CCP_NUM -lt $(( $SYS_PSP_NUM + $SYS_NTB_NUM )) ] && [ ${TPM_DEV} = "${PSP_DEVS[0]}" ]; then
						continue;
				fi
				let i++
				if [ $i -le $CCP_NUM ]; then
					${HCT_SCRIPT_DIR}/hct_ccp_bind.py --bind=$DRIVER_MOD ${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" ] || [ $CCP_NUM == "default" ]; then
			for TPM_DEV in ${NTB_DEVS[*]} ; do
				${HCT_SCRIPT_DIR}/hct_ccp_bind.py --bind=$DRIVER_MOD ${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=$DRIVER_MOD ${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 == "default" ]; 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=$DRIVER_MOD ${TPM_DEV}
			done
		elif [ $CCP_NUM == "all" ]; then
			for TPM_DEV in ${PSP_DEVS[*]} ; do
				${HCT_SCRIPT_DIR}/hct_ccp_bind.py --bind=$DRIVER_MOD ${TPM_DEV}
			done
		elif [ $SYS_PSP_NUM -ge $CCP_NUM ]; then
			local i=0
			for TPM_DEV in ${PSP_DEVS[*]} ; do
				if [ $SYS_PSP_NUM -gt $CCP_NUM ] && [ ${TPM_DEV} = "${PSP_DEVS[0]}" ]; then
						continue;
				fi
				let i++
				if [ $i -le $CCP_NUM ]; then
					${HCT_SCRIPT_DIR}/hct_ccp_bind.py --bind=$DRIVER_MOD ${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

	if [ ' '$DRIVER_MOD == " hct" ]; then
		EFFECT_CCP_NUM=`${HCT_SCRIPT_DIR}/hct_ccp_bind.py -s | grep ${CCPALL} | grep "drv=hct" | awk 'END {print NR}'`
		#create mdev per progress
		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"
	else
		EFFECT_CCP_NUM=`${HCT_SCRIPT_DIR}/hct_ccp_bind.py -s | grep ${CCPALL} | grep "drv=vfio-pci" | awk 'END {print NR}'`
		echo "${EFFECT_CCP_NUM} CCP devices are bound with vfio-pci driver"
	fi

	#chmod /dev/vfio/xx[group_id] to user permissions
	VFIO_PATH=/dev/vfio
	for group_id in `ls ${VFIO_PATH}`;do
		if [ ' '${group_id} != " vfio" ]; then
			chmod go+rw ${VFIO_PATH}/${group_id}
		fi
	done

	# chmod pci resource2 to user permissions
	for TPM_DEV in ${PSP_DEVS[*]} ; do
		chmod go+r /sys/bus/pci/devices/${TPM_DEV}/resource2
	done

	# enable master psp's pci to support the hgsc checking
	if [[ ${PSP_DEVS[0]} != "" ]]; then
		master_psp_status=`cat /sys/bus/pci/devices/${PSP_DEVS[0]}/enable`
		if [ $master_psp_status -eq 0 ]; then
			echo 1 > /sys/bus/pci/devices/${PSP_DEVS[0]}/enable
		fi
	fi
}


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

	if [ ' '$DRIVER_MOD != " hct" ] && [ ' '$DRIVER_MOD != " vfio-pci" ]; then
		echo "driver must be 'hct' or 'vfio-pci'"
		exit 8
	fi

	modprobe vfio | modprobe vfio-pci | modprobe vfio_iommu_type1
	if [ $? -ne 0 ]; then
		echo "lack of basic kernel modules(vfio/vfio-pci/vfio_iommu_type1)!"
		exit 8
	fi

	if [ ' '$DRIVER_MOD == " hct" ]; then
                KER_VER_MJR=`uname -a | awk '{print $3}' | awk -F '.'  '{print $1}'`
                KER_VER_MIN=`uname -a | awk '{print $3}' | awk -F '.'  '{print $2}'`
                if { [ $KER_VER_MJR -eq  5 ] && [ $KER_VER_MIN -ge 16 ]; } || [ $KER_VER_MJR -ge 6 ]; then
                    modprobe mdev | modprobe hct
                else
                    modprobe mdev | modprobe vfio_mdev | modprobe hct
                fi

		if [ $? -ne 0 ]; then
			echo "lack of necessary kernel modules(mdev/vfio_mdev/hct)!"
			exit 8
		fi
	fi

	# unbind ccp first
	unbind_ccp

	# bind ccp now
	bind_ccp

	if [ $HARD_FILE_LIMIT -gt 0 ]; then
		update_fd_limit
	fi
}

#-------------------------------------
# service stop function
#-------------------------------------
stop() {
	echo "Release CCPs..."
	unbind_ccp

	hct_info=`lsmod |grep hct | awk 'END {print NR}'`
	if [ $hct_info -ne 0 ]; then
		rmmod hct
	fi
}

#-------------------------------------
# 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 ${CCPALL} | awk 'END {print NR}'`
		MAX_PROGRESS=`expr ${MDEV_NUM} / ${EFFECT_CCP_NUM}`
	else
		EFFECT_CCP_NUM=`${HCT_SCRIPT_DIR}/hct_ccp_bind.py -s | sed -n '/drv=vfio-pci/p' | grep ${CCPALL} | awk 'END {print NR}'`
		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
	if [ $MDEV_NUM -ne 0 ]; then
		printf "HCT created %s mdev devices to support max %s progress \n"  $MDEV_NUM $MAX_PROGRESS
	else
		printf "Found %s CCP devices are bound with vfio-pci driver\n" $EFFECT_CCP_NUM
	fi
	echo ""
	printf "HCT CCP effective Binding info :\n"
	${HCT_SCRIPT_DIR}/hct_ccp_bind.py -s | grep ${CCPALL}
	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 $*
