// Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
//
// NVIDIA CORPORATION and its licensors retain all intellectual property
// and proprietary rights in and to this software, related documentation
// and any modifications thereto.  Any use, reproduction, disclosure or
// distribution of this software and related documentation without an express
// license agreement from NVIDIA CORPORATION is strictly prohibited.

#include "stdafx.h"
#include "NvcMitigationPolicy.h"
#include <processthreadsapi.h>


namespace nvc
{
    typedef enum _PROCESS_MITIGATION_POLICY
    {
        ProcessDEPPolicy,
        ProcessASLRPolicy,
        ProcessDynamicCodePolicy,
        ProcessStrictHandleCheckPolicy,
        ProcessSystemCallDisablePolicy,
        ProcessMitigationOptionsMask,
        ProcessExtensionPointDisablePolicy,
        ProcessControlFlowGuardPolicy,
        ProcessSignaturePolicy,
        ProcessFontDisablePolicy,
        ProcessImageLoadPolicy,
        ProcessSystemCallFilterPolicy,
        ProcessPayloadRestrictionPolicy,
        ProcessChildProcessPolicy,
        ProcessSideChannelIsolationPolicy,
        ProcessUserShadowStackPolicy,
        ProcessRedirectionTrustPolicy,
        ProcessUserPointerAuthPolicy,
        ProcessSEHOPPolicy,
        MaxProcessMitigationPolicy
    } PROCESS_MITIGATION_POLICY,
        *PPROCESS_MITIGATION_POLICY;

    using PSetProcessMitigationPolicy = BOOL(WINAPI*)(PROCESS_MITIGATION_POLICY MitigationPolicy, PVOID lpBuffer, SIZE_T dwLength);

#pragma warning(push)
#pragma warning(disable : 4201) // nonstandard extension used: nameless struct/union
    typedef struct _PROCESS_MITIGATION_REDIRECTION_TRUST_POLICY
    {
        union
        {
            DWORD Flags;
            struct
            {
                DWORD EnforceRedirectionTrust : 1;
                DWORD AuditRedirectionTrust : 1;
                DWORD ReservedFlags : 30;
            } DUMMYSTRUCTNAME;
        } DUMMYUNIONNAME;
    } PROCESS_MITIGATION_REDIRECTION_TRUST_POLICY, *PPROCESS_MITIGATION_REDIRECTION_TRUST_POLICY;
#pragma warning(pop)

    bool EnableEnforceRedirectionTrust()
    {
        HMODULE kernel32Module = GetModuleHandleA("kernel32");

        if(kernel32Module == NULL)
        {
            return false;
        }
        PSetProcessMitigationPolicy setProcessMitigationPolicy =
            reinterpret_cast<PSetProcessMitigationPolicy>(GetProcAddress(kernel32Module, "SetProcessMitigationPolicy"));

        if(setProcessMitigationPolicy == nullptr)
        {
            return false;
        }

        PROCESS_MITIGATION_REDIRECTION_TRUST_POLICY policy = {};
        policy.EnforceRedirectionTrust = 1;

        return TRUE == setProcessMitigationPolicy(PROCESS_MITIGATION_POLICY::ProcessRedirectionTrustPolicy, &policy, sizeof(policy));
    }
} // namespace nvc