/*
 * Copyright (c) 2024 Apple Inc. All rights reserved.
 *
 * @APPLE_OSREFERENCE_LICENSE_HEADER_START@
 *
 * This file contains Original Code and/or Modifications of Original Code
 * as defined in and that are subject to the Apple Public Source License
 * Version 2.0 (the 'License'). You may not use this file except in
 * compliance with the License. The rights granted to you under the License
 * may not be used to create, or enable the creation or redistribution of,
 * unlawful or unlicensed copies of an Apple operating system, or to
 * circumvent, violate, or enable the circumvention or violation of, any
 * terms of an Apple operating system software license agreement.
 *
 * Please obtain a copy of the License at
 * http://www.opensource.apple.com/apsl/ and read it before using this file.
 *
 * The Original Code and all software distributed under the License are
 * distributed on an 'AS IS' basis, WITHOUT WARRANTY OF ANY KIND, EITHER
 * EXPRESS OR IMPLIED, AND APPLE HEREBY DISCLAIMS ALL SUCH WARRANTIES,
 * INCLUDING WITHOUT LIMITATION, ANY WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE, QUIET ENJOYMENT OR NON-INFRINGEMENT.
 * Please see the License for the specific language governing rights and
 * limitations under the License.
 *
 * @APPLE_OSREFERENCE_LICENSE_HEADER_END@
 */

/*
 * try_read_write.c
 *
 * Helper functions for userspace tests to read or write memory and
 * verify that EXC_BAD_ACCESS is or is not generated by that operation.
 */

#include <assert.h>
#include <stdbool.h>
#include <stdatomic.h>
#include <ptrauth.h>
#include <darwintest.h>
#include <dispatch/dispatch.h>

#include "exc_helpers.h"
#include "try_read_write.h"

/*
 * -- Implementation overview --
 *
 * try_read_byte() and try_write_byte() operate by performing
 * a read or write instruction with a Mach exception handler
 * in place.
 *
 * The exception handler catches EXC_BAD_ACCESS. If the bad access
 * came from our designated read or write instructions then it
 * records the exception that occurred to thread-local storage
 * and moves that thread's program counter to resume execution
 * and recover from the exception.
 *
 * Unrecognized exceptions, and EXC_BAD_ACCESS exceptions from
 * unrecognized instructions, either go uncaught or are caught and
 * re-raised. In either case they lead to an ordinary crash. This
 * means we don't get false positives where the test expects one
 * crash but incorrectly passes after crashing in some unrelated way.
 * We can be precise about what the fault was and where it came from.
 *
 * We use Mach exceptions instead of signals because
 * on watchOS signal handlers do not receive the thread
 * state so they cannot recover from the signal.
 *
 * try_read_write_exception_handler()
 *     our exception handler, installed using tests/exc_helpers.c
 *
 * read_byte() and write_byte()
 *     our designated read and write instructions, recognized by
 *     the exception handler and specially structured to allow
 *     recovery by changing the PC
 *
 * try_read_write_thread_t
 *     thread-local storage to record the caught exception
 */

static dispatch_once_t try_read_write_initializer;
static mach_port_t try_read_write_exc_port;

/*
 * Bespoke thread-local storage for threads inside try_read_write.
 * We can't use pthread local storage because the Mach exception
 * handler needs to access it and that exception handler runs on
 * a different thread.
 *
 * Access by the Mach exception thread is safe because the real thread
 * is suspended at that point. (This scheme would be unsound if the
 * real thread raised an exception while manipulating the thread-local
 * data, but we don't try to cover that case.)
 */
typedef struct {
	mach_port_t thread;
	kern_return_t exception_kr;  /* EXC_BAD_ADDRESS sub-code */
	uint64_t exception_pc;       /* PC of faulting instruction */
	uint64_t exception_memory;   /* Memory address of faulting access */
} try_read_write_thread_t;

#define TRY_READ_WRITE_MAX_THREADS 128
static pthread_mutex_t try_read_write_thread_list_mutex = PTHREAD_MUTEX_INITIALIZER;
static unsigned try_read_write_thread_count = 0;
static try_read_write_thread_t try_read_write_thread_list[TRY_READ_WRITE_MAX_THREADS];
static __thread try_read_write_thread_t *try_read_write_thread_self;

/*
 * Look up the try_read_write_thread_t for a Mach thread.
 * If create == true and no info was found, add it to the list.
 * Returns NULL if no info was found and create == false.
 */
static __attribute__((overloadable))
try_read_write_thread_t *
thread_info_for_mach_thread(mach_port_t thread_port, bool create)
{
	/* first look for a cached value in real thread-local storage */
	if (mach_thread_self() == thread_port) {
		try_read_write_thread_t *info = try_read_write_thread_self;
		if (info) {
			return info;
		}
	}

	int err = pthread_mutex_lock(&try_read_write_thread_list_mutex);
	assert(err == 0);

	/* search the list */
	for (unsigned i = 0; i < try_read_write_thread_count; i++) {
		try_read_write_thread_t *info = &try_read_write_thread_list[i];
		if (info->thread == thread_port) {
			pthread_mutex_unlock(&try_read_write_thread_list_mutex);
			if (mach_thread_self() == thread_port) {
				try_read_write_thread_self = info;
			}
			return info;
		}
	}

	/* not in list - create if requested */
	if (create) {
		assert(try_read_write_thread_count < TRY_READ_WRITE_MAX_THREADS);
		try_read_write_thread_t *info = &try_read_write_thread_list[try_read_write_thread_count++];
		info->thread = thread_port;
		info->exception_kr = 0;
		pthread_mutex_unlock(&try_read_write_thread_list_mutex);
		if (mach_thread_self() == thread_port) {
			try_read_write_thread_self = info;
		}
		return info;
	}

	pthread_mutex_unlock(&try_read_write_thread_list_mutex);
	return NULL;
}

static __attribute__((overloadable))
try_read_write_thread_t *
thread_info_for_mach_thread(mach_port_t thread_port)
{
	return thread_info_for_mach_thread(thread_port, false /* create */);
}


/*
 * read_byte() and write_byte() are functions that
 * read or write memory as their first instruction.
 * Used to test memory access that may provoke an exception.
 *
 * try_read_write_exception_handler() below checks if the exception PC
 * is equal to one of these functions. The first instruction must be
 * the memory access instruction.
 *
 * try_read_write_exception_handler() below increments the PC by four bytes.
 * The memory access instruction must be padded to exactly four bytes.
 */

static uint64_t __attribute__((naked))
read_byte(mach_vm_address_t addr)
{
#if __arm64__
	asm("\n ldrb w0, [x0]"
            "\n ret");
#elif __x86_64__
	asm("\n movb (%rdi), %al"
            "\n nop"  /* pad load to four bytes */
            "\n nop"
            "\n ret");
#else
#       error unknown architecture
#endif
}

static void __attribute__((naked))
write_byte(mach_vm_address_t addr, uint8_t value)
{
#if __arm64__
	asm("\n strb w1, [x0]"
            "\n ret");
#elif __x86_64__
	asm("\n movb %sil, (%rdi)"
            "\n nop"  /* pad store to four bytes */
            "\n ret");
#else
#       error unknown architecture
#endif
}


/*
 * Mach exception handler for EXC_BAD_ACCESS called by exc_helpers.
 * Returns the number of bytes to advance the PC to resolve the exception.
 */
static size_t
try_read_write_exception_handler(
	__unused mach_port_t task,
	mach_port_t thread,
	exception_type_t exception,
	mach_exception_data_t codes,
	uint64_t exception_pc)
{
	assert(exception == EXC_BAD_ACCESS);
	try_read_write_thread_t *info = thread_info_for_mach_thread(thread);
	assert(info);  /* we do not expect exceptions from other threads */

	uint64_t read_byte_pc  = (uint64_t)ptrauth_strip(&read_byte, ptrauth_key_function_pointer);
	uint64_t write_byte_pc = (uint64_t)ptrauth_strip(&write_byte, ptrauth_key_function_pointer);

	if (exception_pc != read_byte_pc && exception_pc != write_byte_pc) {
		/* this exception isn't one of ours - re-raise it */
		if (verbose_exc_helper) {
			T_LOG("not a try_read_write exception");
		}
		return EXC_HELPER_HALT;
	}

	assert(info->exception_kr == 0); /* no nested exceptions allowed */

	info->exception_pc = exception_pc;
	info->exception_kr = codes[0];
	info->exception_memory = codes[1];
	if (verbose_exc_helper) {
		T_LOG("try_read_write exception: pc 0x%llx kr %d mem 0x%llx",
		    info->exception_pc, info->exception_kr, info->exception_memory);
	}

	/* advance pc by 4 bytes to recover */
	return 4;
}

/*
 * Create an exc_helpers exception handler port and thread,
 * and install the exception handler port on this thread.
 */
static void
initialize_exception_handlers(void)
{
	try_read_write_exc_port = create_exception_port(EXC_MASK_BAD_ACCESS);
	repeat_exception_handler(try_read_write_exc_port, try_read_write_exception_handler);
}

/*
 * Begin try_read_write exception handling on this thread.
 */
static void
begin_expected_exceptions(void)
{
	dispatch_once(&try_read_write_initializer, ^{
		initialize_exception_handlers();
	});

	try_read_write_thread_t *info = try_read_write_thread_self;
	if (!info) {
		set_thread_exception_port(try_read_write_exc_port, EXC_MASK_BAD_ACCESS);
		info = thread_info_for_mach_thread(mach_thread_self(), true /* create */);
	}

	info->exception_kr = 0;
	info->exception_pc = 0;
	info->exception_memory = 0;
}

/*
 * End try_read_write exception handling on this thread.
 * Returns the caught exception data, if any.
 */
static void
end_expected_exceptions(
	kern_return_t * const out_kr,
	uint64_t * const out_pc,
	uint64_t * const out_memory)
{
	try_read_write_thread_t *info = try_read_write_thread_self;
	assert(info);
	*out_kr = info->exception_kr;
	*out_pc = info->exception_pc;
	*out_memory = info->exception_memory;
}


extern bool
try_read_byte(
	mach_vm_address_t addr,
	uint8_t * const out_byte,
	kern_return_t * const out_error)
{
	kern_return_t exception_kr;
	uint64_t exception_pc;
	uint64_t exception_memory;

	begin_expected_exceptions();
	*out_byte = read_byte(addr);
	end_expected_exceptions(&exception_kr, &exception_pc, &exception_memory);

	/*
	 * pc was verified inside the exception handler.
	 * kr will be verified by the caller.
	 * Verify address here.
	 */

	if (exception_kr != KERN_SUCCESS) {
		assert(exception_memory == addr);
	}

	*out_error = exception_kr;
	return exception_kr == 0;
}

extern bool
try_write_byte(
	mach_vm_address_t addr,
	uint8_t byte,
	kern_return_t * const out_error)
{
	kern_return_t exception_kr;
	uint64_t exception_pc;
	uint64_t exception_memory;

	begin_expected_exceptions();
	write_byte(addr, byte);
	end_expected_exceptions(&exception_kr, &exception_pc, &exception_memory);

	/*
	 * pc was verified inside the exception handler.
	 * kr will be verified by the caller.
	 * Verify address here.
	 */

	if (exception_kr != KERN_SUCCESS) {
		assert(exception_memory == addr);
	}

	*out_error = exception_kr;
	return exception_kr == 0;
}