ebsnvme-id

#!/usr/bin/bash

# Copyright Amazon.com, Inc. and its affiliates. All Rights Reserved.
#
# Licensed under the MIT License. See the LICENSE accompanying this file
# for the specific language governing permissions and limitations under
# the License.

# Bash/Unix implementation of ebsnvme-id that requires nvme-cli >= 1.13
# Since this has to run inside initramfs, this is not using getopt/getopts.

# Change this in case your distro puts nvme-cli in a different location
NVME=/usr/sbin/nvme

# The EBS model name
EBS_MN="Amazon Elastic Block Store"

##
# Ensures that the kernel device name in $1 is actually an EBS volume.
ensure_ebs_volume() {
  local device=$1
  local mn

  if [[ ! -b /dev/$device ]]; then
    echo "[ERROR] /dev/$device is not a block device" >&2
    exit 1
  fi

  mn="$(read_nvme_attr "$device" mn)"

  if [[ "$mn" != "${EBS_MN}" ]]; then
    echo "[ERROR] /dev/$device is not an EBS device, but is ${mn:-<unknown>}" >&2
    exit 1
  fi
}

##
# Prints the block device mapping name for the device specified in $1
print_block_device_mapping() {
  local device=$1
  local block_device_name

  block_device_name=$(read_block_device_name "$device")

  echo "${block_device_name}"
}

##
# Prints udev env vars for the device specified in $1
print_udev_env() {
  local device=$1
  local bdev
  local volid

  bdev=$(read_block_device_name "$device")
  volid=$(read_volume_id "$device")

  if [[ -z $bdev ]] || [[ -z $volid ]]; then
    echo "[ERROR] /dev/$device cannot determine udev environment" >&2
    exit 1
  fi

  echo "_EC2_BLOCK_DEVICE_MAPPING=${bdev}"
  echo "_EBS_VOLUME_ID=${volid}"
}

##
# Prints the volume ID for the device specified in $1
print_volume_id() {
  local device=$1
  local vol_id

  vol_id=$(read_volume_id "$device")

  echo "Volume ID: ${vol_id}"
}

##
# Reads the block device name from the device specified in $1 and strips off
# any leading /dev/ that might be present
read_block_device_name() {
  local device=$1
  local bdev

  ensure_ebs_volume "$device"

  bdev=$(read_nvme_attr "$device" bdev)

  if [[ -z $bdev ]]; then
    echo "[ERROR] /dev/$device has no block device mapping available" >&2
    exit 1
  fi

  printf "%s" "${bdev#/dev/}"
}

##
# Reads the NVMe attribute specified by $2 from the device specified in $1
read_nvme_attr() {
  local device=$1
  local nvme_attr=$2

  ${NVME} amzn id-ctrl "/dev/${device}" | while read -r line; do
    if [[ $line =~ ^${nvme_attr}[[:space:]]*: ]]; then
      printf "%s" "${line#*: }"
      return 0
    fi
  done

  return 1
}

##
# Reads the volume ID for the device specified in $1 and parses it into a
# regular EBS vol-abcdef12345 format.
read_volume_id() {
  local device=$1
  local sn

  ensure_ebs_volume "$device"

  sn=$(read_nvme_attr "$device" sn)

  if [[ -z $sn ]]; then
    echo "[ERROR] /dev/$device has no volume ID available" >&2
    exit 1
  fi

  # Strip the prefix, which is usually vol, but we should tolerate vol- too
  local without_prefix=${sn#vol}

  # Print out vol-XXX after stripping the - prefix which is probably not present
  printf "vol-%s" "${without_prefix#-}"
}

##
# Prints program usage information and exits
usage() {
  echo "usage: ebsnvme-id [OPTION]

Reads EBS information from NVMe devices.

Options:
  -h            show this help message and exit
  -v DEVICE     Return volume ID
  -b DEVICE     Return block device mapping
  -u DEVICE     Deprecated alias for -b DEVICE
  -U DEVICE     Output data in a format suitable to be used as an IMPORT{program} in a udev rule
"
}

# Check if we have enough arguments
if [[ $# -ne 2 ]]; then
  usage
  exit 1
fi

# Name the arguments and strip any leading /dev
option="$1"
device="${2#/dev/}"

# Parse options
case $option in
  -U)
    print_udev_env "$device"
    ;;
  -b | -u)
    print_block_device_mapping "$device"
    ;;
  -v)
    print_volume_id "$device"
    ;;
  *)
    usage
    exit 1
    ;;
esac