/* _NVRM_COPYRIGHT_BEGIN_
 *
 * Copyright 2020 by NVIDIA Corporation.  All rights reserved.  All
 * information contained herein is proprietary and confidential to NVIDIA
 * Corporation.  Any use, reproduction, or disclosure without the written
 * permission of NVIDIA Corporation is prohibited.
 *
 * _NVRM_COPYRIGHT_END_
 */

#include <stdint.h>
#include <stdbool.h>
#include <misc/gcc_attrs.h>
#include <cpu/shutdown.h>
#include <ssp/ssp.h>
#include <error/common-errors.h>
#if !NVRISCV_FEATURE_SSP_FORCE_SW_CANARY
#include <misc/nvmisc_drf.h>
#include <scp/scp.h>
#endif // NVRISCV_FEATURE_SSP_FORCE_SW_CANARY

START_RFD_BLOCK(MISRA, DEVIATE, Rule_1_2, "Approval: Bug 200531996, DR: SWE-FSP-009-SWSADR.docx",
                MISRA, DEVIATE, Directive_4_9, "Approval: Bug 200531995, DR: SWE-FSP-012-SWSADR.docx",
                MISRA, DEVIATE, Rule_15_1, "Approval: Bug 3176459, DR: SWE-FSP-028-SWSADR.docx")

#define CHECK_ERROR_OR_GOTO(X, Y)   \
({                                  \
    ret = X;                        \
    if (ret != E_SUCCESS)           \
    {                               \
       goto Y;                      \
    }                               \
})

/*
 * For baremetal image, it can stay as is.
 * For multi-partition images, it likely nees to be written in some shared data
 * segment and then stored/restored on partition entry/exit
 */
GCC_ATTR_SECTION(".data.libnvriscv.__stack_chk_guard")
uintptr_t __stack_chk_guard  = 0;

#if !NVRISCV_FEATURE_SSP_FORCE_SW_CANARY
GCC_ATTR_SECTION(".data.libnvriscv.ssp_scpdma_buf")
GCC_ATTR_ALIGNED(SCP_RAND_SIZE) static uintptr_t ssp_scpdma_buf[SCP_RAND_SIZE / sizeof(uint64_t)] = {0,};
#endif

GCC_ATTR_NAKED
GCC_ATTR_NORETURN
GCC_ATTR_NO_SSP
void __stack_chk_fail(void)
{
#if NVRISCV_FEATURE_SSP_ENABLE_FAIL_HOOK
    __asm__("jal ssp_check_fail_hook");
#endif // NVRISCV_FEATURE_SSP_ENABLE_FAIL_HOOK
    __asm__("j riscv_panic");
}

GCC_ATTR_NO_SSP
void ssp_set_canary(uintptr_t canary)
{
    __stack_chk_guard = canary;
}

uintptr_t ssp_get_canary(void)
{
    return __stack_chk_guard;
}

#if !NVRISCV_FEATURE_SSP_FORCE_SW_CANARY
GCC_ATTR_NO_SSP
error_t ssp_generate_and_set_canary_with_init(void)
{
    error_t ret = E_SUCCESS;

    CHECK_ERROR_OR_GOTO((scp_Init(SCP_INIT_FLAG_DEFAULT)), out);

    CHECK_ERROR_OR_GOTO((scp_Configure_Rand(SCP_RAND_CONFIG_DEFAULT)), out);

    CHECK_ERROR_OR_GOTO((scp_Start_Rand()), out);

    CHECK_ERROR_OR_GOTO((ssp_generate_and_set_canary()), out);

    ret = scp_Shutdown();

out:
    return ret;

}

GCC_ATTR_NO_SSP
error_t ssp_generate_and_set_canary(void)
{
    error_t ret = E_SUCCESS;
    uint32_t attempts=10; // Try 10 times to get proper number

    ssp_scpdma_buf[0] = 0;

    while (attempts--)
    {
        CHECK_ERROR_OR_GOTO((scp_Get_Rand((uintptr_t)&ssp_scpdma_buf, sizeof(ssp_scpdma_buf), 0)), out);

        if (ssp_scpdma_buf[0] != 0)
        {
            ssp_set_canary(ssp_scpdma_buf[0]);
            goto out;
        }
    }

    ret = E_TIMEOUT;

out:
    // Wipe buffer to avoid canary leak
    ssp_scpdma_buf[0] = 0;
    return ret;
}
#endif // NVRISCV_FEATURE_SSP_FORCE_SW_CANARY

END_RFD_BLOCK(MISRA, DEVIATE, Rule_1_2, "Approval: Bug 200531996, DR: SWE-FSP-009-SWSADR.docx",
              MISRA, DEVIATE, Directive_4_9, "Approval: Bug 200531995, DR: SWE-FSP-012-SWSADR.docx",
              MISRA, DEVIATE, Rule_15_1, "Approval: Bug 3176459, DR: SWE-FSP-028-SWSADR.docx")
