/*
 * 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 storage prepared by
 * try_read_byte() or try_write_byte() 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_exception_received_t
 *     storage to record the caught exception
 */

typedef struct {
	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_exception_received_t;

/*
 * try_read_write_read_byte() and try_read_write_write_byte() are functions that
 * read or write memory as their first instruction.
 * Used to test memory access that may provoke an exception.
 *
 * If the memory operation completes successfully,
 * *out_exception_received is unchanged.
 *
 * If an exception is received during the memory operation,
 * *out_exception_received is populated with the exception's
 * contents by the exception handler itself.
 *
 * try_read_write_exception_handler() below checks if the exception PC
 * is equal to one of these functions. The first instruction here 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))
try_read_write_read_byte(
	try_read_write_exception_received_t * const out_exception_received,
	mach_vm_address_t addr)
{
#if __arm64__
	asm("\n ldrb w0, [x1]"
            "\n ret");
#elif __x86_64__
	asm("\n movb (%rsi), %al"
            "\n nop"  /* pad load to four bytes */
            "\n nop"
            "\n ret");
#else
#       error unknown architecture
#endif
}

static void __attribute__((naked))
try_read_write_write_byte(
	try_read_write_exception_received_t * const out_exception_received,
	mach_vm_address_t addr,
	uint8_t value)
{
#if __arm64__
	asm("\n strb w2, [x1]"
            "\n ret");
#elif __x86_64__
	asm("\n movb %dl, (%rsi)"
            "\n nop"  /* pad store to four bytes */
            "\n nop"
            "\n ret");
#else
#       error unknown architecture
#endif
}

/*
 * Given thread state for an exception at read_byte() or write_byte(),
 * return the pointer to the exception received info
 * that the exception handler should populate.
 */
static try_read_write_exception_received_t *
get_exception_received_info(thread_state_t in_state)
{
	/*
	 * Pointer to info is in the first parameter register
	 * for both read_byte() and write_byte().
	 */
#if __arm64__
	arm_thread_state64_t *arm_state = (arm_thread_state64_t *)in_state;
	return (try_read_write_exception_received_t *)arm_state->__x[0];
#elif __x86_64__
	x86_thread_state64_t *x86_state = (x86_thread_state64_t *)in_state;
	return (try_read_write_exception_received_t *)x86_state->__rdi;
#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(
	exception_type_t exception,
	mach_exception_data_t codes,
	uint64_t exception_pc,
	thread_state_t in_state,
	mach_msg_type_number_t in_state_count __unused)
{
	assert(exception == EXC_BAD_ACCESS);

	uint64_t read_byte_pc  = (uint64_t)ptrauth_strip(&try_read_write_read_byte, ptrauth_key_function_pointer);
	uint64_t write_byte_pc = (uint64_t)ptrauth_strip(&try_read_write_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;
	}

	try_read_write_exception_received_t *info = get_exception_received_info(in_state);
	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;
}

/*
 * Begin try_read_write exception handling on this thread.
 */
static void
begin_expected_exceptions(void)
{
	/* global state */
	static dispatch_once_t try_read_write_initializer;
	static mach_port_t try_read_write_exc_port;

	/* thread-local state */
	static __thread bool try_read_write_thread_is_initialized;

	dispatch_once(&try_read_write_initializer, ^{
		/*
		 * Create an exc_helper exception handler for all try_read_write threads.
		 * We don't use exc_helper's default EXCEPTION_STATE_IDENTITY
		 * because that flavor can't change thread state to recover
		 * when Developer Mode is off (such as Release builds).
		 */
		try_read_write_exc_port = create_exception_port_behavior64(EXC_MASK_BAD_ACCESS, EXCEPTION_STATE);
		repeat_exception_handler_behavior64(try_read_write_exc_port, try_read_write_exception_handler, EXCEPTION_STATE);
	});

	if (try_read_write_thread_is_initialized == false) {
		/* Install the exception handler on this thread. */
		set_thread_exception_port_behavior64(try_read_write_exc_port, EXC_MASK_BAD_ACCESS, EXCEPTION_STATE);
		try_read_write_thread_is_initialized = true;
	}
}

/*
 * End try_read_write exception handling on this thread and evaluate the result.
 * Returns true if the operation was successful. Sets *out_error = KERN_SUCCESS.
 * Returns false if there was an exception. Sets *out_error to the exception's error.
 */
static bool
end_expected_exceptions(
	try_read_write_exception_received_t info,
	mach_vm_address_t addr,
	kern_return_t * const out_error)
{
	/*
	 * exception_pc was verified inside the exception handler.
	 * exception_kr will be verified by the caller of try_read/write_byte.
	 * Verify exception_memory here.
	 */
	if (info.exception_kr != KERN_SUCCESS) {
		assert(info.exception_memory == addr);
	}

	*out_error = info.exception_kr;
	return info.exception_kr == KERN_SUCCESS;
}

extern bool
try_read_byte(
	mach_vm_address_t addr,
	uint8_t * const out_byte,
	kern_return_t * const out_error)
{
	try_read_write_exception_received_t info = { KERN_SUCCESS, 0, 0 };

	begin_expected_exceptions();
	*out_byte = try_read_write_read_byte(&info, addr);
	return end_expected_exceptions(info, addr, out_error);
}

extern bool
try_write_byte(
	mach_vm_address_t addr,
	uint8_t byte,
	kern_return_t * const out_error)
{
	try_read_write_exception_received_t info = { KERN_SUCCESS, 0, 0 };

	begin_expected_exceptions();
	try_read_write_write_byte(&info, addr, byte);
	return end_expected_exceptions(info, addr, out_error);
}