NVME_ADMIN_IDENTIFY = 0x06
NVME_IOCTL_ADMIN_CMD = 0xC0484E41
AMZN_NVME_VID = 0x1D0F
AMZN_NVME_EBS_MN = "Amazon Elastic Block Store"

NVME_IDENTIFY_CONTROLLER_PSD = [0, 0, 0, 0, 0, 0, 0, 0, " "].pack("SSLLCCCCa16")
NVME_IDENTIFY_CONTROLLER_AMZN_VS = [" ", " "].pack("a32a#{1024 - 32}")

NVME_ADMIN_CMD_TMPL = [
  "C",                                          # ("opcode", c_uint8),                        # op code
  "C",                                          # ("flags", c_uint8),                         # fused operation
  "S",                                          # ("cid", c_uint16),                          # command id
  "L",                                          # ("nsid", c_uint32),                         # namespace id
  "Q",                                          # ("reserved0", c_uint64),
  "P",                                          # ("mptr", c_uint64),                         # metadata pointer
  "P",                                          # ("addr", c_uint64),                         # data pointer
  "L",                                          # ("mlen", c_uint32),                         # metadata length
  "L",                                          # ("alen", c_uint32),                         # data length
  "L",                                          # ("cdw10", c_uint32),
  "L",                                          # ("cdw11", c_uint32),
  "L",                                          # ("cdw12", c_uint32),
  "L",                                          # ("cdw13", c_uint32),
  "L",                                          # ("cdw14", c_uint32),
  "L",                                          # ("cdw15", c_uint32),
  "Q",                                          # ("reserved1", c_uint64)]
].join

NVME_ID_CTRL_TMPL = [
  "S",                                          # ("vid", c_uint16),                          # PCI Vendor ID
  "S",                                          # ("ssvid", c_uint16),                        # PCI Subsystem Vendor ID
  "a20",                                        # ("sn", c_char * 20),                        # Serial Number
  "a40",                                        # ("mn", c_char * 40),                        # Module Number
  "a8",                                         # ("fr", c_char * 8),                         # Firmware Revision
  "C",                                          # ("rab", c_uint8),                           # Recommend Arbitration Burst
  "a3",                                         # ("ieee", c_uint8 * 3),                      # IEEE OUI Identifier
  "C",                                          # ("mic", c_uint8),                           # Multi-Interface Capabilities
  "C",                                          # ("mdts", c_uint8),                          # Maximum Data Transfer Size
  "a#{256 - 78}",                               # ("reserved0", c_uint8 * (256 - 78)),
  "S",                                          # ("oacs", c_uint16),                         # Optional Admin Command Support
  "C",                                          # ("acl", c_uint8),                           # Abort Command Limit
  "C",                                          # ("aerl", c_uint8),                          # Asynchronous Event Request Limit
  "C",                                          # ("frmw", c_uint8),                          # Firmware Updates
  "C",                                          # ("lpa", c_uint8),                           # Log Page Attributes
  "C",                                          # ("elpe", c_uint8),                          # Error Log Page Entries
  "C",                                          # ("npss", c_uint8),                          # Number of Power States Support
  "C",                                          # ("avscc", c_uint8),                         # Admin Vendor Specific Command Configuration
  "a#{512 - 265}",                              # ("reserved1", c_uint8 * (512 - 265)),
  "C",                                          # ("sqes", c_uint8),                          # Submission Queue Entry Size
  "C",                                          # ("cqes", c_uint8),                          # Completion Queue Entry Size
  "S",                                          # ("reserved2", c_uint16),
  "L",                                          # ("nn", c_uint32),                           # Number of Namespaces
  "S",                                          # ("oncs", c_uint16),                         # Optional NVM Command Support
  "S",                                          # ("fuses", c_uint16),                        # Fused Operation Support
  "C",                                          # ("fna", c_uint8),                           # Format NVM Attributes
  "C",                                          # ("vwc", c_uint8),                           # Volatile Write Cache
  "S",                                          # ("awun", c_uint16),                         # Atomic Write Unit Normal
  "S",                                          # ("awupf", c_uint16),                        # Atomic Write Unit Power Fail
  "C",                                          # ("nvscc", c_uint8),                         # NVM Vendor Specific Command Configuration
  "a#{704 - 531}",                              # ("reserved3", c_uint8 * (704 - 531)),
  "a#{2048 - 704}",                             # ("reserved4", c_uint8 * (2048 - 704)),
  "a#{NVME_IDENTIFY_CONTROLLER_PSD.size * 32}", # ("psd", nvme_identify_controller_psd * 32), # Power State Descriptor
  "a#{NVME_IDENTIFY_CONTROLLER_AMZN_VS.size}",  # ("vs", nvme_identify_controller_amzn_vs)]   # Vendor Specific
].join

# allocate 4K of memory for nvme_id_ctrl struct (see <linux/nvme.h>)
nvme_id_ctrl_buf = ' ' * 4096

NVME_ADMIN_COMMAND = [
  NVME_ADMIN_IDENTIFY,      # opcode
  0,
  0,
  0,
  0,
  " ",
  nvme_id_ctrl_buf,         # pass pointer to the buffer
                            # where nvme_id_ctrl struct will be stored
  0,
  nvme_id_ctrl_buf.size,    # pass size of the buffer
  1,
  0,
  0,
  0,
  0,
  0,
  0
].pack(NVME_ADMIN_CMD_TMPL) # make binary sequence

def nvme_ioctl!(buf, device_path)
  File.open(device_path, "r") do |f|
    # call `ioctl` system call via ruby's `ioctl` wrapper
    # the first arg is the request code, the second is untyped pointer to memory
    f.ioctl(NVME_IOCTL_ADMIN_CMD, NVME_ADMIN_COMMAND)
  end

  buf.unpack(NVME_ID_CTRL_TMPL)
end

Facter.add(:nvme_device_name_map) do
  setcode do
    device_map = {}

    # for each block device which matches glob pattern
    Dir.glob("/dev/nvme[0-9]{[0-9],}") { |device|
      nvme_id_ctrl_struct = nvme_ioctl!(nvme_id_ctrl_buf, device)

      if nvme_id_ctrl_struct[0] == AMZN_NVME_VID or nvme_id_ctrl_struct[3] == AMZN_NVME_EBS_MN
        # amazon device name is a 33th element of array
        # the former 32 bytes are a device name, the latter bytes are reserved
        device_name = nvme_id_ctrl_struct[33].strip

        # prepend "/dev/" if for some reason amazon desided to trim it before
        device_name = "/dev/" + device_name if not device_name.start_with?("/dev/")

        # we are interested in block devices which are named like nvme1n1
        # and not in nvme1 which is nvme subsystem
        device_map[device_name] = device + "n1"
      end
    }

    # puppet facts find | jq '.. | .nvme_device_name_map? | select(. != null)'
    # {
    #   "/dev/xvdf": "/dev/nvme1n1",
    #   "/dev/xvdh": "/dev/nvme3n1",
    #   "/dev/xvdg": "/dev/nvme2n1",
    #   "sda1": "/dev/nvme0n1"
    # }
    device_map
  end
end