// 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.Security.Cryptography;
using Microsoft.AspNetCore.Cryptography;
using Microsoft.AspNetCore.DataProtection.Managed;
 
namespace Microsoft.AspNetCore.DataProtection.SP800_108;
 
internal static class ManagedSP800_108_CTR_HMACSHA512
{
    public static void DeriveKeys(byte[] kdk, ArraySegment<byte> label, ArraySegment<byte> context, Func<byte[], HashAlgorithm> prfFactory, ArraySegment<byte> output)
    {
        // make copies so we can mutate these local vars
        var outputOffset = output.Offset;
        var outputCount = output.Count;
 
        using (var prf = prfFactory(kdk))
        {
            // See SP800-108, Sec. 5.1 for the format of the input to the PRF routine.
            var prfInput = new byte[checked(sizeof(uint) /* [i]_2 */ + label.Count + 1 /* 0x00 */ + context.Count + sizeof(uint) /* [K]_2 */)];
 
            // Copy [L]_2 to prfInput since it's stable over all iterations
            uint outputSizeInBits = (uint)checked((int)outputCount * 8);
            prfInput[prfInput.Length - 4] = (byte)(outputSizeInBits >> 24);
            prfInput[prfInput.Length - 3] = (byte)(outputSizeInBits >> 16);
            prfInput[prfInput.Length - 2] = (byte)(outputSizeInBits >> 8);
            prfInput[prfInput.Length - 1] = (byte)(outputSizeInBits);
 
            // Copy label and context to prfInput since they're stable over all iterations
            Buffer.BlockCopy(label.Array!, label.Offset, prfInput, sizeof(uint), label.Count);
            Buffer.BlockCopy(context.Array!, context.Offset, prfInput, sizeof(int) + label.Count + 1, context.Count);
 
            var prfOutputSizeInBytes = prf.GetDigestSizeInBytes();
            for (uint i = 1; outputCount > 0; i++)
            {
                // Copy [i]_2 to prfInput since it mutates with each iteration
                prfInput[0] = (byte)(i >> 24);
                prfInput[1] = (byte)(i >> 16);
                prfInput[2] = (byte)(i >> 8);
                prfInput[3] = (byte)(i);
 
                // Run the PRF and copy the results to the output buffer
                var prfOutput = prf.ComputeHash(prfInput);
                CryptoUtil.Assert(prfOutputSizeInBytes == prfOutput.Length, "prfOutputSizeInBytes == prfOutput.Length");
                var numBytesToCopyThisIteration = Math.Min(prfOutputSizeInBytes, outputCount);
                Buffer.BlockCopy(prfOutput, 0, output.Array!, outputOffset, numBytesToCopyThisIteration);
                Array.Clear(prfOutput, 0, prfOutput.Length); // contains key material, so delete it
 
                // adjust offsets
                outputOffset += numBytesToCopyThisIteration;
                outputCount -= numBytesToCopyThisIteration;
            }
        }
    }
 
    public static void DeriveKeysWithContextHeader(byte[] kdk, ArraySegment<byte> label, byte[] contextHeader, ArraySegment<byte> context, Func<byte[], HashAlgorithm> prfFactory, ArraySegment<byte> output)
    {
        var combinedContext = new byte[checked(contextHeader.Length + context.Count)];
        Buffer.BlockCopy(contextHeader, 0, combinedContext, 0, contextHeader.Length);
        Buffer.BlockCopy(context.Array!, context.Offset, combinedContext, contextHeader.Length, context.Count);
        DeriveKeys(kdk, label, new ArraySegment<byte>(combinedContext), prfFactory, output);
    }
}