#!/bin/sh
#################################################################
#
#   SPDX-License-Identifier: BSD-3-Clause
#   Copyright(c) 2007-2026 Intel Corporation
# 
#   These contents may have been developed with support from one or more
#   Intel-operated generative artificial intelligence solutions.
#
#################################################################
#
#
# qat_init.sh Setup drivers for Intel QAT.
#

VFIO_DRIVER=${VFIO_DRIVER-vfio-pci}
QAT_USER_GROUP=${QAT_USER_GROUP-qat}

INTEL_VENDORID="0x8086"
DH895_DEVICE_PCI_ID="0x0435"
DH895_DEVICE_PCI_ID_VM="0x0443"
DH895_DEVICE_NAME="dh895xcc"
DH895_DRIVER_NAME="qat_dh895xcc"
C62X_DEVICE_PCI_ID="0x37c8"
C62X_DEVICE_PCI_ID_VM="0x37c9"
C62X_DEVICE_NAME="c6xx"
C62X_DRIVER_NAME="qat_c62x"
C3XX_DEVICE_PCI_ID="0x19e2"
C3XX_DEVICE_PCI_ID_VM="0x19e3"
C3XX_DEVICE_NAME="c3xxx"
C3XX_DRIVER_NAME="qat_c3xxx"
D15XX_DEVICE_PCI_ID="0x6f54"
D15XX_DEVICE_PCI_ID_VM="0x6f55"
D15XX_DEVICE_NAME="d15xx"
D15XX_DRIVER_NAME="qat_d15xx"
QAT_4XXX_DEVICE_PCI_ID="0x4940"
QAT_4XXX_DEVICE_PCI_ID_VM="0x4941"
QAT_401XX_DEVICE_PCI_ID="0x4942"
QAT_401XX_DEVICE_PCI_ID_VM="0x4943"
QAT_402XX_DEVICE_PCI_ID="0x4944"
QAT_402XX_DEVICE_PCI_ID_VM="0x4945"
QAT_4XXX_DEVICE_NAME="4xxx"
QAT_4XXX_DRIVER_NAME="qat_4xxx"
PF_NAMES="$QAT_4XXX_DEVICE_NAME"
VF_DEVICE_IDS="$QAT_4XXX_DEVICE_PCI_ID_VM $QAT_401XX_DEVICE_PCI_ID_VM $QAT_402XX_DEVICE_PCI_ID_VM"
QAT_420XX_DEVICE_PCI_ID="0x4946"
QAT_420XX_DEVICE_PCI_ID_VM="0x4947"
QAT_420XX_DEVICE_NAME="420xx"
QAT_420XX_DRIVER_NAME="qat_420xx"
PF_NAMES="$PF_NAMES $QAT_420XX_DEVICE_NAME"
VF_DEVICE_IDS="$VF_DEVICE_IDS $QAT_420XX_DEVICE_PCI_ID_VM"
SERVICES_LIST="sym asym sym;asym dc sym;dc asym;dc dcc"
SERVICES_ENABLED="NONE"
SERVICES_ENABLED_FOUND="FALSE"
SYSTEMD_FILE="/usr/lib/systemd/system/qat.service"
QAT_HUGEPAGES_DIR="/dev/hugepages/qat"

# space separated array for supporting new devices pci ids
SUPPORTED_DEVICE_PCI_IDS="$QAT_4XXX_DEVICE_PCI_ID"
SUPPORTED_DEVICE_PCI_IDS="$SUPPORTED_DEVICE_PCI_IDS $QAT_401XX_DEVICE_PCI_ID"
SUPPORTED_DEVICE_PCI_IDS="$SUPPORTED_DEVICE_PCI_IDS $QAT_402XX_DEVICE_PCI_ID"
SUPPORTED_DEVICE_PCI_IDS="$SUPPORTED_DEVICE_PCI_IDS $QAT_420XX_DEVICE_PCI_ID"

# space separated array for supporting new devices pci names
SUPPORTED_DRIVER_NAMES="$QAT_4XXX_DRIVER_NAME"
SUPPORTED_DRIVER_NAMES="$SUPPORTED_DRIVER_NAMES $QAT_420XX_DRIVER_NAME"

check_config() {
    CONFIG_FILE=$(awk -F'EnvironmentFile=-' '{print $2}' $SYSTEMD_FILE | grep '\S')
    if [ ! -f "$CONFIG_FILE" ]; then
        return
    fi

    SERVICES_ENABLED=$(awk -F'ServicesEnabled=' '{print $2}' $CONFIG_FILE | grep '\S')
    SERVICES_ENABLED_FOUND="FALSE"
    for SERVICE in $SERVICES_LIST
    do
        if [ "$SERVICE" = "$SERVICES_ENABLED" ]; then
            SERVICES_ENABLED_FOUND="TRUE"
            break
        fi
    done
}

get_module_state() {
    CMD=""
    for SUPPORTED_DRIVER_NAME in $SUPPORTED_DRIVER_NAMES;
    do
        CMD="$CMD -e ^$SUPPORTED_DRIVER_NAME"
    done
    echo "$(cat /proc/modules | grep $CMD | cut -d' ' -f5)"
}

check_driver() {
    # check if is running on VM (no PFs)
    PF_AVAILABLE=0
    for SUPPORTED_PCI_ID in $SUPPORTED_DEVICE_PCI_IDS;
    do
        if [ ! -z "$(lspci -nD | grep ${SUPPORTED_PCI_ID#0x})" ]; then
            PF_AVAILABLE=1
            break
        fi
    done

    if [ $PF_AVAILABLE -ne 1 ]; then
        echo "lspci: No PFs found, so assume qatlib is running on a VM"
        return
    fi
    # qat driver needs 0.2s per PF, so for 8S machine it needs ~ 6.4 s
    # TIMEOUT = TIMEOUT_CNT * ATTEMPT_INTERVAL so there is a 20s
    TIMEOUT_CNT=40
    ATTEMPT_INTERVAL=0.5
    ATTEMPT_CNT=0
    CURRENT_STATE=$(get_module_state)

    while [ "$CURRENT_STATE" != "Live" ]
    do
        if [ $ATTEMPT_CNT -ge $TIMEOUT_CNT ]; then
            TIMEOUT=$(awk "BEGIN {print $ATTEMPT_CNT * $ATTEMPT_INTERVAL}")
            echo "QAT driver is still not present after ${TIMEOUT}s. Aborting qat_init"
            exit 1
        fi
        ATTEMPT_CNT=`expr $ATTEMPT_CNT + 1`
        sleep $ATTEMPT_INTERVAL
        CURRENT_STATE=$(get_module_state)

    done
}

sysfs_config() {
    if [ "$SERVICES_ENABLED_FOUND" != "TRUE" ]; then
        return
    fi

    PCI_DEV_SUPPORTED=0
    PCI_DEVS=`ls -d /sys/bus/pci/devices/* | awk 'BEGIN{FS="/"} {print $NF}'`
    for DEV_PCI_ADDR in $PCI_DEVS
    do
        DEV=/sys/bus/pci/devices/$DEV_PCI_ADDR
        PCI_DEV=`cat $DEV/device 2> /dev/null`

        for SUPPORTED_PCI_ID in $SUPPORTED_DEVICE_PCI_IDS;
        do
            if [ "$PCI_DEV" = "$SUPPORTED_PCI_ID" ]; then
                PCI_DEV_SUPPORTED=1
                break
            fi
        done

        if [ $PCI_DEV_SUPPORTED -ne 1 ]; then
            continue
        fi
        PCI_DEV_SUPPORTED=0

        if [ ! -e "$DEV/qat/cfg_services" ]; then
            echo "Cannot access $DEV/qat/cfg_services: No such file or directory."
            echo "You may need to update your kernel driver. See QAT kernel release notes."
            continue
        fi
        CURRENT_SERVICES=`cat $DEV/qat/cfg_services`
        if [ "$CURRENT_SERVICES" != "$SERVICES_ENABLED" ]; then
            CURRENT_STATE=`cat $DEV/qat/state`
            if [ "$CURRENT_STATE" = "up" ]; then
                echo down > $DEV/qat/state
            fi
            echo $SERVICES_ENABLED > $DEV/qat/cfg_services
            if [ $? != 0 ]; then
                echo "Can't write to cfg_services file: $SERVICES_ENABLED is not supported."
                echo "You may need to update your kernel driver. See QAT kernel release notes."
                continue
            fi
            CURRENT_SERVICES=`cat $DEV/qat/cfg_services`
        fi
        echo "$0, device $DEV_PCI_ADDR configured with services: $CURRENT_SERVICES"
    done
}

unbind() {
    BSF=$1
    OLD_DRIVER=$2

    LOOP=0
    while [ $LOOP -lt  5 ]
    do
        echo -n 2> /dev/null $BSF > /sys/bus/pci/drivers/$OLD_DRIVER/unbind
        if [ $? -eq 0 ]; then
            break
        else
           LOOP=`expr $LOOP + 1`
            sleep 1
        fi
    done
}

override() {
    BSF=$1
    NEW_DRIVER=$2

    LOOP=0
    while [ $LOOP -lt  5 ]
    do
        echo -n 2> /dev/null  $NEW_DRIVER > /sys/bus/pci/devices/$BSF/driver_override
        if [ $? -eq 0 ]; then
                break
        else
           LOOP=`expr $LOOP + 1`
            sleep 1
        fi
    done
}

bind() {
    BSF=$1
    DRIVER=$2

    LOOP=0
    while [ $LOOP -lt  5 ]
    do
        echo -n 2> /dev/null $BSF > /sys/bus/pci/drivers/$DRIVER/bind
        if [ $? -eq 0 ]; then
            break
        else
           LOOP=`expr $LOOP + 1`
           echo -n 2> /dev/null $BSF > /sys/bus/pci/devices/$BSF/driver/unbind
            sleep 1
        fi
    done
}

bind_driver() {
    BSF=$1
    DEVICE=$2
    NEW_DRIVER=$VFIO_DRIVER
    MODULE=$VFIO_DRIVER

    # Check if this device should be bound to qat VF driver
    LKCF_DEV=
    for DEV in $LKCF_LIST
    do
        if echo $BSF | grep -q -E $DEV; then
            LKCF_DEV=$DEV
            break
        fi
    done
    if [ $LKCF_DEV ]; then
        # Find the qat vf driver
        case "$DEVICE" in
        $DH895_DEVICE_PCI_ID_VM )
            NEW_DRIVER=${DH895_DEVICE_NAME}vf
            MODULE=${DH895_DRIVER_NAME}vf
            ;;
        $C62X_DEVICE_PCI_ID_VM )
            NEW_DRIVER=${C62X_DEVICE_NAME}vf
            MODULE=${C62X_DRIVER_NAME}vf
            ;;
        $C3XX_DEVICE_PCI_ID_VM )
            NEW_DRIVER=${C3XX_DEVICE_NAME}vf
            MODULE=${C3XX_DRIVER_NAME}vf
            ;;
        $D15XX_DEVICE_PCI_ID_VM )
            NEW_DRIVER=${D15XX_DEVICE_NAME}vf
            MODULE=${D15XX_DRIVER_NAME}vf
             ;;
        $QAT_4XXX_DEVICE_PCI_ID_VM )
            NEW_DRIVER=${QAT_4XXX_DEVICE_NAME}vf
            MODULE=${QAT_4XXX_DRIVER_NAME}vf
             ;;
        $QAT_401XX_DEVICE_PCI_ID_VM )
            NEW_DRIVER=${QAT_4XXX_DEVICE_NAME}vf
            MODULE=${QAT_4XXX_DRIVER_NAME}vf
             ;;
        $QAT_402XX_DEVICE_PCI_ID_VM )
            NEW_DRIVER=${QAT_4XXX_DEVICE_NAME}vf
            MODULE=${QAT_4XXX_DRIVER_NAME}vf
             ;;
        $QAT_420XX_DEVICE_PCI_ID_VM )
            NEW_DRIVER=${QAT_420XX_DEVICE_NAME}vf
            MODULE=${QAT_420XX_DRIVER_NAME}vf
             ;;
        * )
             echo Unsupported PCI device $DEVICE
             ;;
        esac
    fi

    VF_DEV=/sys/bus/pci/devices/$BSF

    if [ ! -d /sys/bus/pci/drivers/$NEW_DRIVER ]; then
        modprobe $MODULE
    fi

    # What driver is currently bound to the device?
    if [ -e $VF_DEV/driver ]; then
         VF_DRIVER=`readlink $VF_DEV/driver | awk 'BEGIN {FS="/"} {print $NF}'`
    else
         VF_DRIVER=
    fi
    if [ x$VF_DRIVER != x$NEW_DRIVER ]; then
        if [ $VF_DRIVER ]; then
            # Unbind from existing driver
            unbind $BSF $VF_DRIVER
        fi

        # Bind to $NEW_DRIVER
        override $BSF $NEW_DRIVER
        bind $BSF $NEW_DRIVER

        # Change permissions on the device,
        # a delay is needed to allow the init caused
        # by the bind to complete before the permissions
        # can be changed
        GROUP=`readlink $VF_DEV/iommu_group | awk 'BEGIN {FS="/"} {print $NF}'`
        if [ -e /dev/vfio/$GROUP ]; then
            sleep 0.1
            chown :$QAT_USER_GROUP /dev/vfio/$GROUP
            chmod +060 /dev/vfio/$GROUP
        fi
    fi
}

enable_sriov() {
    PF_LIST=
    for NAME in $PF_NAMES
    do
        for PF in `ls -d /sys/bus/pci/drivers/$NAME/????:??:??.? 2> /dev/null`
        do
            PF_LIST="$PF_LIST $PF"
        done
    done

    if [ "$PF_LIST" ]; then
        for PF_DEV in $PF_LIST
        do
            # Enable sriov on the PF_DEV
            if [ -r $PF_DEV/sriov_totalvfs -a -w $PF_DEV/sriov_numvfs ]; then
                TOTALVFS=`cat $PF_DEV/sriov_totalvfs`
                NUMVFS=`cat $PF_DEV/sriov_numvfs`
                if [ $TOTALVFS -ne $NUMVFS ]; then
                    echo $TOTALVFS > $PF_DEV/sriov_numvfs
                fi
            fi

            for VF_LINK in `ls -d $PF_DEV/virtfn* 2> /dev/null`
            do
               BSF=`readlink $VF_LINK | awk 'BEGIN {FS="/"} {print $NF}'`
               DEVICE=`cat /sys/bus/pci/devices/$BSF/device`
               bind_driver $BSF $DEVICE &
            done
        done
    else
        # No PFs.  Find by pci device id.
        PCI_DEVICES=`ls -d /sys/bus/pci/devices/* | awk 'BEGIN{FS="/"} {print $NF}'`
        for BSF in $PCI_DEVICES
        do
          DEVICE=`cat /sys/bus/pci/devices/$BSF/device`
          if echo $VF_DEVICE_IDS | grep -q $DEVICE; then
              VENDOR=`cat /sys/bus/pci/devices/$BSF/vendor`
              if [ $VENDOR = $INTEL_VENDORID ]; then
                  bind_driver $BSF $DEVICE &
              fi
          fi
        done
    fi
}

setup_hugepages_directory() {
    # Check if hugepages present
    if [ -d "/dev/hugepages" ]; then
        if [ ! -d $QAT_HUGEPAGES_DIR ]; then
            mkdir "$QAT_HUGEPAGES_DIR"
            # Set group ownership and permissions
            chown :$QAT_USER_GROUP "$QAT_HUGEPAGES_DIR"
            chmod +070 "$QAT_HUGEPAGES_DIR"
        fi
    fi
}

remove_hugepages_directory() {
    # Remove the qat hugepages directory if it exists
    if [ -d "$QAT_HUGEPAGES_DIR" ]; then
        rm -rf "$QAT_HUGEPAGES_DIR"
    fi
}

# Function to update VFIO permissions for QAT VF devices
update_vfio_permissions() {
    # Find all QAT VF devices using /sys/bus/pci/devices
    for PCI_DEV in /sys/bus/pci/devices/*; do
        [ -e "$PCI_DEV/vendor" ] || continue
        VENDOR=`cat "$PCI_DEV/vendor" 2>/dev/null`
        DEVICE=`cat "$PCI_DEV/device" 2>/dev/null`

        # Check if this is a QAT VF device
        MATCHED=0
        for VF_ID in $VF_DEVICE_IDS; do
            if [ "$DEVICE" = "$VF_ID" ] && [ "$VENDOR" = "$INTEL_VENDORID" ]; then
                MATCHED=1
                break
            fi
        done

        if [ $MATCHED -eq 1 ]; then
            # Found a QAT VF, get its IOMMU group
            if [ -e "$PCI_DEV/iommu_group" ]; then
                IOMMU_GROUP=`basename $(readlink "$PCI_DEV/iommu_group")`
                VFIO_NODE="/dev/vfio/$IOMMU_GROUP"

                if [ -e "$VFIO_NODE" ]; then
                    echo "Setting permissions for $VFIO_NODE" \
                         "(QAT VF `basename $PCI_DEV`)"
                    chown :$QAT_USER_GROUP "$VFIO_NODE" 2>/dev/null && \
                        chmod +060 "$VFIO_NODE" 2>/dev/null
                fi
            fi
        fi
    done
}

# Handle reload-permissions mode
if [ "$1" = "reload-permissions" ]; then
    update_vfio_permissions
    exit 0
fi

# Handle command line arguments
if [ "$1" = "remove-hugepages-dir" ]; then
    remove_hugepages_directory
    exit 0
fi

check_config
setup_hugepages_directory
check_driver
sysfs_config
enable_sriov
wait

exit 0
