/*
 * Copyright (C) 2018 Richard Hughes <richard@hughsie.com>
 *
 * SPDX-License-Identifier: LGPL-2.1+
 */

#include "config.h"

#include <fwupdplugin.h>

#include <linux/nvme_ioctl.h>
#include <sys/ioctl.h>

#include "fu-nvme-common.h"
#include "fu-nvme-device.h"

#define FU_NVME_ID_CTRL_SIZE 0x1000

struct _FuNvmeDevice {
	FuUdevDevice parent_instance;
	guint pci_depth;
	guint64 write_block_size;
};

/**
 * FU_NVME_DEVICE_FLAG_FORCE_ALIGN:
 *
 * Force alignment of the firmware file.
 */
#define FU_NVME_DEVICE_FLAG_FORCE_ALIGN (1 << 0)

G_DEFINE_TYPE(FuNvmeDevice, fu_nvme_device, FU_TYPE_UDEV_DEVICE)

static void
fu_nvme_device_to_string(FuDevice *device, guint idt, GString *str)
{
	FuNvmeDevice *self = FU_NVME_DEVICE(device);
	FU_DEVICE_CLASS(fu_nvme_device_parent_class)->to_string(device, idt, str);
	fu_common_string_append_ku(str, idt, "PciDepth", self->pci_depth);
}

/* @addr_start and @addr_end are *inclusive* to match the NMVe specification */
static gchar *
fu_nvme_device_get_string_safe(const guint8 *buf, guint16 addr_start, guint16 addr_end)
{
	GString *str;

	g_return_val_if_fail(buf != NULL, NULL);
	g_return_val_if_fail(addr_start < addr_end, NULL);

	str = g_string_new_len(NULL, addr_end + addr_start + 1);
	for (guint16 i = addr_start; i <= addr_end; i++) {
		gchar tmp = (gchar)buf[i];
		/* skip leading spaces */
		if (g_ascii_isspace(tmp) && str->len == 0)
			continue;
		if (g_ascii_isprint(tmp))
			g_string_append_c(str, tmp);
	}

	/* nothing found */
	if (str->len == 0) {
		g_string_free(str, TRUE);
		return NULL;
	}
	return g_strchomp(g_string_free(str, FALSE));
}

static gchar *
fu_nvme_device_get_guid_safe(const guint8 *buf, guint16 addr_start)
{
	if (!fu_common_guid_is_plausible(buf + addr_start))
		return NULL;
	return fwupd_guid_to_string((const fwupd_guid_t *)(buf + addr_start),
				    FWUPD_GUID_FLAG_MIXED_ENDIAN);
}

static gboolean
fu_nvme_device_submit_admin_passthru(FuNvmeDevice *self, struct nvme_admin_cmd *cmd, GError **error)
{
	gint rc = 0;
	guint32 err;

	/* submit admin command */
	if (!fu_udev_device_ioctl(FU_UDEV_DEVICE(self),
				  NVME_IOCTL_ADMIN_CMD,
				  (guint8 *)cmd,
				  &rc,
				  error)) {
		g_prefix_error(error, "failed to issue admin command 0x%02x: ", cmd->opcode);
		return FALSE;
	}

	/* check the error code */
	err = rc & 0x3ff;
	switch (err) {
	case NVME_SC_SUCCESS:
	/* devices are always added with _NEEDS_REBOOT, so ignore */
	case NVME_SC_FW_NEEDS_CONV_RESET:
	case NVME_SC_FW_NEEDS_SUBSYS_RESET:
	case NVME_SC_FW_NEEDS_RESET:
		return TRUE;
	default:
		break;
	}
	g_set_error(error,
		    FWUPD_ERROR,
		    FWUPD_ERROR_NOT_SUPPORTED,
		    "Not supported: %s",
		    fu_nvme_status_to_string(err));
	return FALSE;
}

static gboolean
fu_nvme_device_identify_ctrl(FuNvmeDevice *self, guint8 *data, GError **error)
{
	struct nvme_admin_cmd cmd = {
	    .opcode = 0x06,
	    .nsid = 0x00,
	    .addr = 0x0, /* memory address of data */
	    .data_len = FU_NVME_ID_CTRL_SIZE,
	    .cdw10 = 0x01,
	    .cdw11 = 0x00,
	};
	memcpy(&cmd.addr, &data, sizeof(gpointer));
	return fu_nvme_device_submit_admin_passthru(self, &cmd, error);
}

static gboolean
fu_nvme_device_fw_commit(FuNvmeDevice *self,
			 guint8 slot,
			 guint8 action,
			 guint8 bpid,
			 GError **error)
{
	struct nvme_admin_cmd cmd = {
	    .opcode = 0x10,
	    .cdw10 = (bpid << 31) | (action << 3) | slot,
	};
	return fu_nvme_device_submit_admin_passthru(self, &cmd, error);
}

static gboolean
fu_nvme_device_fw_download(FuNvmeDevice *self,
			   guint32 addr,
			   const guint8 *data,
			   guint32 data_sz,
			   GError **error)
{
	struct nvme_admin_cmd cmd = {
	    .opcode = 0x11,
	    .addr = 0x0, /* memory address of data */
	    .data_len = data_sz,
	    .cdw10 = (data_sz >> 2) - 1, /* convert to DWORDs */
	    .cdw11 = addr >> 2,		 /* convert to DWORDs */
	};
	memcpy(&cmd.addr, &data, sizeof(gpointer));
	return fu_nvme_device_submit_admin_passthru(self, &cmd, error);
}

static void
fu_nvme_device_parse_cns_maybe_dell(FuNvmeDevice *self, const guint8 *buf)
{
	g_autofree gchar *component_id = NULL;
	g_autofree gchar *devid = NULL;
	g_autofree gchar *guid_efi = NULL;
	g_autofree gchar *guid = NULL;

	/* add extra component ID if set */
	component_id = fu_nvme_device_get_string_safe(buf, 0xc36, 0xc3d);
	if (component_id == NULL || !g_str_is_ascii(component_id) || strlen(component_id) < 6) {
		g_debug("invalid component ID, skipping");
		return;
	}

	/* do not add the FuUdevDevice instance IDs as generic firmware
	 * should not be used on these OEM-specific devices */
	fu_device_add_internal_flag(FU_DEVICE(self), FU_DEVICE_INTERNAL_FLAG_NO_AUTO_INSTANCE_IDS);

	/* add instance ID *and* GUID as using no-auto-instance-ids */
	devid = g_strdup_printf("STORAGE-DELL-%s", component_id);
	fu_device_add_instance_id(FU_DEVICE(self), devid);
	guid = fwupd_guid_hash_string(devid);
	fu_device_add_guid(FU_DEVICE(self), guid);

	/* also add the EFI GUID */
	guid_efi = fu_nvme_device_get_guid_safe(buf, 0x0c26);
	if (guid_efi != NULL)
		fu_device_add_guid(FU_DEVICE(self), guid_efi);
}

static gboolean
fu_nvme_device_parse_cns(FuNvmeDevice *self, const guint8 *buf, gsize sz, GError **error)
{
	guint8 fawr;
	guint8 fwug;
	guint8 nfws;
	guint8 s1ro;
	g_autofree gchar *gu = NULL;
	g_autofree gchar *mn = NULL;
	g_autofree gchar *sn = NULL;
	g_autofree gchar *sr = NULL;

	/* wrong size */
	if (sz != FU_NVME_ID_CTRL_SIZE) {
		g_set_error(error,
			    G_IO_ERROR,
			    G_IO_ERROR_FAILED,
			    "failed to parse blob, expected 0x%04x bytes",
			    (guint)FU_NVME_ID_CTRL_SIZE);
		return FALSE;
	}

	/* get sanitiezed string from CNS -- see the following doc for offsets:
	 * NVM-Express-1_3c-2018.05.24-Ratified.pdf */
	sn = fu_nvme_device_get_string_safe(buf, 4, 23);
	if (sn != NULL)
		fu_device_set_serial(FU_DEVICE(self), sn);
	mn = fu_nvme_device_get_string_safe(buf, 24, 63);
	if (mn != NULL)
		fu_device_set_name(FU_DEVICE(self), mn);
	sr = fu_nvme_device_get_string_safe(buf, 64, 71);
	if (sr != NULL)
		fu_device_set_version(FU_DEVICE(self), sr);

	/* firmware update granularity (FWUG) */
	fwug = buf[319];
	if (fwug != 0x00 && fwug != 0xff)
		self->write_block_size = ((guint64)fwug) * 0x1000;

	/* firmware slot information */
	fawr = (buf[260] & 0x10) >> 4;
	nfws = (buf[260] & 0x0e) >> 1;
	s1ro = buf[260] & 0x01;
	g_debug("fawr: %u, nr fw slots: %u, slot1 r/o: %u", fawr, nfws, s1ro);

	/* FRU globally unique identifier (FGUID) */
	gu = fu_nvme_device_get_guid_safe(buf, 127);
	if (gu != NULL)
		fu_device_add_guid(FU_DEVICE(self), gu);

	/* Dell helpfully provide an EFI GUID we can use in the vendor offset,
	 * but don't have a header or any magic we can use -- so check if the
	 * component ID looks plausible and the GUID is "sane" */
	fu_nvme_device_parse_cns_maybe_dell(self, buf);

	/* fall back to the device description */
	if (fu_device_get_guids(FU_DEVICE(self))->len == 0) {
		g_debug("no vendor GUID, falling back to mn");
		fu_device_add_instance_id(FU_DEVICE(self), mn);
	}
	return TRUE;
}

static void
fu_nvme_device_dump(const gchar *title, const guint8 *buf, gsize sz)
{
	if (g_getenv("FWUPD_NVME_VERBOSE") == NULL)
		return;
	g_print("%s (%" G_GSIZE_FORMAT "):", title, sz);
	for (gsize i = 0; i < sz; i++) {
		if (i % 64 == 0)
			g_print("\naddr 0x%04x: ", (guint)i);
		g_print("%02x", buf[i]);
	}
	g_print("\n");
}

static gboolean
fu_nvme_device_probe(FuDevice *device, GError **error)
{
	FuNvmeDevice *self = FU_NVME_DEVICE(device);

	/* FuUdevDevice->probe */
	if (!FU_DEVICE_CLASS(fu_nvme_device_parent_class)->probe(device, error))
		return FALSE;

	/* set the physical ID */
	if (!fu_udev_device_set_physical_id(FU_UDEV_DEVICE(device), "pci", error))
		return FALSE;

	/* look at the PCI depth to work out if in an external enclosure */
	self->pci_depth = fu_udev_device_get_slot_depth(FU_UDEV_DEVICE(device), "pci");
	if (self->pci_depth <= 2) {
		fu_device_add_flag(device, FWUPD_DEVICE_FLAG_INTERNAL);
		fu_device_add_flag(device, FWUPD_DEVICE_FLAG_USABLE_DURING_UPDATE);
	}

	/* all devices need at least a warm reset, but some quirked drives
	 * need a full "cold" shutdown and startup */
	if (!fu_device_has_flag(self, FWUPD_DEVICE_FLAG_NEEDS_SHUTDOWN))
		fu_device_add_flag(device, FWUPD_DEVICE_FLAG_NEEDS_REBOOT);

	return TRUE;
}

static gboolean
fu_nvme_device_setup(FuDevice *device, GError **error)
{
	FuNvmeDevice *self = FU_NVME_DEVICE(device);
	guint8 buf[FU_NVME_ID_CTRL_SIZE] = {0x0};

	/* get and parse CNS */
	if (!fu_nvme_device_identify_ctrl(self, buf, error)) {
		g_prefix_error(error,
			       "failed to identify %s: ",
			       fu_device_get_physical_id(FU_DEVICE(self)));
		return FALSE;
	}
	fu_nvme_device_dump("CNS", buf, sizeof(buf));
	if (!fu_nvme_device_parse_cns(self, buf, sizeof(buf), error))
		return FALSE;

	/* success */
	return TRUE;
}

static gboolean
fu_nvme_device_write_firmware(FuDevice *device,
			      FuFirmware *firmware,
			      FwupdInstallFlags flags,
			      GError **error)
{
	FuNvmeDevice *self = FU_NVME_DEVICE(device);
	g_autoptr(GBytes) fw2 = NULL;
	g_autoptr(GBytes) fw = NULL;
	g_autoptr(GPtrArray) chunks = NULL;
	guint64 block_size = self->write_block_size > 0 ? self->write_block_size : 0x1000;

	/* get default image */
	fw = fu_firmware_get_bytes(firmware, error);
	if (fw == NULL)
		return FALSE;

	/* some vendors provide firmware files whose sizes are not multiples
	 * of blksz *and* the device won't accept blocks of different sizes */
	if (fu_device_has_private_flag(device, FU_NVME_DEVICE_FLAG_FORCE_ALIGN)) {
		fw2 = fu_common_bytes_align(fw, block_size, 0xff);
	} else {
		fw2 = g_bytes_ref(fw);
	}

	/* build packets */
	chunks = fu_chunk_array_new_from_bytes(fw2,
					       0x00,	    /* start_addr */
					       0x00,	    /* page_sz */
					       block_size); /* block size */

	/* write each block */
	fu_device_set_status(device, FWUPD_STATUS_DEVICE_WRITE);
	for (guint i = 0; i < chunks->len; i++) {
		FuChunk *chk = g_ptr_array_index(chunks, i);
		if (!fu_nvme_device_fw_download(self,
						fu_chunk_get_address(chk),
						fu_chunk_get_data(chk),
						fu_chunk_get_data_sz(chk),
						error)) {
			g_prefix_error(error, "failed to write chunk %u: ", i);
			return FALSE;
		}
		fu_device_set_progress_full(device, (gsize)i, (gsize)chunks->len + 1);
	}

	/* commit */
	if (!fu_nvme_device_fw_commit(self,
				      0x00, /* let controller choose */
				      0x01, /* download replaces, activated on reboot */
				      0x00, /* boot partition identifier */
				      error)) {
		g_prefix_error(error, "failed to commit to auto slot: ");
		return FALSE;
	}

	/* success! */
	fu_device_set_progress(device, 100);
	return TRUE;
}

static gboolean
fu_nvme_device_set_quirk_kv(FuDevice *device, const gchar *key, const gchar *value, GError **error)
{
	FuNvmeDevice *self = FU_NVME_DEVICE(device);
	if (g_strcmp0(key, "NvmeBlockSize") == 0) {
		self->write_block_size = fu_common_strtoull(value);
		return TRUE;
	}

	g_set_error_literal(error, G_IO_ERROR, G_IO_ERROR_NOT_SUPPORTED, "quirk key not supported");
	return FALSE;
}

static void
fu_nvme_device_init(FuNvmeDevice *self)
{
	fu_device_add_flag(FU_DEVICE(self), FWUPD_DEVICE_FLAG_REQUIRE_AC);
	fu_device_add_flag(FU_DEVICE(self), FWUPD_DEVICE_FLAG_UPDATABLE);
	fu_device_set_version_format(FU_DEVICE(self), FWUPD_VERSION_FORMAT_PLAIN);
	fu_device_set_summary(FU_DEVICE(self), "NVM Express solid state drive");
	fu_device_add_icon(FU_DEVICE(self), "drive-harddisk");
	fu_device_add_protocol(FU_DEVICE(self), "org.nvmexpress");
	fu_udev_device_set_flags(FU_UDEV_DEVICE(self),
				 FU_UDEV_DEVICE_FLAG_OPEN_READ |
				     FU_UDEV_DEVICE_FLAG_VENDOR_FROM_PARENT);
	fu_device_register_private_flag(FU_DEVICE(self),
					FU_NVME_DEVICE_FLAG_FORCE_ALIGN,
					"force-align");
}

static void
fu_nvme_device_finalize(GObject *object)
{
	G_OBJECT_CLASS(fu_nvme_device_parent_class)->finalize(object);
}

static void
fu_nvme_device_class_init(FuNvmeDeviceClass *klass)
{
	GObjectClass *object_class = G_OBJECT_CLASS(klass);
	FuDeviceClass *klass_device = FU_DEVICE_CLASS(klass);
	object_class->finalize = fu_nvme_device_finalize;
	klass_device->to_string = fu_nvme_device_to_string;
	klass_device->set_quirk_kv = fu_nvme_device_set_quirk_kv;
	klass_device->setup = fu_nvme_device_setup;
	klass_device->write_firmware = fu_nvme_device_write_firmware;
	klass_device->probe = fu_nvme_device_probe;
}

FuNvmeDevice *
fu_nvme_device_new_from_blob(FuContext *ctx, const guint8 *buf, gsize sz, GError **error)
{
	g_autoptr(FuNvmeDevice) self = NULL;
	self = g_object_new(FU_TYPE_NVME_DEVICE, "context", ctx, NULL);
	if (!fu_nvme_device_parse_cns(self, buf, sz, error))
		return NULL;
	return g_steal_pointer(&self);
}