#!/bin/ksh93

MAPFILE='/lib/udev/baymapNvme'
RULEFILE='/lib/udev/rules.d/61-baylinksNvme.rules'
DEVPREFIX='chassis'
RAMFSHOOK='/usr/share/initramfs-tools/hooks/baymapsNvme'
ROOTFS=
PRODUCT=
BAYLABEL_PREFIX='SSD'
MAPSCRIPT='/lib/udev/bayLinksNvme.sh'

function getSmBios {
	typeset -n TXT=$1
	if [[ -n ${BIOS} ]]; then
		if [[ ! -e ${BIOS} ]]; then
			print -u2 "ERROR: dmidecode output file '${BIOS}' not found"
			return 1
		fi
		TXT=$(<${BIOS})
		[[ -z ${TXT} ]] && return 2
		return 0
	fi

	DMIDECODE=${ whence dmidecode ; }
	[[ -z ${DMIDECODE} && -x /usr/sbin/dmidecode ]] && \
		DMIDECODE=/usr/sbin/dmidecode
	[[ -z ${DMIDECODE} && -x /sbin/dmidecode ]] && DMIDECODE=/sbin/dmidecode
	[[ -z ${DMIDECODE} ]] && print -u2 "ERROR: dmidecode node found" && return 3
	TXT=${ ${DMIDECODE} -t 9 ; }
}

# $1 .. addr to analyze
# $2 .. where to store the result
# $3 .. SSD slots assigne to the related AOC card
function addAddrMap {
	typeset -n ADDR=$1
	typeset -n MAP=$2
	typeset -n SLOTS=$3
	[[ -z ${ADDR} ]] && return

	typeset -i16 H=16#${ADDR:5:2}
	typeset A=( ${SLOTS} )	ADDR_PREFIX=${ADDR:0:4} ADDR_SUFFIX='00' P S

	if (( ${#A[@]} == 2 )); then
		 (( (H & 1) == 0 )) && (( H-- ))
	elif (( H % 4 != 1 )); then
		(( H % 4 == 0 )) && (( H -= 3 )) || (( H = (H & 16#fc) + 1 ))
	fi
	for S in ${SLOTS[@]} ; do
		(( H < 10 )) && P='0' || P=
		MAP["${ADDR_PREFIX}:${P}${H:3}:${ADDR_SUFFIX}"]="$S"
		(( H++ ))
	done
}

function map_AS2124USTNRP {
	typeset SMBIOS
	typeset BSLOT= ADDR D E K V
	typeset -n MAP=$1
	typeset -i16 BUS

	[[ -f $F ]] || F=

	getSmBios SMBIOS || return 1

	print "${SMBIOS}" | while read K V ; do
		if [[ $K == 'Designation:' ]]; then
			if [[ ${V:0:22} == 'AOC-2UR68G4-i4XTS Slot' ]]; then
				V=${V:22:2}
				if [[ $V == '1 ' ]]; then
					BSLOT='10 11'
				elif [[ $V == '2 ' ]]; then
					BSLOT='4 5 6 7'
				elif [[ $V == '3 ' ]]; then
					BSLOT='8 9'
				else
					BSLOT=
				fi
			elif  [[ ${V:0:17} == 'RSC-W2-66G4 SLOT2' ]]; then
				BSLOT='16 17 18 19'
			elif [[ ${V:0:14} == 'RSC-WR-6 SLOT1' ]]; then
				BSLOT='20 21 22 23'
			else
				BSLOT=
			fi
		fi
		[[ -z ${BSLOT} ]] && continue
		if [[ $K == 'Bus' && ${V:0:9} == 'Address: ' ]]; then
			ADDR=
			D="/sys/bus/pci/drivers/pcieport/${V:9}/pci_bus"
			if [[ -d $D ]]; then
				# occures if all NVME ports of the related riser card are empty
				cd $D
				for E in ~(N)0* ; do
					[[ -d $E ]] && ADDR="${E}:00" && break
				done
				cd ~-
			else
				ADDR=${V:9}
			fi
			addAddrMap ADDR MAP BSLOT
			BSLOT=
		fi
	done
	if (( 0 )); then
		D=
		for V in ${!MAP[@]}; do
			D+="$V\t${MAP[$V]}\n"
		done
		print "Riser cards analyzed:\n$D"
	fi

	# now we need to find out P1_NVME0/1 + P1_NVME2/3 and
	# P2_NVME0 + P2_NVME1 + P2_NVME2/3 are mapped to. 
	# all P1_* are below pci0000:80, all P2_* above pci0000:80
	for V in ~(N)/sys/bus/pci/slots/*/address ; do
		ADDR=$(<$V)
		[[ -n ${MAP["${ADDR}"]} ]] && continue
		# if not yet found, it must be a P*_NVMEn assuming no other additional
		# PCI-NVMe cards (e.g. M.2 based) are installed.
		BUS=16#${ADDR:5:2}
		(( BUS < 16#80 )) && BSLOT='0 1 2 3' || BSLOT='12 13 14 15'
		addAddrMap ADDR MAP BSLOT
	done
	# set -s ${!MAP[@]} ; for V ; do print "$V\t${MAP[$V]}" ; done
}

function showUsage {
	typeset WHAT="$1" X='--man'
	[[ -z ${WHAT} ]] && WHAT='MAIN' && X='-?'
    getopts -a "${PROG}" "${ print ${USAGE}; }" OPT $X
}

function initrdNotice {
	print '
If you are booting via initrd, you should update the related image. E.g.:
	update-initramfs -k all -u
'
}

function installHook {
	HOOK='#!/bin/sh -e
PREREQS="udev"

# Since this stuff gets sourced in rather than execed by a very brain damaged,
# limited and buggy shell, we need to use its poor, bloated syntax and can not
# make use of advanced features provides by ksh93 ... :((((((((
MAPSCRIPT="'"${MAPSCRIPT}"'"

# Just to avoid copy over itself
if [ -z "${DESTDIR}" ]; then
	DESTDIR=/tmp
fi

if [ -n "$1" -a "$1" = "prereqs" ]; then
	echo "${PREREQS}"
    exit 0
fi

[ ! -e "${MAPSCRIPT}" ] && return

. /usr/share/initramfs-tools/hook-functions

while read LINE ; do
	if [ "${LINE#MAPFILE=}" != "${LINE}" ]; then
		MAPFILE="${LINE#MAPFILE=?}"
		MAPFILE="${MAPFILE%?}"
	elif [ "${LINE#RULEFILE=}" != "${LINE}" ]; then
		RULEFILE="${LINE#RULEFILE=?}"
		RULEFILE="${RULEFILE%?}"
		break
	fi
done < "${MAPSCRIPT}"

if [ -z "${MAPFILE}" -o -z "${RULEFILE}" ] ; then
	echo "${MAPSCRIPT} script does not contain MAPFILE and RULEFILE settings!"
	exit 1
fi

cp -p "${MAPSCRIPT}" "${DESTDIR}/lib/udev/"
if [ -e "${RULEFILE}" ]; then
	cp -pL "${RULEFILE}" "${DESTDIR}/lib/udev/rules.d/"
fi
if [ -e "${MAPFILE}" ]; then
	cp -pL "${MAPFILE}" "${DESTDIR}/lib/udev/"
fi
[ -e /bin/ksh93 ] || copy_exec /bin/ksh93 /bin
'
	if (( DRY )); then
		print "=== schnipp ===\n${HOOK}\n=== schnapp ==="
		print "Hook would be installed as '${ROOTFS}${RAMFSHOOK}'"
	else
		print "${HOOK}" >"${ROOTFS}${RAMFSHOOK}" && \
			print "Hook installed as '${ROOTFS}${RAMFSHOOK}'"
		chmod 755 "${ROOTFS}${RAMFSHOOK}"
	fi
}

function makeRule {
	typeset RULE='# NVMe PCI address <-> HDD label mapper
SUBSYSTEM!="block", GOTO="baylabel_end"
KERNEL!="nvme?*", GOTO="baylabel_end"

# NOTE: ACTION=="add" is not sufficient, because a simple "parted -lm" would
# trigger a change event, which would remove the symlink, if it is not added
# again by the script! "disk" and "partition" makes also sure, that we do not
# get the nvme controller itself.
ENV{DEVTYPE}=="disk", PROGRAM="bayLinksNvme.sh $devpath", RESULT=="?*", SYMLINK+="%c{1}"
# Dueto limitations of initramfs wrt. zfs as well as zfs its own pecularities,
# we have to flatten the structure
ACTION!="remove", ENV{DEVTYPE}=="partition", PROGRAM="bayLinksNvme.sh $devpath", RESULT=="?*", SYMLINK+="%c{1}p%n %c{2}p%n"

LABEL="baylabel_end"
'
	if (( DRY )); then
		print "=== schnipp ===\n${RULE}\n=== schnapp ==="
		print "Rule would be installed as '${ROOTFS}${RULEFILE}'"
	else
		print -n "${RULE}" >"${ROOTFS}${RULEFILE}" && \
			print "Rule installed as '${ROOTFS}${RULEFILE}'"
	fi
	initrdNotice
}

typeset -A ADDR2SLOT ADDR2NVME ADDR2BAY

function getProduct {
	[[ -n ${PRODUCT} ]] && return
	typeset F='/sys/class/dmi/id/product_name' RAW= NAME= K V M
	integer I

	if [[ -n ${BIOS} ]]; then
		I=0
		while read K V M ; do
			if (( I )); then
				[[ $K == 'Product' && $V == 'Name:' ]] && RAW="$M" && break
			fi
			[[ $K == 'System' && $V == 'Information' ]] && I=1
		done<${BIOS}
	elif [[ -f $F ]]; then
		RAW=$(<$F)
	fi
	for (( I=0; I < ${#RAW}; I++ )); do
		[[ ${RAW:I:1} =~ ^[0-9a-zA-Z]$ ]] && NAME+="${RAW:I:1}"
	done
	PRODUCT="${NAME}"
}

function populateBaymap {
	typeset F
	if [[ -z "$1" ]]; then
		if [[ -f ${MAPFILE} ]]; then
			F=$(<${MAPFILE})
		else
			print -u2 "Unknown product and no '${MAPFILE}' -> bogus mapping."
		fi
	else
		typeset -n F=$1
	fi
	typeset K V
	print "$F" | while read  V K ; do
		[[ $V =~ ^[A-Z]+:[0-9a-f]+:[0-9a-f]+:[0-9a-f]+=[0-9a-f]+$ ]] || continue
		K=${V%=*}
		V=${.sh.match:1}
		ADDR2BAY["$K"]="$V"
	done
}

function makeMap {
	typeset -A NMAP  SEEN
	typeset MAP DEVPATH DEVNAME ADDR B E MINOR
	integer N=0

	MAP="# ${MAPFILE}\n\n"
	MAP+="# Generated by ${FPROG} on ${ date ; }\n"
	MAP+='# Format: CONTAINER_LABEL:PCI_ADDRESS=BAY_NUM\t# COMMENT\n'
	for F in ~(N)/dev/nvme[0-9]* ; do
		DEVPATH= DEVTYPE= MINOR= DEVNAME=
		udevadm info $F | while read K V ; do
			[[ $K == 'N:' ]] && DEVNAME="$V" && continue
			[[ ${V:0:8} == 'DEVPATH=' ]] && DEVPATH="${V:8}" && continue
			[[ ${V:0:8} == 'DEVTYPE=' ]] && DEVTYPE="${V:8}" && continue
			[[ ${V:0:6} == 'MINOR=' ]] && MINOR="${V:6}"
		done
		# only interested in info for DEVTYPE=disk
		[[ -n ${MINOR} && ${DEVTYPE} == 'disk' ]] || continue
		# e.g. /devices/pci0000:f0/0000:f0:03.2/0000:f2:00.0/nvme/nvme9/nvme9n1
		A=( ${DEVPATH//\// } )
		(( ${#A[@]} != 7 )) && continue
		ADDR="${A[3]%.*}"				# e.g.: 0000:f2:00.0
		ADDR2BAY["SYS:${ADDR}"]="${MINOR}"		# just an enumerator
		NMAP["SYS:${ADDR}"]="${DEVNAME}"		# e.g. /dev/nvme9
	done

	getProduct
	E=${ whence -v map_${PRODUCT} 2>/dev/null ; }
	if [[ ${E: -9:9} == ' function' ]]; then
		print -u2 "Using product specific map ${PRODUCT} ..."
		typeset -A RMAP
		map_${PRODUCT} RMAP || return 1
		set -s ${!RMAP[@]}
		for E ; do
			B="${RMAP[$E]}"
			[[ -n ${SEEN["_$B"]} ]] && \
				MAP+='# Duplicate - next line disabled\n#' || SEEN["_$B"]=1
			MAP+="SYS:${E}=$B\t# ${NMAP[SYS:$E]}\n"
			(( N++ ))
		done
	else
		set -s ${!ADDR2BAY[@]}
		for E ; do
			B="${ADDR2BAY[$E]}"
			[[ -n ${SEEN["_$B"]} ]] && \
				MAP+='# Duplicate - next line disabled\n#' || SEEN["_$B"]=1
			MAP+="${E}=$B\t# ${NMAP[$E]}\n"
			(( N++ ))
		done
	fi
	(( N == 0 )) && UDEV=0 && INSTHOOK=0 && return 0

	E='
Please check and adjust the BAY_NUM as needed. If not all slots are populated,
some mappings might be missing or mappings might be wrong.'

	if (( DRY )); then
		print "=== schnipp ===\n${MAP}\n=== schnapp ===\n"
		print "Baymap would be installed as '${ROOTFS}${MAPFILE}'$E"
	elif print -n "${MAP}" >"${ROOTFS}${MAPFILE}" ; then
		print "Baymap installed as '${ROOTFS}${MAPFILE}'$E"
	else
		print "\nBaymap NOT installed!"
		return
	fi
	initrdNotice
}

function prepareMaps {
	(( ${#ADDR2SLOT[@]} )) && return

	typeset SLOT ADDR F
	for F in ~(N)/sys/bus/pci/slots/*/address ; do
		SLOT=${F%/*}
		[[ -e ${SLOT}/power ]] || continue	# not managed by a real NVME ctrl
		ADDR=$(<$F)
		ADDR2SLOT["${ADDR}"]="${SLOT##*/}"
	done
	for F in ~(N)/sys/class/nvme/nvme* ; do
		ADDR=$(<$F/address)
		ADDR=${ADDR%.*}
		ADDR2NVME["${ADDR}"]="${F##*/}"
	done

	# 
	if [[ -f ${ROOTFS}${MAPFILE} ]]; then
		populateBaymap
	else
		getProduct
		if [[ ${PRODUCT} == 'AS2124USTNRP' ]]; then
			populateBaymap AS2124USTNRP
		else
			makeMap
		fi
	fi
}

function listMap {
	prepareMaps
	typeset A ADDR
	print "Slot\tAddress   \tBay"
	for ADDR in ${!ADDR2BAY[@]}; do
		A=${ADDR#*:}
		printf "%5s\t%10s\t%3s\t%s\n" "${ADDR2SLOT[$A]}" "$A" \
			"${ADDR2BAY[$ADDR]}" "${ADDR2NVME[$A]}"
	done | sort -k3,3n
	print "\nUse 'echo 0 >/sys/bus/pci/slots/\${Slot}/power to power off SSD."
}

function lookup {
	integer DEBUG=1
	prepareMaps

	if (( DEBUG )); then
		typeset D='/dev/.initramfs' T
		[ -d $D ] || D='/tmp'
		T="$D/${ printf 'baymap-%(%s)T' now ; }.lst"
		set >"$T" || true
		exec 2>"$T".stderr
	fi
	# nothing to do ?
	[[ -z $1 ]] && return 0

	typeset DEVPATH="$1" A S BOX BAY ADDR
	# env example:
	# DEVPATH=/devices/pci0000:60/0000:60:03.1/0000:61:00.0/nvme/nvme0/nvme0n1
	# ID_PATH=pci-0000:61:00.0-nvme-1
	# DEVTYPE=disk
	# DEVNAME=/dev/nvme0n1

	A=( ${DEVPATH//\// } )
	if [[ ${DEVTYPE} == 'partition' ]]; then
		(( ${#A[@]} != 8 )) && return 0
	elif (( ${#A[@]} != 7 )) ; then
		return 0	# DEVTYPE disk
	fi
	[[ ${A[4]} == 'nvme' ]] || return
	S="${A[3]%.*}"
	for ADDR in ${!ADDR2BAY[@]}; do
		A=${ADDR#*:}
		BOX=${.sh.match%:}
		[[ $A == $S ]] || continue
		BAY="${ADDR2BAY[${ADDR}]}"
		if [[ -n ${BAY} ]]; then
			BAY="${BAYLABEL_PREFIX}${BAY}"
			print "${DEVPREFIX}/${BOX}/${BAY} ${BAY}"
			(( DEBUG )) && print "${DEVPREFIX}/${BOX}/${BAY} ${BAY}" \
				>>"$T" || true
			break
		fi
	done
	return 0
}

function fireUdev {
	typeset DEVS=${ ls -1 /dev/nvme[0-9]?([0-9])n1 ; } F
	[[ -z ${DEVS} ]] && return 0
	for D in ${DEVS} ; do
		udevadm trigger --action=add $D
	done
}

integer MKMAP=0 MKRULE=0 DRY=0 INSTHOOK=0 UDEV=0 LIST=0 IDX BUGGY=0
typeset -r FPROG=${.sh.file} PROG=${FPROG##*/}
typeset BIOS=

USAGE='[-?$Id$ ]
[-copyright?Copyright (c) 2024 Jens Elkner. All rights reserved.]
[-license?CDDL 1.0]
[+NAME?'"${PROG}"' - manage nvme adresses wrt. bay alias SDD labels]
[+DESCRIPTION?This script can be used on \bLinux\b by \budev\b(7) to create \b/dev/'"${DEVPREFIX}"'/\b\abox\a/\abay\a\b/disk\b links to the related \b/dev/nvme\b\a?\a to obtain a \bstable naming\b. This is required, because Linux assigns device names on a first-come, first-served basis, and thus e.g. /dev/nvme1 might be currently assigned to the NVME in bay 0 and on the next boot it might be assigned to the NVME in bay 7. Also it helps an admin to easily determine, which device is located in which bay and avoids pulling the wrong disk, etc. accidents.]
[+?Since NVME device are basically PCIe devices, they get accessed via PCIe bus:function pair and the related PCIe cable is usually directly connected to the NVME bay backplane of the server and never gets changed. So if one moves the NVME controller to a different PCIe slot or someone connects the related cables to different sockets on the backplane, one probably needs to adjust the baymap, too.]
[+?If no option is given, the script assumes, that it got called by \budev\b(7). In this case the operand \bdevpath\b as well as the baymap file \b'"${MAPFILE}"'\b are required. They are used to lookup the related mapping and spit out the symlink to add by udev. The baymap file can be produced with this script. Because not all nvme controllers might be initialized when the event for a certain disk gets fired, the mapping cannot be generated on the fly.]
[+FILES?]{
	[+'"${MAPSCRIPT}"'?Where this file should be copied to make bay mapping work.]
}
[h:help?Print this help and exit.]
[F:functions?Print a list of all functions of this script.]
[T:trace]:[functionList?A comma separated list of functions of this script to trace (convinience for troubleshooting).] 
[+?]
[D:destdir]:[path?The directory to use as virtual root when writing/reading files except those beneath /dev, /sys and /proc.]
[b:bios]:[path?Use the file \apath\a as output of dmidecode (smbios data) instead of invoking \bdmidecode\b(8) directly, which would usually require root privileges to read /dev/mem and /sys/firmware/dmi/tables/smbios_entry_point. The BIOS output for type 9 is required when using the map command (-m) for a Supermicro AS-2124US-TNRP, otherwise ignored.]
[l:list?List current mapping and exit. This is useful to find easily the NVME related PCIe slot for hot-plug.]
[m:map?Generate the baymap file to be used by this script. If the product is unknown the map reflects the current state of the system, i.e. the bay number is probably not correct because of unpopulated bays and detection order of the OS. However it might be ok as a starting point.]
[n:dry?Just print the related map or rule to stdout, but do not install it.]
[R:ramfs?Generate the initramfs hook, so that all related files get copied to the initrd image and thus the script can fire on boot.]
[r:rule?Generate the rule file, which makes use of this script.]
[u:udevs?Triggers an udev add event for all /dev/sd[a-z]][a-z]]*.]
\n\n[\adevpath\a]
'
while getopts -a "${PROG}" "${ print ${USAGE}; }" OPT ; do
	case ${OPT} in
		h) showUsage M ;;
		F) typeset +f && exit 0 ;;
		T) if [[ ${OPTARG} == 'ALL' ]]; then
                typeset -ft ${ typeset +f ; }
            else 
                typeset -ft ${OPTARG//,/ } 
            fi
            ;;
		b) BIOS="${OPTARG}" ;;
		l) LIST=1 ;;
		n) DRY=1 ;;
		m) MKMAP=1 ;;
		r) MKRULE=1 ;;
		R) INSTHOOK=1 ;;
		D)
			if [[ ! -d ${OPTARG} ]]; then
				print -u2 "Directory '${OPTARG}' does not exist - exiting!"
				exit 1
			fi
			ROOTFS="${OPTARG%%/}" ;;
		u) UDEV=1 ;;
		*) showUsage ;;
	esac
done
(( IDX=OPTIND-1 ))
shift $IDX

(( LIST )) && listMap
(( MKRULE )) && makeRule
(( MKMAP )) && makeMap
(( INSTHOOK )) && installHook
(( UDEV )) && fireUdev 
if [[ -n $1 ]]; then
	lookup "$1"
elif (( MKRULE + MKMAP + INSTHOOK + UDEV + LIST == 0 )); then
	showUsage
fi