/**
@file			hook_trampoline_common.h
@brief    Common trampoline declaration
@details  Copyright (c) 2024 Acronis International GmbH
@author   Denis Kopyrin ([email protected])
@since    $Id: $
*/

#pragma once

#include "stringify.h"

#ifdef CONFIG_RETPOLINE
#include <asm/nospec-branch.h>
#endif

// If CONFIG_RETPOLINE is on, this will be magic to mute indirect 'jmp'
#ifndef ANNOTATE_RETPOLINE_SAFE
#define ANNOTATE_RETPOLINE_SAFE ""
#endif

// If CONFIG_RETHUNK is on, this will be magic that expands in a suitable 'ret'
#ifndef ASM_RET
#define ASM_RET "ret\n\t"
#endif

// Syscall hook may either decide to call to the original function 'fn' or return value 'ret'.
// If 'fn' is not NULL, syscall hook trampoline will jump to this function.
// If 'fn' is NULL, syscall hook trampoline will return 'ret'.
typedef struct {
	long fn;
	long ret;
} hook_ret_t;

// This function is called when x86_64 enters syscall.
// All parameters are in registers, push them on stack.
// As no params were changed, just call our pre-handler.
// Pop back the arguments and either call original function
// with the restored arguments or override return value.
// This asm function is the following C snippet coded to use tail jmp.
#if 0
typedef hook_ret_t (*syscall_hook_generic_t)(long a, long b, long c, long d, long e, long f);
typedef long (*syscall_generic_t)(long a, long b, long c, long d, long e, long f);
#define HOOK_TRAMPOLINE(abi, tag)
long name(long a, long b, long c, long d, long e, long f)
{
  syscall_hook_generic_t fn = (syscall_hook_generic_t)(void*) fn;
	hook_ret_t ret = fn(a, b, c, d, e, f);
	if (ret.fn)
	{
		syscall_generic_t orig = (syscall_generic_t)(void*) ret.fn;
		return orig(a, b, c, d, e, f);
	}
	else
	{
		return ret.ret;
	}
}
#endif

// In SysV 'RDI, RSI, RDX, RCX, R8, R9' are used for passing arguments, in kernel 'RDI, RSI, RDX, R10, R8, R9'
// That's due to the fact that 'RCX' is used for syscall passing in kernel so 'RCX' is clobbered
#ifndef KERNEL_MOCK
#define HOOK_TRAMPOLINE_SIZE_DECL(name) ".size " STRINGIFY(name) ", .-" STRINGIFY(name) "\n\t"
#else
#define HOOK_TRAMPOLINE_SIZE_DECL(name)
#endif

// R10 is being pushed twice to keep the stack aligned
#define HOOK_TRAMPOLINE_ASM(name, fn) __asm__( \
	".align 8;"                             "\n\t" \
	".pushsection .text;"                   "\n\t" \
	STRINGIFY(name) ":"                     "\n\t" \
	"push %rbp"                             "\n\t" \
	"mov %rsp, %rbp"                        "\n\t" \
	"push %rdx"                             "\n\t" \
	"push %rdi"                             "\n\t" \
	"push %rsi"                             "\n\t" \
	"push %r10"                             "\n\t" \
	"push %r10"                             "\n\t" \
	"push %rcx"                             "\n\t" \
	"push %r8"                              "\n\t" \
	"push %r9"                              "\n\t" \
	"call " STRINGIFY(fn)                   "\n\t" \
	"test %rax, %rax"                       "\n\t" \
	"je ._ret_" STRINGIFY(name)             "\n\t" \
	"pop %r9"                               "\n\t" \
	"pop %r8"                               "\n\t" \
	"pop %rcx"                              "\n\t" \
	"pop %r10"                              "\n\t" \
	"pop %r10"                              "\n\t" \
	"pop %rsi"                              "\n\t" \
	"pop %rdi"                              "\n\t" \
	"pop %rdx"                              "\n\t" \
	"pop %rbp"                              "\n\t" \
	"" ANNOTATE_RETPOLINE_SAFE "jmpq *%rax" "\n\t" \
	"._ret_" STRINGIFY(name) ":"            "\n\t" \
	"mov %rdx, %rax"                        "\n\t" \
	"pop %r9"                               "\n\t" \
	"pop %r8"                               "\n\t" \
	"pop %rcx"                              "\n\t" \
	"pop %r10"                              "\n\t" \
	"pop %r10"                              "\n\t" \
	"pop %rsi"                              "\n\t" \
	"pop %rdi"                              "\n\t" \
	"pop %rdx"                              "\n\t" \
	"pop %rbp"                              "\n\t" \
	"" ASM_RET ""                           "\n\t" \
	".popsection;"                          "\n\t" \
	".type " STRINGIFY(name) ", @function;" "\n\t" \
	HOOK_TRAMPOLINE_SIZE_DECL(name)         "\n\t" \
);