/*
 * This file and its contents are supplied under the terms of the
 * Common Development and Distribution License ("CDDL"), version 1.0.
 * You may only use this file in accordance with the terms of version
 * 1.0 of the CDDL.
 *
 * A full copy of the text of the CDDL should have accompanied this
 * source.  A copy of the CDDL is also available via the Internet at
 * http://www.illumos.org/license/CDDL.
 */

/*
 * Copyright 2024 Oxide Computer Company
 */

/*
 * Iterate over the namespaces and ensure that the information we get in the
 * discovery is the same as when we take a snapshot. To ensure that nothing
 * changes out from under us, we take a controller write lock which will ensure
 * that no modifications can occur.
 *
 * In addition, we want to test that discovery filters work so we first go
 * through and count the different levels.
 */

#include <err.h>
#include <string.h>
#include <umem.h>

#include "libnvme_test_common.h"

static bool
ns_disc_count_cb(nvme_ctrl_t *ctrl, const nvme_ns_disc_t *disc, void *arg)
{
	uint32_t *valp = arg;
	*valp = *valp + 1;
	return (true);
}

static bool
ns_disc_count(nvme_ctrl_t *ctrl, nvme_ns_disc_level_t level, uint32_t exp)
{
	uint32_t count = 0;

	if (!nvme_ns_discover(ctrl, level, ns_disc_count_cb, &count)) {
		libnvme_test_ctrl_warn(ctrl, "failed to discover at level %u",
		    level);
		return (false);
	} else if (count != exp) {
		warnx("TEST FAILED: ns discovery level %u found 0x%x "
		    "namespaces, but expected 0x%x", level, count, exp);
		return (false);
	} else {
		(void) printf("TEST PASSED: ns discovery level %u had correct "
		    "count (0x%x)\n", level, exp);
		return (true);
	}
}

static bool
ns_disc_blkdev_cb(nvme_ctrl_t *ctrl, const nvme_ns_disc_t *disc, void *arg)
{
	int *ret = arg;
	nvme_ns_info_t *info;
	const uint32_t nsid = nvme_ns_disc_nsid(disc);
	const char *addr;

	if (nvme_ns_disc_level(disc) < NVME_NS_DISC_F_BLKDEV) {
		warnx("TEST FAILED: ns %u has level %u, but filtering on "
		    "blkdev (%u)", nsid, nvme_ns_disc_level(disc),
		    NVME_NS_DISC_F_BLKDEV);
		*ret = EXIT_FAILURE;
		return (true);
	}

	if (!nvme_ctrl_ns_info_snap(ctrl, nsid, &info)) {
		libnvme_test_ctrl_warn(ctrl, "failed to get info snapshot for "
		    "nsid %u", nsid);
		*ret = EXIT_FAILURE;
		return (true);
	}

	if (!nvme_ns_info_bd_addr(info, &addr)) {
		libnvme_test_ctrl_warn(ctrl, "failed to get bd addr for nsid "
		    "%u", nsid);
		*ret = EXIT_FAILURE;
	} else if (addr[0] == '\0') {
		warnx("TEST FAILED: nsid %u has invalid bd addr", nsid);
		*ret = EXIT_FAILURE;
	} else {
		(void) printf("TEST PASSED: nsid %u bd addr valid\n", nsid);
	}

	nvme_ns_info_free(info);
	return (true);
}

static bool
ns_disc_guids_cb(nvme_ctrl_t *ctrl, const nvme_ns_disc_t *disc, void *arg)
{
	int *ret = arg;
	nvme_ns_info_t *info;
	const uint32_t nsid = nvme_ns_disc_nsid(disc);
	const nvme_ns_disc_flags_t flags = nvme_ns_disc_flags(disc);
	uint8_t id[16];
	bool bret;

	if (nvme_ns_disc_level(disc) < NVME_NS_DISC_F_ACTIVE) {
		warnx("TEST FAILED: ns %u has level %u, but filtering on "
		    "active (%u)", nsid, nvme_ns_disc_level(disc),
		    NVME_NS_DISC_F_ACTIVE);
		*ret = EXIT_FAILURE;
		return (true);
	}

	if (!nvme_ctrl_ns_info_snap(ctrl, nsid, &info)) {
		libnvme_test_ctrl_warn(ctrl, "failed to get info snapshot for "
		    "nsid %u", nsid);
		*ret = EXIT_FAILURE;
		return (true);
	}

	bret = nvme_ns_info_eui64(info, id);
	if (bret != ((flags & NVME_NS_DISC_F_EUI64_VALID) != 0)) {
		warnx("TEST FAILED: nvme_ns_info_eui64() returned %s, but "
		    "expected %s from discovery information for nsid %u",
		    bret ? "true" : "false",
		    (flags & NVME_NS_DISC_F_EUI64_VALID) != 0 ? "true" :
		    "false", nsid);
		*ret = EXIT_FAILURE;
	} else {
		(void) printf("TEST PASSED: namespace snapshot and discovery "
		    "agreed on EUI64 presence for nsid %u\n", nsid);
	}

	if (bret) {
		const uint8_t *eui64 = nvme_ns_disc_eui64(disc);
		const uint8_t zero[8] = { 0 };

		if (memcmp(eui64, id, sizeof (zero)) != 0) {
			warnx("TEST FAILED: EUI64 differs between "
			    "discovery and info snapshot for nsid %u", nsid);
			*ret = EXIT_FAILURE;
		} else {
			(void) printf("TEST PASSED: EUI64 equal between "
			    "discovery and info snapshot for nsid %u\n", nsid);
		}

		if (memcmp(id, zero, sizeof (zero)) == 0) {
			warnx("TEST FAILED: Found invalid zero EUI64 for nsid "
			    "%u", nsid);
			*ret = EXIT_FAILURE;
		} else {
			(void) printf("TEST PASSED: EUI64 is non-zero for "
			    "nsid %u\n", nsid);
		}
	} else {
		if (nvme_ns_disc_eui64(disc) != NULL) {
			warnx("TEST FAILED: discovery EUI64 was valid, but "
			    "should be NULL for nsid %u", nsid);
			*ret = EXIT_FAILURE;
		} else {
			(void) printf("TEST PASSED: discovery EUI64 correctly "
			    "returned NULL for nsid %u\n", nsid);
		}

		switch (nvme_ns_info_err(info)) {
		case NVME_INFO_ERR_VERSION:
		case NVME_INFO_ERR_MISSING_CAP:
			(void) printf("TEST PASSED: nvme_ns_info_eui64() "
			    "returned a valid error for nsid %u\n", nsid);
			break;
		default:
			warnx("TEST FAILED: nvme_ns_info_eui64() returned an "
			    "invalid error for nsid %u: %s (%u)", nsid,
			    nvme_ns_info_errtostr(info, nvme_ns_info_err(info)),
			    nvme_ns_info_err(info));
			*ret = EXIT_FAILURE;
			break;
		}
	}

	bret = nvme_ns_info_nguid(info, id);
	if (bret != ((flags & NVME_NS_DISC_F_NGUID_VALID) != 0)) {
		warnx("TEST FAILED: nvme_ns_info_nguid() returned %s, but "
		    "expected %s from discovery information for nsid %u",
		    bret ? "true" : "false",
		    (flags & NVME_NS_DISC_F_NGUID_VALID) != 0 ? "true" :
		    "false", nsid);
		*ret = EXIT_FAILURE;
	} else {
		(void) printf("TEST PASSED: namespace snapshot and discovery "
		    "agreed on NGUID presence for nsid %u\n", nsid);
	}

	if (bret) {
		const uint8_t *nguid = nvme_ns_disc_nguid(disc);
		const uint8_t zero[16] = { 0 };

		if (memcmp(nguid, id, sizeof (zero)) != 0) {
			warnx("TEST FAILED: NGUID differs between "
			    "discovery and info snapshot for nsid %u", nsid);
			*ret = EXIT_FAILURE;
		} else {
			(void) printf("TEST PASSED: NGUID equal between "
			    "discovery and info snapshot for nsid %u\n", nsid);
		}

		if (memcmp(id, zero, sizeof (zero)) == 0) {
			warnx("TEST FAILED: Found invalid zero NGUID for nsid "
			    "%u", nsid);
			*ret = EXIT_FAILURE;
		} else {
			(void) printf("TEST PASSED: NGUID is non-zero for "
			    "nsid %u\n", nsid);
		}
	} else {
		if (nvme_ns_disc_nguid(disc) != NULL) {
			warnx("TEST FAILED: discovery NGUID was valid, but "
			    "should be NULL for nsid %u", nsid);
			*ret = EXIT_FAILURE;
		} else {
			(void) printf("TEST PASSED: discovery NGUID correctly "
			    "returned NULL for nsid %u\n", nsid);
		}

		switch (nvme_ns_info_err(info)) {
		case NVME_INFO_ERR_VERSION:
		case NVME_INFO_ERR_MISSING_CAP:
			(void) printf("TEST PASSED: nvme_ns_info_nguid() "
			    "returned a valid error for nsid %u\n", nsid);
			break;
		default:
			warnx("TEST FAILED: nvme_ns_info_nguid() returned an "
			    "invalid error for nsid %u: %s (%u)", nsid,
			    nvme_ns_info_errtostr(info, nvme_ns_info_err(info)),
			    nvme_ns_info_err(info));
			*ret = EXIT_FAILURE;
			break;
		}
	}
	nvme_ns_info_free(info);
	return (true);
}

static bool
ns_disc_level_cb(nvme_ctrl_t *ctrl, const nvme_ns_disc_t *disc, void *arg)
{
	int *ret = arg;
	nvme_ns_info_t *info;
	const uint32_t nsid = nvme_ns_disc_nsid(disc);

	if (!nvme_ctrl_ns_info_snap(ctrl, nsid, &info)) {
		libnvme_test_ctrl_warn(ctrl, "failed to get info snapshot for "
		    "nsid %u", nsid);
		*ret = EXIT_FAILURE;
		return (true);
	}

	if (nvme_ns_disc_level(disc) != nvme_ns_info_level(info)) {
		warnx("TEST FAILED: discovery and ns info snapshot disagree "
		    "on discovery level: disc has %u, info has %u",
		    nvme_ns_disc_level(disc), nvme_ns_info_level(info));
		*ret = EXIT_FAILURE;
	} else {
		(void) printf("TEST PASSED: discovery and ns info snapshot "
		    "agree for nsid %u\n", nsid);
	}

	nvme_ns_info_free(info);
	return (true);
}

static bool
ns_disc_bad_disc_init(nvme_ctrl_t *ctrl, nvme_ns_disc_level_t level,
    nvme_ns_iter_t **iterp, nvme_err_t exp_err, const char *desc)
{
	if (nvme_ns_discover_init(ctrl, level, iterp)) {
		warnx("TEST FAILED: nvme_ns_discover_init() erroneously "
		    "passed despite %s", desc);
		nvme_ns_discover_fini(*iterp);
		return (false);
	} else if (nvme_ctrl_err(ctrl) != exp_err) {
		warnx("TEST FAILED: nvme_ns_discover_init() returned "
		    "wrong error %s (0x%x), not %s (0x%x)",
		    nvme_ctrl_errtostr(ctrl, nvme_ctrl_err(ctrl)),
		    nvme_ctrl_err(ctrl), nvme_ctrl_errtostr(ctrl,
		    exp_err), exp_err);
		return (false);
	} else {
		(void) printf("TEST PASSED: nvme_ns_discover_init() failed "
		    "correctly for %s\n", desc);
		return (true);
	}
}

static bool
ns_disc_bad_disc(nvme_ctrl_t *ctrl, nvme_ns_disc_level_t level,
    nvme_ns_disc_f func, nvme_err_t exp_err, const char *desc)
{
	if (nvme_ns_discover(ctrl, level, func, NULL)) {
		warnx("TEST FAILED: nvme_ns_discover() erroneously "
		    "passed despite %s", desc);
		return (false);
	} else if (nvme_ctrl_err(ctrl) != exp_err) {
		warnx("TEST FAILED: nvme_ns_discover() returned "
		    "wrong error %s (0x%x), not %s (0x%x)",
		    nvme_ctrl_errtostr(ctrl, nvme_ctrl_err(ctrl)),
		    nvme_ctrl_err(ctrl), nvme_ctrl_errtostr(ctrl,
		    exp_err), exp_err);
		return (false);
	} else {
		(void) printf("TEST PASSED: nvme_ns_discover() failed "
		    "correctly for %s\n", desc);
		return (true);
	}
}

static bool
ns_disc_nop_cb(nvme_ctrl_t *ctrl, const nvme_ns_disc_t *disc, void *arg)
{
	return (true);
}

int
main(void)
{
	int ret = EXIT_SUCCESS;
	nvme_t *nvme;
	nvme_ctrl_t *ctrl;
	nvme_ctrl_info_t *info;
	nvme_iter_t iret;
	nvme_ns_iter_t *iter;
	const nvme_ns_disc_t *disc;
	uint32_t nbd = 0, nni = 0, nact = 0, nalloc = 0, nns = 0;

	libnvme_test_init(&nvme, &ctrl);
	if (!nvme_ctrl_lock(ctrl, NVME_LOCK_L_WRITE, NVME_LOCK_F_DONT_BLOCK)) {
		libnvme_test_ctrl_fatal(ctrl, "failed to obtain write lock");
	}

	if (!nvme_ns_discover_init(ctrl, NVME_NS_DISC_F_ALL, &iter)) {
		libnvme_test_ctrl_fatal(ctrl, "failed to initialize initial "
		    "ns discovery");
	}

	while ((iret = nvme_ns_discover_step(iter, &disc)) == NVME_ITER_VALID) {
		switch (nvme_ns_disc_level(disc)) {
		case NVME_NS_DISC_F_BLKDEV:
			nbd++;
			/* FALLTHROUGH */
		case NVME_NS_DISC_F_NOT_IGNORED:
			nni++;
			/* FALLTHROUGH */
		case NVME_NS_DISC_F_ACTIVE:
			nact++;
			/* FALLTHROUGH */
		case NVME_NS_DISC_F_ALLOCATED:
			nalloc++;
			/* FALLTHROUGH */
		case NVME_NS_DISC_F_ALL:
			nns++;
			break;
		}
	}

	nvme_ns_discover_fini(iter);
	if (iret != NVME_ITER_DONE) {
		libnvme_test_ctrl_fatal(ctrl, "initial ns discovery failed");
	}

	if (!nvme_ctrl_info_snap(ctrl, &info)) {
		libnvme_test_ctrl_fatal(ctrl, "failed to get info snapshot");
	}

	if (nns != nvme_ctrl_info_nns(info)) {
		warnx("TEST FAILED: discovery found %u namespaces, but the "
		    "identify controller suggests there are %u", nns,
		    nvme_ctrl_info_nns(info));
	}

	if (!ns_disc_count(ctrl, NVME_NS_DISC_F_BLKDEV, nbd)) {
		ret = EXIT_FAILURE;
	}

	if (!ns_disc_count(ctrl, NVME_NS_DISC_F_NOT_IGNORED, nni)) {
		ret = EXIT_FAILURE;
	}

	if (!ns_disc_count(ctrl, NVME_NS_DISC_F_ACTIVE, nact)) {
		ret = EXIT_FAILURE;
	}

	if (!ns_disc_count(ctrl, NVME_NS_DISC_F_ALLOCATED, nalloc)) {
		ret = EXIT_FAILURE;
	}

	if (!ns_disc_count(ctrl, NVME_NS_DISC_F_ALL, nns)) {
		ret = EXIT_FAILURE;
	}

	/*
	 * For anything that has a blkdev address, ensure that our info snapshot
	 * has a valid blkdev address.
	 */
	if (!nvme_ns_discover(ctrl, NVME_NS_DISC_F_BLKDEV, ns_disc_blkdev_cb,
	    &ret)) {
		libnvme_test_ctrl_warn(ctrl, "discovery failed for blkdev "
		    "test");
		ret = EXIT_FAILURE;
	}

	/*
	 * For anything active, check if there are guids and that the
	 * information snapshot matches the same logic.
	 */
	if (!nvme_ns_discover(ctrl, NVME_NS_DISC_F_ACTIVE, ns_disc_guids_cb,
	    &ret)) {
		libnvme_test_ctrl_warn(ctrl, "discovery failed for guids "
		    "test");
		ret = EXIT_FAILURE;
	}

	/*
	 * For everything, make sure the levels match.
	 */
	if (!nvme_ns_discover(ctrl, NVME_NS_DISC_F_ALL, ns_disc_level_cb,
	    &ret)) {
		libnvme_test_ctrl_warn(ctrl, "discovery failed for levels "
		    "test");
		ret = EXIT_FAILURE;
	}

	nvme_ctrl_unlock(ctrl);

	if (!ns_disc_bad_disc_init(ctrl, INT32_MAX, &iter, NVME_ERR_BAD_FLAG,
	    "invalid level")) {
		ret = EXIT_FAILURE;
	}

	if (!ns_disc_bad_disc_init(ctrl, NVME_NS_DISC_F_ALL, NULL,
	    NVME_ERR_BAD_PTR, "invalid iter pointer")) {
		ret = EXIT_FAILURE;
	}

	if (!ns_disc_bad_disc(ctrl, UINT32_MAX, ns_disc_nop_cb,
	    NVME_ERR_BAD_FLAG, "invalid level")) {
		ret = EXIT_FAILURE;
	}

	if (!ns_disc_bad_disc(ctrl, NVME_NS_DISC_F_ALL, NULL,
	    NVME_ERR_BAD_PTR, "invalid function pointer")) {
		ret = EXIT_FAILURE;
	}

	umem_setmtbf(1);
	if (!ns_disc_bad_disc_init(ctrl, NVME_NS_DISC_F_ALL, &iter,
	    NVME_ERR_NO_MEM, "no memory")) {
		ret = EXIT_FAILURE;
	}

	if (!ns_disc_bad_disc(ctrl, NVME_NS_DISC_F_ALL, ns_disc_nop_cb,
	    NVME_ERR_NO_MEM, "no memory")) {
		ret = EXIT_FAILURE;
	}
	umem_setmtbf(0);

	if (ret == EXIT_SUCCESS) {
		(void) printf("All tests passed successfully\n");
	}

	nvme_ctrl_info_free(info);
	nvme_ctrl_fini(ctrl);
	nvme_fini(nvme);
	return (ret);
}