// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
 
using System;
using System.Runtime.CompilerServices;
using System.Runtime.ConstrainedExecution;
using System.Threading;
using Microsoft.AspNetCore.Cryptography.SafeHandles;
 
namespace Microsoft.AspNetCore.Cryptography;
 
internal static unsafe class UnsafeBufferUtil
{
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
#if NETSTANDARD2_0
    [ReliabilityContract(Consistency.WillNotCorruptState, Cer.Success)]
#endif
    public static void BlockCopy(void* from, void* to, int byteCount)
    {
        BlockCopy(from, to, checked((uint)byteCount)); // will be checked before invoking the delegate
    }
 
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
#if NETSTANDARD2_0
    [ReliabilityContract(Consistency.WillNotCorruptState, Cer.Success)]
#endif
    public static void BlockCopy(void* from, void* to, uint byteCount)
    {
        if (byteCount != 0)
        {
            BlockCopyCore((byte*)from, (byte*)to, byteCount);
        }
    }
 
#if NETSTANDARD2_0
    [ReliabilityContract(Consistency.WillNotCorruptState, Cer.MayFail)]
#endif
    public static void BlockCopy(LocalAllocHandle from, void* to, uint byteCount)
    {
        bool refAdded = false;
        try
        {
            from.DangerousAddRef(ref refAdded);
            BlockCopy((void*)from.DangerousGetHandle(), to, byteCount);
        }
        finally
        {
            if (refAdded)
            {
                from.DangerousRelease();
            }
        }
    }
 
#if NETSTANDARD2_0
    [ReliabilityContract(Consistency.WillNotCorruptState, Cer.MayFail)]
#endif
    public static void BlockCopy(void* from, LocalAllocHandle to, uint byteCount)
    {
        bool refAdded = false;
        try
        {
            to.DangerousAddRef(ref refAdded);
            BlockCopy(from, (void*)to.DangerousGetHandle(), byteCount);
        }
        finally
        {
            if (refAdded)
            {
                to.DangerousRelease();
            }
        }
    }
 
#if NETSTANDARD2_0
    [ReliabilityContract(Consistency.WillNotCorruptState, Cer.MayFail)]
#endif
    public static void BlockCopy(LocalAllocHandle from, LocalAllocHandle to, IntPtr length)
    {
        if (length == IntPtr.Zero)
        {
            return;
        }
 
        bool fromRefAdded = false;
        bool toRefAdded = false;
        try
        {
            from.DangerousAddRef(ref fromRefAdded);
            to.DangerousAddRef(ref toRefAdded);
            if (sizeof(IntPtr) == 4)
            {
                BlockCopyCore(from: (byte*)from.DangerousGetHandle(), to: (byte*)to.DangerousGetHandle(), byteCount: (uint)length.ToInt32());
            }
            else
            {
                BlockCopyCore(from: (byte*)from.DangerousGetHandle(), to: (byte*)to.DangerousGetHandle(), byteCount: (ulong)length.ToInt64());
            }
        }
        finally
        {
            if (fromRefAdded)
            {
                from.DangerousRelease();
            }
            if (toRefAdded)
            {
                to.DangerousRelease();
            }
        }
    }
 
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private static void BlockCopyCore(byte* from, byte* to, uint byteCount)
    {
        Buffer.MemoryCopy(from, to, (ulong)byteCount, (ulong)byteCount);
    }
 
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private static void BlockCopyCore(byte* from, byte* to, ulong byteCount)
    {
        Buffer.MemoryCopy(from, to, byteCount, byteCount);
    }
 
    /// <summary>
    /// Securely clears a memory buffer.
    /// </summary>
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
#if NETSTANDARD2_0
    [ReliabilityContract(Consistency.WillNotCorruptState, Cer.Success)]
#endif
    public static void SecureZeroMemory(byte* buffer, int byteCount)
    {
        SecureZeroMemory(buffer, checked((uint)byteCount));
    }
 
    /// <summary>
    /// Securely clears a memory buffer.
    /// </summary>
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
#if NETSTANDARD2_0
    [ReliabilityContract(Consistency.WillNotCorruptState, Cer.Success)]
#endif
    public static void SecureZeroMemory(byte* buffer, uint byteCount)
    {
        if (byteCount != 0)
        {
            do
            {
                buffer[--byteCount] = 0;
            } while (byteCount != 0);
 
            // Volatile to make sure the zero-writes don't get optimized away
            Volatile.Write(ref *buffer, 0);
        }
    }
 
    /// <summary>
    /// Securely clears a memory buffer.
    /// </summary>
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
#if NETSTANDARD2_0
    [ReliabilityContract(Consistency.WillNotCorruptState, Cer.Success)]
#endif
    public static void SecureZeroMemory(byte* buffer, ulong byteCount)
    {
        if (byteCount != 0)
        {
            do
            {
                buffer[--byteCount] = 0;
            } while (byteCount != 0);
 
            // Volatile to make sure the zero-writes don't get optimized away
            Volatile.Write(ref *buffer, 0);
        }
    }
 
    /// <summary>
    /// Securely clears a memory buffer.
    /// </summary>
#if NETSTANDARD2_0
    [ReliabilityContract(Consistency.WillNotCorruptState, Cer.Success)]
#endif
    public static void SecureZeroMemory(byte* buffer, IntPtr length)
    {
        if (sizeof(IntPtr) == 4)
        {
            SecureZeroMemory(buffer, (uint)length.ToInt32());
        }
        else
        {
            SecureZeroMemory(buffer, (ulong)length.ToInt64());
        }
    }
}