#!/bin/bash
#
# Copyright (c) 2020-2025, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
#

# This script helps enable/disable PCIe Relaxed Ordering on
# the NVMe disks on NVIDIA platforms

function usage {
    echo ""
    echo "Usage: $(basename $0) enable|disable"
    echo "  Enable or Disable PCIe Relaxed Ordering on NVMe data disks"
    echo ""
    exit 1
}

function is_less_than {
    # use sort -V to check if $1 is less than $2
    if [[ "${1}" = "`echo -e "${1}\n${2}" | sort -V | head -n1`" ]]; then
        echo "true"
    else
        echo "false"
    fi
}

# Argument sanity check
ENABLE_RO="false"
if [[ $# -ne 1 ]]; then
    usage
else
    if [ "$1" != "enable" ] && [ "$1" != "disable" ]; then
        usage
    elif [ "$1" = "enable" ]; then
        ENABLE_RO="true"
    fi
fi

# Use common APIs to figure out platform
NV_SCRIPT_DIR="/usr/local/sbin/nv_scripts/"
general_funcs="${NV_SCRIPT_DIR}/general_funcs.bash"
device_funcs="${NV_SCRIPT_DIR}/device_funcs.bash"
plat_funcs="${NV_SCRIPT_DIR}/plat_funcs.bash"
. ${general_funcs}
. ${device_funcs}
. ${plat_funcs}

must_run_as_root

# Only allow enabling Relaxed Ordering on some platforms
if ! plat_needs_nvme_relaxed_ordering_config; then
    echo "This command is not supported on this platform"
    exit 2
fi

# Check that all data disks are the same vendor

# Form a list of disks "/dev/nvme0 /dev/nvme1 /dev/nvme2 ..."
DATA_DISKS=`${NV_SCRIPT_DIR}/get_data_drives.bash | awk -F':' '{print $NF}' | sed 's/n1//g' | sed 's/,/ /g'`

# get_data_drives.bash can return unknown here if there are no
# data disks
if [ "${DATA_DISKS}" = "unknown" ]; then
    echo "No data disks found"
    exit 2
fi

PREVIOUS_VID="0xFFFF"
NVME_DATA_DISKS=""
for x in ${DATA_DISKS}; do
    # If the data disk is not an NVMe disk (for instance /dev/sda), disregard it by
    # ignoring it if /sys/class/nvme/<DATA_DISK_NAME>/device/vendor file does not exist
    # Note: The ${x#\/dev\/} below removes "/dev/" from the device name: 
    #       For instance: "/dev/nvme1" will become "nvme1"
    if [ -f /sys/class/nvme/${x#\/dev\/}/device/vendor ]; then
        VID=`cat /sys/class/nvme/${x#\/dev\/}/device/vendor`
        if [ ${VID} != ${PREVIOUS_VID} ]; then
            if [ ${PREVIOUS_VID} = "0xFFFF" ]; then
                PREVIOUS_VID=${VID}
                NVME_DATA_DISKS="$NVME_DATA_DISKS $x"
            else
                echo "Relaxed Ordering with mixed vendor NVMe disks is not supported"
                exit 2
            fi
        else
            NVME_DATA_DISKS="$NVME_DATA_DISKS $x"
        fi
    fi
done

# Check if nvme-cli is installed
CHECK_NVME_INSTALLED=$(nvme version 2>/dev/null)
if [ $? -ne 0 ]; then
    echo "nvme-cli utility is not installed. Please install nvme-cli and try again."
    exit 1
fi

# Check version of nvme-cli for cli option change from -v to -V
NVME_CLI_VERSION=$(echo ${CHECK_NVME_INSTALLED} | grep -E '^nvme version' | awk '{print $3}')
NVME_NEW="2.7"
USE_OLD_CLI=$(is_less_than "${NVME_CLI_VERSION}" "${NVME_NEW}")
if [[ "${USE_OLD_CLI}" = "true" ]]; then
    echo "Using nvme-cli version less than ${NVME_NEW} with -v"
    VALUE_CLI="-v"
else
    echo "Using nvme-cli version greater than ${NVME_NEW} with -V"
    VALUE_CLI="-V"
fi

# Enable Relaxed Ordering
VID_SAMSUNG="0x144d"
case "${VID}" in
    "${VID_SAMSUNG}")
        # For x in the list of disks ("/dev/nvme0 /dev/nvme1 /dev/nvme2 ...")
        for x in ${NVME_DATA_DISKS}; do
            # Check if the drive is PM1743 or PM1733
            if [ -f "/sys/class/nvme/${x#/dev/}/model" ]; then
                MODEL=$(cat "/sys/class/nvme/${x#/dev/}/model")
                if [[ "$MODEL" == *"MZWLJ3T8HBLS-00007"* ]]; then
                    # PM1733 - use feature 198
                    FEATURE_ID="198"
                    logger "PM1733 detected on ${x}, using feature ID ${FEATURE_ID}"
                    if [ "${ENABLE_RO}" = "true" ]; then
                        logger "Enabling RO on ${x} with feature ${FEATURE_ID}"
                        nvme set-feature ${x} -f ${FEATURE_ID} "${VALUE_CLI}" 1 -s
                    else
                        logger "Disabling RO on ${x} with feature ${FEATURE_ID}"
                        nvme set-feature ${x} -f ${FEATURE_ID} "${VALUE_CLI}" 0 -s
                    fi
                else
                    echo "Skipping Samsung model:" ${MODEL} "on" ${x}
                    echo "Relaxed Ordering is already enabled on this drive"
                    continue
                fi
                continue
            fi
        done
        ;;
    *)
        echo "Relaxed Ordering is only supported on Samsung NVMe disks"
        exit 3
esac

exit 0