/*
 * @author  xamidev <xamidev@riseup.net>
 * @brief   Undefined behavior sanitization hooks
 * @license GPL-3.0-only
 */

#include <stdint.h>
#include <security/ubsan.h>
#include <kernel.h>
#include <stddef.h>

/*
 * Hooks were inspired from the Sortix implementation:
 * https://gitlab.com/sortix/sortix/-/blob/main/libc/ubsan/ubsan.c
 *
 * Not all hooks are implemented here.
 */

extern struct init_status init;

/*
 * __ubsan_handle_type_mismatch_v1 - Hook for UBSan
 * @data_raw:    pointer to error data
 * @pointer_raw: faulty pointer
 *
 * This function is executed when the UBSan library detects
 * following undefined behavior: type mismatch, null pointer
 * access, and unaligned access. It halts the system and
 * gives the location of the code that triggered it.
 */
void __ubsan_handle_type_mismatch_v1(void* data_raw, void* pointer_raw)
{
	struct ubsan_type_mismatch_v1_data* data = (struct ubsan_type_mismatch_v1_data*) data_raw;
	uintptr_t pointer = (uintptr_t)pointer_raw;
	uintptr_t alignment = (uintptr_t)1UL << data->log_alignment;
	const char* violation = "type mismatch";
	if ( !pointer ) {
        violation = "null pointer access";
    }
	else if ( alignment && (pointer & (alignment - 1)) ) {
        violation = "unaligned access";
    }
	
    struct ubsan_source_location* loc = &data->location;

    DEBUG("\x1b[38;5;231mUBSan: %s (ptr=%p) at %s:%u:%u\x1b[0m", violation, pointer, loc->filename, loc->line, loc->column);

    if (init.all) {
        printf("\x1b[38;5;231mUBSan: %s (ptr=%p) at %s:%u:%u\x1b[0m\r\n", violation, pointer, loc->filename, loc->line, loc->column);
    }

    panic(NULL, "Undefined Behavior Sanitization error");
}

/*
 * __ubsan_handle_pointer_overflow - Hook for UBSan
 * @data_raw:   pointer to error data
 * @base_raw:   base pointer
 * @result_raw: pointer after faulty operation
 *
 * This function is executed when the UBSan library detects
 * following undefined behavior: pointer overflow.
 * It halts the system and gives the location of the code
 * that triggered it.
 */
void __ubsan_handle_pointer_overflow(void* data_raw, void* base_raw, void* result_raw)
{
	struct ubsan_pointer_overflow_data* data = (struct ubsan_pointer_overflow_data*) data_raw;
	uintptr_t base = (uintptr_t)base_raw;
	uintptr_t result = (uintptr_t)result_raw;
    
    struct ubsan_source_location* loc = &data->location;

    DEBUG("\x1b[38;5;231mUBSan: pointer overflow (base=%p, result=%p) at %s:%u:%u\x1b[0m", base, result, loc->filename, loc->line, loc->column);

    if (init.all) {
        printf("\x1b[38;5;231mUBSan: pointer overflow (base=%p, result=%p) at %s:%u:%u\x1b[0m", base, result, loc->filename, loc->line, loc->column);
    }

    panic(NULL, "Undefined Behavior Sanitization error: pointer overflow");
}

/*
 * __ubsan_handle_shift_out_of_bounds - Hook for UBSan
 * @data_raw: pointer to error data
 * @lhs_raw:  left hand side (value being shifted)
 * @rhs_raw:  right hand side (shift amount)
 *
 * This function is executed when the UBSan library detects
 * following undefined behavior: shift out of bounds.
 * It halts the system and gives the location of the code
 * that triggered it.
 */
void __ubsan_handle_shift_out_of_bounds(void* data_raw, void* lhs_raw, void* rhs_raw)
{
	struct ubsan_shift_out_of_bounds_data* data = (struct ubsan_shift_out_of_bounds_data*) data_raw;
	uintptr_t lhs = (uintptr_t) lhs_raw;
	uintptr_t rhs = (uintptr_t) rhs_raw;
	
    struct ubsan_source_location* loc = &data->location;

    DEBUG("\x1b[38;5;231mUBSan: shift out of bounds (lhs=%p, rhs=%p) at %s:%u:%u\x1b[0m", lhs, rhs, loc->filename, loc->line, loc->column);

    if (init.all) {
        printf("\x1b[38;5;231mUBSan: shift out of bounds (lhs=%p, rhs=%p) at %s:%u:%u\x1b[0m", lhs, rhs, loc->filename, loc->line, loc->column);
    }

    panic(NULL, "Undefined Behavior Sanitization error: shift out of bounds");
}

/*
 * __ubsan_handle_load_invalid_value - Hook for UBSan
 * @data_raw:    pointer to error data
 * @value_raw:   value loaded
 *
 * This function is executed when the UBSan library detects
 * following undefined behavior: invalid value load.
 * It halts the system and gives the location of the code
 * that triggered it.
 */
void __ubsan_handle_load_invalid_value(void* data_raw, void* value_raw)
{
	struct ubsan_invalid_value_data* data = (struct ubsan_invalid_value_data*) data_raw;
	uintptr_t value = (uintptr_t) value_raw;

    struct ubsan_source_location* loc = &data->location;

    DEBUG("\x1b[38;5;231mUBSan: invalid value load (value=%p) at %s:%u:%u\x1b[0m", value, loc->filename, loc->line, loc->column);

    if (init.all) {
        printf("\x1b[38;5;231mUBSan: invalid value load (value=%p) at %s:%u:%u\x1b[0m", value, loc->filename, loc->line, loc->column);
    }

    panic(NULL, "Undefined Behavior Sanitization error: invalid value load");
}

/*
 * __ubsan_handle_out_of_bounds - Hook for UBSan
 * @data_raw:    pointer to error data
 * @index_raw:   out-of-bounds index in array
 *
 * This function is executed when the UBSan library detects
 * following undefined behavior: access out of bounds.
 * It halts the system and gives the location of the code
 * that triggered it.
 */
void __ubsan_handle_out_of_bounds(void* data_raw, void* index_raw)
{
	struct ubsan_out_of_bounds_data* data = (struct ubsan_out_of_bounds_data*) data_raw;
	uintptr_t index = (uintptr_t) index_raw;
    struct ubsan_source_location* loc = &data->location;

    DEBUG("\x1b[38;5;231mUBSan: out of bounds (index=%p) at %s:%u:%u\x1b[0m", index, loc->filename, loc->line, loc->column);

    if (init.all) {
        printf("\x1b[38;5;231mUBSan: out of bounds (index=%p) at %s:%u:%u\x1b[0m", index, loc->filename, loc->line, loc->column);
    }

    panic(NULL, "Undefined Behavior Sanitization error: out of bounds");
}

/*
 * __ubsan_handle_divrem_overflow - Hook for UBSan
 * @data_raw:    pointer to error data
 * @lhs_raw:     left hand side operator
 * @rhs_raw:     right hand side operator
 *
 * This function is executed when the UBSan library detects
 * following undefined behavior: division remainder overflow.
 * It halts the system and gives the location of the code
 * that triggered it.
 */
void __ubsan_handle_divrem_overflow(void* data_raw, void* lhs_raw, void* rhs_raw)
{
	struct ubsan_overflow_data* data = (struct ubsan_overflow_data*) data_raw;
	uintptr_t lhs = (uintptr_t) lhs_raw;
	uintptr_t rhs = (uintptr_t) rhs_raw;
    struct ubsan_source_location* loc = &data->location;

	DEBUG("\x1b[38;5;231mUBSan: division remainder overflow (lhs=%p, rhs=%p) at %s:%u:%u\x1b[0m", lhs, rhs, loc->filename, loc->line, loc->column);

    if (init.all) {
        printf("\x1b[38;5;231mUBSan: division remainder overflow (lhs=%p, rhs=%p) at %s:%u:%u\x1b[0m", lhs, rhs, loc->filename, loc->line, loc->column);
    }

    panic(NULL, "Undefined Behavior Sanitization error: division remainder overflow");
}

/*
 * __ubsan_handle_sub_overflow - Hook for UBSan
 * @data_raw:    pointer to error data
 * @lhs_raw:     left hand side operator
 * @rhs_raw:     right hand side operator
 *
 * This function is executed when the UBSan library detects
 * following undefined behavior: subtraction overflow.
 * It halts the system and gives the location of the code
 * that triggered it.
 */
void __ubsan_handle_sub_overflow(void* data_raw, void* lhs_raw, void* rhs_raw)
{
	struct ubsan_overflow_data* data = (struct ubsan_overflow_data*) data_raw;
	uintptr_t lhs = (uintptr_t) lhs_raw;
	uintptr_t rhs = (uintptr_t) rhs_raw;
	struct ubsan_source_location* loc = &data->location;

	DEBUG("\x1b[38;5;231mUBSan: subtraction overflow (lhs=%p, rhs=%p) at %s:%u:%u\x1b[0m", lhs, rhs, loc->filename, loc->line, loc->column);

    if (init.all) {
        printf("\x1b[38;5;231mUBSan: subtraction overflow (lhs=%p, rhs=%p) at %s:%u:%u\x1b[0m", lhs, rhs, loc->filename, loc->line, loc->column);
    }

    panic(NULL, "Undefined Behavior Sanitization error: subtraction overflow");
}

/*
 * __ubsan_handle_add_overflow - Hook for UBSan
 * @data_raw:    pointer to error data
 * @lhs_raw:     left hand side operator
 * @rhs_raw:     right hand side operator
 *
 * This function is executed when the UBSan library detects
 * following undefined behavior: addition overflow.
 * It halts the system and gives the location of the code
 * that triggered it.
 */
void __ubsan_handle_add_overflow(void* data_raw, void* lhs_raw, void* rhs_raw)
{
	struct ubsan_overflow_data* data = (struct ubsan_overflow_data*) data_raw;
	uintptr_t lhs = (uintptr_t) lhs_raw;
	uintptr_t rhs = (uintptr_t) rhs_raw;
    struct ubsan_source_location* loc = &data->location;

	DEBUG("\x1b[38;5;231mUBSan: addition overflow (lhs=%p, rhs=%p) at %s:%u:%u\x1b[0m", lhs, rhs, loc->filename, loc->line, loc->column);

    if (init.all) {
        printf("\x1b[38;5;231mUBSan: addition overflow (lhs=%p, rhs=%p) at %s:%u:%u\x1b[0m", lhs, rhs, loc->filename, loc->line, loc->column);
    }

    panic(NULL, "Undefined Behavior Sanitization error: addition overflow");
}

/*
 * __ubsan_handle_mul_overflow - Hook for UBSan
 * @data_raw:    pointer to error data
 * @lhs_raw:     left hand side operator
 * @rhs_raw:     right hand side operator
 *
 * This function is executed when the UBSan library detects
 * following undefined behavior: multiplication overflow.
 * It halts the system and gives the location of the code
 * that triggered it.
 */
void __ubsan_handle_mul_overflow(void* data_raw, void* lhs_raw, void* rhs_raw)
{
	struct ubsan_overflow_data* data = (struct ubsan_overflow_data*) data_raw;
	uintptr_t lhs = (uintptr_t) lhs_raw;
	uintptr_t rhs = (uintptr_t) rhs_raw;
	struct ubsan_source_location* loc = &data->location;

	DEBUG("\x1b[38;5;231mUBSan: multiplication overflow (lhs=%p, rhs=%p) at %s:%u:%u\x1b[0m", lhs, rhs, loc->filename, loc->line, loc->column);

    if (init.all) {
        printf("\x1b[38;5;231mUBSan: multiplication overflow (lhs=%p, rhs=%p) at %s:%u:%u\x1b[0m", lhs, rhs, loc->filename, loc->line, loc->column);
    }

    panic(NULL, "Undefined Behavior Sanitization error: multiplication overflow");
}

/*
 * __ubsan_handle_negate_overflow - Hook for UBSan
 * @data_raw:      pointer to error data
 * @old_value_raw: value before overflow
 *
 * This function is executed when the UBSan library detects
 * following undefined behavior: negation overflow.
 * It halts the system and gives the location of the code
 * that triggered it.
 */
void __ubsan_handle_negate_overflow(void* data_raw, void* old_value_raw)
{
	struct ubsan_overflow_data* data = (struct ubsan_overflow_data*) data_raw;
	uintptr_t old_value = (uintptr_t) old_value_raw;
    struct ubsan_source_location* loc = &data->location;

	DEBUG("\x1b[38;5;231mUBSan: negation overflow (old_value=%p) at %s:%u:%u\x1b[0m", old_value, loc->filename, loc->line, loc->column);

    if (init.all) {
        printf("\x1b[38;5;231mUBSan: negation overflow (old_value=%p) at %s:%u:%u\x1b[0m", old_value, loc->filename, loc->line, loc->column);
    }

    panic(NULL, "Undefined Behavior Sanitization error: negation overflow");
}