#!/bin/bash
# Copyright 2020 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Used to generate symlinks for NVMe devices (both local SSD and
# persistent disk) using the disk names reported by the metadata server.

# Locations of the script's dependencies
nvme_cli_bin=/usr/sbin/nvme
sysfs_nvme_class_dir=/sys/class/nvme
sysfs_pci_devices_dir=/sys/bus/pci/devices

# Bash regex to parse device paths and controller identification
readonly PD_CONTROLLER_REGEX="nvme_card-pd"
readonly SSD_CONTROLLER_REGEX="nvme_card[0-9]*"
readonly CONTROLLER_NUMBER_REGEX="nvme_card([[:digit:]]+)"
readonly NAMESPACE_NUMBER_REGEX="/dev/nvme[[:digit:]]+n([[:digit:]]+).*"
readonly PARTITION_NUMBER_REGEX="/dev/nvme[[:digit:]]+n[[:digit:]]+p([[:digit:]]+)"
readonly PD_NVME_REGEX="sn[[:space:]]+:[[:space]]+nvme_card-pd"

# Globals used to generate the symlinks for a NVMe disk.  These are populated
# by the identify_pd_disk function and exported for consumption by udev rules.
ID_SERIAL=''
ID_SERIAL_SHORT=''

#######################################
# Helper function to log an error message to stderr.
# Globals:
#   None
# Arguments:
#   String to print as the log message
# Outputs:
#   Writes error to STDERR
#######################################
function err() {
  echo "[$(date +'%Y-%m-%dT%H:%M:%S%z')]: $*" >&2
}

#######################################
# Retrieves the device name for an NVMe namespace using nvme-cli.
# Globals:
#   Uses nvme_cli_bin
# Arguments:
#   The path to the nvme namespace (/dev/nvme0n?)
# Outputs:
#   The device name parsed from the JSON in the vendor ext of the ns-id command.
# Returns:
#   0 if the device name for the namespace could be retrieved, 1 otherwise
#######################################
function get_namespace_device_name() {
  local nvme_json
  nvme_json="$("$nvme_cli_bin" id-ns -b "$1" | dd bs=1 skip=384 2>/dev/null)"
  if [[ $? -ne 0 ]]; then
    return 1
  fi

  if [[ -z "$nvme_json" ]]; then
    err "NVMe Vendor Extension disk information not present"
    return 1
  fi

  local device_name
  device_name="$(echo "$nvme_json" | grep device_name | sed -e 's/.*"device_name":[  \t]*"\([.a-zA-Z0-9_-]\+\)".*/\1/')"

  # Error if our device name is empty
  if [[ -z "$device_name" ]]; then
    err "Empty name"
    return 1
  fi

  echo "$device_name"
  return 0
}

#######################################
# Retrieves the nsid for an NVMe namespace
# Globals:
#   None
# Arguments:
#   The path to the nvme namespace (/dev/nvme0n*)
# Outputs:
#   The namespace number/id
# Returns:
#   0 if the namespace id could be retrieved, 1 otherwise
#######################################
function get_namespace_number() {
  local dev_path="$1"
  local namespace_number
  if [[ "$dev_path" =~ $NAMESPACE_NUMBER_REGEX ]]; then
    namespace_number="${BASH_REMATCH[1]}"
  else
    return 1
  fi

  echo "$namespace_number"
  return 0
}

#######################################
# Retrieves the partition number for a device path if it exists
# Globals:
#   None
# Arguments:
#   The path to the device partition (/dev/nvme0n*p*)
# Outputs:
#   The value after 'p' in the device path, or an empty string if the path has
#   no partition.
#######################################
function get_partition_number() {
  local dev_path="$1"
  local partition_number
  if [[ "$dev_path" =~ $PARTITION_NUMBER_REGEX ]]; then
    partition_number="${BASH_REMATCH[1]}"
    echo "$partition_number"
  else
    echo ''
  fi
  return 0
}

#######################################
# Retrieves the controller number from the device model if it exists
# Globals:
#   None
# Arguments:
#   The NVMe device model (nvme_card or nvme_card1/2/3/...)
# Outputs:
#   The controller id/number
#######################################
function get_controller_number() {
  local dev_model="$1"
  local controller_number
  if [[ "$dev_model" =~ $CONTROLLER_NUMBER_REGEX ]]; then
    controller_number="${BASH_REMATCH[1]}"
    echo "$controller_number"
  else
    # if it's 'nvme_card', echo 0. This is for backward compatibility.
    echo '0'
  fi
  return 0
}

#######################################
# Generates a symlink for a PD-NVMe device using the metadata's disk name.
# Primarily used for testing but can be used if the script is directly invoked.
# Globals:
#   Uses ID_SERIAL_SHORT (can be populated by identify_pd_disk)
# Arguments:
#   The device path for the disk
#######################################
function gen_symlink() {
  local dev_path="$1"
  local partition_number="$(get_partition_number "$dev_path")"

  if [[ -n "$partition_number" ]]; then
    ln -s "$dev_path" /dev/disk/by-id/google-"$ID_SERIAL_SHORT"-part"$partition_number" > /dev/null 2>&1
  else
    ln -s "$dev_path" /dev/disk/by-id/google-"$ID_SERIAL_SHORT" > /dev/null 2>&1
  fi

  return 0
}

#######################################
# Populates the ID_* global variables with a disk's device name and namespace
# Globals:
#   Populates ID_SERIAL_SHORT, and ID_SERIAL
# Arguments:
#   The device path for the disk
# Returns:
#   0 on success and 1 if an error occurs
#######################################
function identify_pd_disk() {
  local dev_path="$1"
  local dev_name
  dev_name="$(get_namespace_device_name "$dev_path")"
  if [[ $? -ne 0 ]]; then
    return 1
  fi

  ID_SERIAL_SHORT="$dev_name"
  ID_SERIAL="Google_PersistentDisk_${ID_SERIAL_SHORT}"
  return 0
}

#######################################
# Populates the ID_* global variables with a disk's device name and namespace
# Globals:
#   Populates ID_SERIAL_SHORT, and ID_SERIAL
# Arguments:
#   The device path for the disk
# Returns:
#   0 on success and 1 if an error occurs
#######################################
function identify_local_ssd_disk() {
  local dev_model="$1"
  local dev_path="$2"
  local controller_number
  controller_number="$(get_controller_number "$dev_model")"
  if [[ $? -ne 0 ]]; then
    return 1
  fi

  local namespace_number
  namespace_number="$(get_namespace_number "$dev_path")"
  if [[ $? -ne 0 ]]; then
    return 1
  fi

  ID_SERIAL_SHORT="local-nvme-ssd-$(($controller_number+$namespace_number-1))"
  ID_SERIAL="Google_EphemeralDisk_${ID_SERIAL_SHORT}"
  return 0
}

#######################################
# Populates the ID_* global variables for a local SSD from the NVMe Identify
# Namespace vendor extension.
#
# identify_local_ssd_disk derives the index arithmetically from the controller
# and namespace numbers. That works on a VM, where a single "nvme_card"
# controller exposes each local SSD as its own namespace, but not on Bare Metal
# (e.g. Z3 Metal), where every physical drive is an independent "nvme_card"
# controller with nsid=1 and the arithmetic returns 0 for all of them.
# Firmware instead stamps a stable per-drive name into the Identify Namespace
# vendor extension, which is what this function reads.
#
# This function makes no attempt to decide whether the vendor extension is the
# right source for a given device; that judgement belongs to the caller. In the
# udev path it is 65-gce-disk-naming.rules that decides, via its match clauses.
#
# Globals:
#   Populates ID_SERIAL_SHORT, and ID_SERIAL
# Arguments:
#   The device path for the disk
# Returns:
#   0 if a name was resolved, 1 otherwise. Returning 1 is the ordinary outcome
#   for any device that does not publish a device name; it is not an error.
#######################################
function identify_local_ssd_from_vendor_ext() {
  local dev_path="$1"

  local dev_name
  if ! dev_name="$(get_namespace_device_name "$dev_path")"; then
    return 1
  fi

  ID_SERIAL_SHORT="$dev_name"
  ID_SERIAL="Google_EphemeralDisk_${ID_SERIAL_SHORT}"
  return 0
}

#######################################
# Populates the ID_* global variables by naming a bare-metal Local SSD after
# the lexicographical position of its serial number among all NVMe drives that
# share its PCI vendor ID.
#
# Only relevant for bare-metal instances whose Local SSDs have neither an
# "nvme_card*" model nor a device name in the Identify Namespace vendor
# extension. A drive's vendor ID and serial number never change, so sorting
# the serials yields a stable 0-based index for each drive.
#
# No name is assigned unless every NVMe PCI function with the drive's vendor ID
# has registered a controller with a serial number: an incomplete list would
# shift the indices and attach symlinks to the wrong drives.
#
# Globals:
#   Uses sysfs_nvme_class_dir, sysfs_pci_devices_dir
#   Populates ID_SERIAL_SHORT, and ID_SERIAL
# Arguments:
#   The nvme id-ctrl output for the disk
#   The device path for the disk
# Returns:
#   0 on success and 1 if an error occurs
#######################################
function identify_local_ssd_by_serial() {
  local dev_model="$1"
  local dev_path="$2"

  local target_sn
  target_sn="$(echo "$dev_model" | awk '/^sn[[:space:]]*:/ {print $3}' | tr -d '[:space:]')"
  if [[ -z "$target_sn" ]]; then
    err "Failed to extract serial number from id-ctrl for $dev_path"
    return 1
  fi

  local target_vid
  target_vid="$(echo "$dev_model" | awk '/^vid[[:space:]]*:/ {print $3}' | tr -d '[:space:]')"
  if [[ -z "$target_vid" ]]; then
    err "Failed to extract vendor ID from id-ctrl for $dev_path"
    return 1
  fi

  # Count the NVMe PCI functions (class 0x010802) with the same vendor ID. This
  # is the number of Local SSDs the scan below must find; any that are still
  # probing at boot have no NVMe controller yet.
  local expected_count=0
  local pci_dir
  for pci_dir in "$sysfs_pci_devices_dir"/*; do
    local pci_vendor pci_class
    pci_vendor="$(cat "$pci_dir/vendor" 2>/dev/null | tr -d '[:space:]')"
    pci_class="$(cat "$pci_dir/class" 2>/dev/null | tr -d '[:space:]')"
    if [[ "${pci_vendor,,}" == "${target_vid,,}" && "${pci_class,,}" == "0x010802" ]]; then
      expected_count=$((expected_count + 1))
    fi
  done
  if [[ "$expected_count" -eq 0 ]]; then
    err "No NVMe PCI functions with vendor ID $target_vid found for $dev_path"
    return 1
  fi

  # Wait up to 5 seconds for every Local SSD controller to register at boot.
  local -a all_serials=()
  local attempt
  for (( attempt=1; attempt<=50; attempt++ )); do
    all_serials=()
    for ctrl_dir in "$sysfs_nvme_class_dir"/nvme*; do
      [[ -d "$ctrl_dir" ]] || continue
      local ctrl_model
      ctrl_model="$(cat "$ctrl_dir/model" 2>/dev/null | tr -d '[:space:]')"
      # Skip Google virtual NVMe controllers (e.g. the "nvme_card-pd"
      # Persistent Disk) so they can never take a Local SSD index.
      if [[ "$ctrl_model" == nvme_card* ]]; then
        continue
      fi
      local ctrl_vendor
      ctrl_vendor="$(cat "$ctrl_dir/device/vendor" 2>/dev/null | tr -d '[:space:]')"
      if [[ "${ctrl_vendor,,}" != "${target_vid,,}" ]]; then
        continue
      fi
      local sn
      sn="$(cat "$ctrl_dir/serial" 2>/dev/null | tr -d '[:space:]')"
      if [[ -n "$sn" ]]; then
        all_serials+=("$sn")
      fi
    done
    if [[ "${#all_serials[@]}" -ge "$expected_count" ]]; then
      break
    fi
    sleep 0.1
  done

  if [[ "${#all_serials[@]}" -lt "$expected_count" ]]; then
    err "Found ${#all_serials[@]} of $expected_count Local SSD serial numbers \
for vendor ID $target_vid; not naming $dev_path"
    return 1
  fi

  # Sort serial numbers to assign deterministic 0-based index
  local -a sorted_serials=()
  mapfile -t sorted_serials < <(printf "%s\n" "${all_serials[@]}" | LC_ALL=C sort)

  local ssd_index=-1
  local i
  for (( i=0; i<${#sorted_serials[@]}; i++ )); do
    if [[ "${sorted_serials[$i]}" == "$target_sn" ]]; then
      ssd_index="$i"
      break
    fi
  done

  if [[ "$ssd_index" -lt 0 ]]; then
    err "Target serial '$target_sn' not found among detected local SSDs (${sorted_serials[*]})"
    return 1
  fi

  ID_SERIAL_SHORT="local-nvme-ssd-${ssd_index}"
  ID_SERIAL="Google_EphemeralDisk_${ID_SERIAL_SHORT}"
  return 0
}

function print_help_message() {
  echo "Usage: google_nvme_id [-s] [-e] [-o] [-h] -d device_path"
  echo "  -d <device_path> (Required): Specifies the path to generate a name"
  echo "        for.  This needs to be a path to an nvme device or namespace"
  echo "  -s: Create symbolic link for the disk under /dev/disk/by-id."
  echo "        Otherwise, the disk name will be printed to STDOUT"
  echo "  -e: Take the local SSD name from the NVMe Identify Namespace vendor"
  echo "        extension instead of deriving it from the controller and"
  echo "        namespace numbers.  Prints nothing and exits 0 if the device"
  echo "        publishes no name, so that callers keep any name they have"
  echo "        already derived."
  echo "  -o: Generate symlinks based on lexicographical ordering of serial"
  echo "        number.  Only relevant for bare-metal instances."
  echo "  -h: Print this help message"
}

function main() {
  local opt_gen_symlink='false'
  local opt_vendor_ext='false'
  local opt_order_by_serial='false'
  local device_path=''
  local OPTIND=1

  while getopts :d:esoh flag; do
    case "$flag" in
      d) device_path="$OPTARG";;
      e) opt_vendor_ext='true';;
      s) opt_gen_symlink='true';;
      o) opt_order_by_serial='true';;
      h) print_help_message
         return 0
         ;;
      :) echo "Invalid option: ${OPTARG} requires an argument" 1>&2
         return 1
         ;;
      *) return 1
    esac
  done

  if [[ -z "$device_path" ]]; then
    echo "Device path (-d) argument required. Use -h for full usage." 1>&2
    exit 1
  fi

  # Ensure the nvme-cli command is installed
  command -v "$nvme_cli_bin" > /dev/null 2>&1
  if [[ $? -ne 0 ]]; then
    err "The nvme utility (/usr/sbin/nvme) was not found. You may need to run \
with sudo or install nvme-cli."
    return 1
  fi

  # Ensure the passed device is actually an NVMe device
  "$nvme_cli_bin" id-ctrl "$device_path" &>/dev/null
  if [[ $? -ne 0 ]]; then
    err "Passed device was not an NVMe device.  (You may need to run this \
script as root/with sudo)."
    return 1
  fi

  # Detect the type of attached nvme device
  local controller_id
  controller_id=$("$nvme_cli_bin" id-ctrl "$device_path")

  if [[ "$opt_vendor_ext" == 'true' ]]; then
    # Name the disk from the Identify Namespace vendor extension, as the caller
    # asked. This branch makes no judgement of its own about what kind of
    # device or machine this is; deciding when the vendor extension is the
    # right source is the caller's job. 65-gce-disk-naming.rules invokes it
    # only for the single-controller local SSD case.
    #
    # Emitting nothing is the ordinary outcome for a device that publishes no
    # name, and is not an error: udev imports only the keys a program actually
    # prints, so the caller keeps the name it had.
    if ! identify_local_ssd_from_vendor_ext "$device_path"; then
      return 0
    fi
  elif [[ "$opt_order_by_serial" == 'true' ]]; then
    # Name a bare-metal local SSD by the lexicographical position of its serial
    # number among all NVMe drives with the same PCI vendor ID.
    identify_local_ssd_by_serial "$controller_id" "$device_path"
    if [[ $? -ne 0 ]]; then
      return $?
    fi
  elif [[ "$controller_id" =~ $PD_CONTROLLER_REGEX ]] ; then
    # Fill the global variables for the id command for the given disk type
    # Error messages will be printed closer to error, no need to reprint here
    identify_pd_disk "$device_path"
    if [[ $? -ne 0 ]]; then
      return $?
    fi
  elif [[ "$controller_id" =~ $SSD_CONTROLLER_REGEX ]] ; then
    identify_local_ssd_disk "$controller_id" "$device_path"
    if [[ $? -ne 0 ]]; then
      return $?
    fi
  else
    err "Device is not a NVMe device"
    return 1
  fi

  # Gen symlinks or print out the globals set by the identify command
  if [[ "$opt_gen_symlink" == 'true' ]]; then
    gen_symlink "$device_path"
  else
    # These will be consumed by udev
    echo "ID_SERIAL_SHORT=${ID_SERIAL_SHORT}"
    echo "ID_SERIAL=${ID_SERIAL}"
  fi

  return $?

}

# Give the ability to test the script without running the main logic
if [[ ! $TEST == 'test' ]]; then
  main "$@"
fi
