/* SPDX-License-Identifier: GPL-2.0-only */
/*
 * Storage Key migration tests
 *
 * There are two variants of this test:
 * - sequential: set storage keys on some pages, migrates the VM and then
 *   verifies that the storage keys are still as we expect.
 * - parallel: start migration and set and check storage keys on some
 *   pages while migration is in process.
 *
 * Copyright IBM Corp. 2022
 *
 * Authors:
 *  Nico Boehr <nrb@linux.ibm.com>
 */

#include <libcflat.h>
#include <migrate.h>
#include <asm/facility.h>
#include <asm/page.h>
#include <asm/mem.h>
#include <asm/barrier.h>
#include <hardware.h>
#include <smp.h>

struct verify_result {
	bool verify_failed;
	union skey expected_key;
	union skey actual_key;
	unsigned long page_mismatch_idx;
	unsigned long page_mismatch_addr;
};

#define NUM_PAGES 128
static uint8_t pagebuf[NUM_PAGES * PAGE_SIZE] __attribute__((aligned(PAGE_SIZE)));

static struct verify_result result;

static unsigned int thread_iters;
static bool thread_should_exit;
static bool thread_exited;

static enum {
	TEST_INVALID,
	TEST_SEQUENTIAL,
	TEST_PARALLEL
} arg_test_to_run;

/*
 * Set storage key test pattern on pagebuf with a seed for the storage keys.
 *
 * Each page's storage key is generated by taking the page's index in pagebuf,
 * XOR-ing that with the given seed and then multipling the result with two.
 *
 * Only the lower seven bits of the seed are considered.
 */
static void set_test_pattern(unsigned char seed)
{
	unsigned char key_to_set;
	unsigned long i;

	for (i = 0; i < NUM_PAGES; i++) {
		/*
		 * Storage keys are 7 bit, lowest bit is always returned as zero
		 * by iske.
		 * This loop will set all 7 bits which means we set fetch
		 * protection as well as reference and change indication for
		 * some keys.
		 */
		key_to_set = (i ^ seed) * 2;
		set_storage_key(pagebuf + i * PAGE_SIZE, key_to_set, 1);
	}
}

/*
 * Verify storage keys on pagebuf.
 * Storage keys must have been set by set_test_pattern on pagebuf before.
 * set_test_pattern must have been called with the same seed value.
 *
 * If storage keys match the expected result, will return a verify_result
 * with verify_failed false. All other fields are then invalid.
 * If there is a mismatch, returned struct will have verify_failed true and will
 * be filled with the details on the first mismatch encountered.
 */
static struct verify_result verify_test_pattern(unsigned char seed)
{
	union skey expected_key, actual_key;
	struct verify_result result = {
		.verify_failed = true
	};
	uint8_t *cur_page;
	unsigned long i;

	for (i = 0; i < NUM_PAGES; i++) {
		cur_page = pagebuf + i * PAGE_SIZE;
		actual_key.val = get_storage_key(cur_page);
		expected_key.val = (i ^ seed) * 2;

		/*
		 * The PoP neither gives a guarantee that the reference bit is
		 * accurate nor that it won't be cleared by hardware. Hence we
		 * don't rely on it and just clear the bits to avoid compare
		 * errors.
		 */
		actual_key.str.rf = 0;
		expected_key.str.rf = 0;

		if (actual_key.val != expected_key.val) {
			result.expected_key.val = expected_key.val;
			result.actual_key.val = actual_key.val;
			result.page_mismatch_idx = i;
			result.page_mismatch_addr = (unsigned long)cur_page;
			return result;
		}
	}

	result.verify_failed = false;
	return result;
}

static void report_verify_result(const struct verify_result * result)
{
	if (result->verify_failed)
		report_fail("page skey mismatch: first page idx = %lu, addr = 0x%lx, "
			    "expected_key = 0x%02x, actual_key = 0x%02x",
			    result->page_mismatch_idx, result->page_mismatch_addr,
			    result->expected_key.val, result->actual_key.val);
	else
		report_pass("skeys match");
}

static void test_skey_migration_sequential(void)
{
	report_prefix_push("sequential");

	set_test_pattern(0);

	migrate_once();

	result = verify_test_pattern(0);
	report_verify_result(&result);

	report_prefix_pop();
}

static void set_skeys_thread(void)
{
	while (!READ_ONCE(thread_should_exit)) {
		set_test_pattern(thread_iters);

		result = verify_test_pattern(thread_iters);

		/*
		 * Always increment even if the verify fails. This ensures primary CPU knows where
		 * we left off and can do an additional verify round after migration finished.
		 */
		thread_iters++;

		if (result.verify_failed)
			break;
	}

	WRITE_ONCE(thread_exited, 1);
}

static void test_skey_migration_parallel(void)
{
	report_prefix_push("parallel");

	if (smp_query_num_cpus() == 1) {
		report_skip("need at least 2 cpus for this test");
		goto error;
	}

	smp_cpu_setup(1, PSW_WITH_CUR_MASK(set_skeys_thread));

	migrate_once();

	WRITE_ONCE(thread_should_exit, 1);

	while (!READ_ONCE(thread_exited))
		;

	/* Ensure we read result and thread_iters below from memory after thread exited */
	mb();
	report_info("thread completed %u iterations", thread_iters);

	report_prefix_push("during migration");
	report_verify_result(&result);
	report_prefix_pop();

	/*
	 * Verification of skeys occurs on the thread. We don't know if we
	 * were still migrating during the verification.
	 * To be sure, make another verification round after the migration
	 * finished to catch skeys which might not have been migrated
	 * correctly.
	 */
	report_prefix_push("after migration");
	assert(thread_iters > 0);
	result = verify_test_pattern(thread_iters - 1);
	report_verify_result(&result);
	report_prefix_pop();

error:
	report_prefix_pop();
}

static void print_usage(void)
{
	report_info("Usage: migration-skey [--parallel|--sequential]");
}

static void parse_args(int argc, char **argv)
{
	if (argc < 2) {
		/* default to sequential since it only needs one cpu */
		arg_test_to_run = TEST_SEQUENTIAL;
		return;
	}

	if (!strcmp("--parallel", argv[1]))
		arg_test_to_run = TEST_PARALLEL;
	else if (!strcmp("--sequential", argv[1]))
		arg_test_to_run = TEST_SEQUENTIAL;
	else
		arg_test_to_run = TEST_INVALID;
}

int main(int argc, char **argv)
{
	report_prefix_push("migration-skey");

	if (test_facility(169)) {
		report_skip("storage key removal facility is active");
		goto error;
	}

	parse_args(argc, argv);

	switch (arg_test_to_run) {
	case TEST_SEQUENTIAL:
		test_skey_migration_sequential();
		break;
	case TEST_PARALLEL:
		test_skey_migration_parallel();
		break;
	default:
		print_usage();
		break;
	}

error:
	migrate_once();
	report_prefix_pop();
	return report_summary();
}