SHA256.h

#include <ntddk.h>

#define SHA256_BLOCK_SIZE 64
#define SHA256_DIGEST_SIZE 32

static const ULONG Sha256InitialHashValue[8] = {
    0x6a09e667,
    0xbb67ae85,
    0x3c6ef372,
    0xa54ff53a,
    0x510e527f,
    0x9b05688c,
    0x1f83d9ab,
    0x5be0cd19
};

static const ULONG Sha256K[64] = {
    0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5,
    0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5,
    0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3,
    0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174,
    0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc,
    0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da,
    0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7,
    0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967,
    0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13,
    0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85,
    0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3,
    0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070,
    0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5,
    0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3,
    0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208,
    0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2
};

typedef struct _SHA256_CONTEXT
{
    ULONG State[8];
    UCHAR Buffer[SHA256_BLOCK_SIZE];
    ULONGLONG BitLength;
    ULONG BufferLength;
} SHA256_CONTEXT;

static __forceinline ULONG Rotr32(ULONG x, int n) { return (x >> n) | (x << (32 - n)); }

static __forceinline ULONG Ch(ULONG x, ULONG y, ULONG z) { return ((x & y) ^ ((~x) & z)); }

static __forceinline ULONG Maj(ULONG x, ULONG y, ULONG z) { return ((x & y) ^ (x & z) ^ (y & z)); }

static __forceinline ULONG Sigma0(ULONG x) { return (Rotr32(x, 2) ^ Rotr32(x, 13) ^ Rotr32(x, 22)); }

static __forceinline ULONG Sigma1(ULONG x) { return (Rotr32(x, 6) ^ Rotr32(x, 11) ^ Rotr32(x, 25)); }

static __forceinline ULONG sigma0(ULONG x) { return (Rotr32(x, 7) ^ Rotr32(x, 18) ^ (x >> 3)); }

static __forceinline ULONG sigma1(ULONG x) { return (Rotr32(x, 17) ^ Rotr32(x, 19) ^ (x >> 10)); }

static __forceinline 
VOID
Sha256Transform(
    _In_ SHA256_CONTEXT* Context,
    _In_ const UCHAR Block[SHA256_BLOCK_SIZE]
)
{
    ULONG W[64];
    ULONG a, b, c, d, e, f, g, h;
    ULONG t1, t2;
    int i;

    for (i = 0; i < 16; i++)
    {
        W[i] = ((ULONG)Block[i * 4] << 24) |
            ((ULONG)Block[i * 4 + 1] << 16) |
            ((ULONG)Block[i * 4 + 2] << 8) |
            ((ULONG)Block[i * 4 + 3]);
    }

    for (; i < 64; i++)
    {
        W[i] = sigma1(W[i - 2]) + W[i - 7] + sigma0(W[i - 15]) + W[i - 16];
    }

    a = Context->State[0];
    b = Context->State[1];
    c = Context->State[2];
    d = Context->State[3];
    e = Context->State[4];
    f = Context->State[5];
    g = Context->State[6];
    h = Context->State[7];

    for (i = 0; i < 64; i++)
    {
        t1 = h + Sigma1(e) + Ch(e, f, g) + Sha256K[i] + W[i];
        t2 = Sigma0(a) + Maj(a, b, c);
        h = g;
        g = f;
        f = e;
        e = d + t1;
        d = c;
        c = b;
        b = a;
        a = t1 + t2;
    }

    Context->State[0] += a;
    Context->State[1] += b;
    Context->State[2] += c;
    Context->State[3] += d;
    Context->State[4] += e;
    Context->State[5] += f;
    Context->State[6] += g;
    Context->State[7] += h;
}

static __forceinline
VOID
Sha256Init(
    _In_ SHA256_CONTEXT* Context
)
{
    RtlCopyMemory(Context->State, Sha256InitialHashValue, sizeof(Sha256InitialHashValue));
    RtlZeroMemory(Context->Buffer, sizeof(Context->Buffer));
    Context->BitLength = 0;
    Context->BufferLength = 0;
}

static __forceinline
VOID
Sha256Update(
    _In_ SHA256_CONTEXT* Context,
    _In_ const UCHAR* Data,
    _In_ SIZE_T Length
)
{
    SIZE_T i;
    for (i = 0; i < Length; i++)
    {
        Context->Buffer[Context->BufferLength] = Data[i];
        Context->BufferLength++;
        if (Context->BufferLength == SHA256_BLOCK_SIZE)
        {
            Sha256Transform(Context, Context->Buffer);
            Context->BitLength += (SHA256_BLOCK_SIZE * 8);
            Context->BufferLength = 0;
        }
    }
}

static __forceinline
VOID
Sha256Final(
    _In_ SHA256_CONTEXT* Context,
    _Out_ UCHAR Hash[SHA256_DIGEST_SIZE]
)
{
    ULONG i;
    ULONG padLen;
    UCHAR pad[SHA256_BLOCK_SIZE];

    Context->BitLength += (Context->BufferLength * 8);

    // 0x80 패딩 추가
    pad[0] = 0x80;
    RtlZeroMemory(&pad[1], SHA256_BLOCK_SIZE - 1);

    padLen = (Context->BufferLength < 56) ? (56 - Context->BufferLength) : (64 - Context->BufferLength + 56);

    Sha256Update(Context, pad, padLen);

    UCHAR lengthBytes[8];
    ULONGLONG bitLengthBE = _byteswap_uint64(Context->BitLength);
    RtlCopyMemory(lengthBytes, &bitLengthBE, 8);
    Sha256Update(Context, lengthBytes, 8);

    for (i = 0; i < 8; i++)
    {
        Hash[i * 4] = (UCHAR)(Context->State[i] >> 24);
        Hash[i * 4 + 1] = (UCHAR)(Context->State[i] >> 16);
        Hash[i * 4 + 2] = (UCHAR)(Context->State[i] >> 8);
        Hash[i * 4 + 3] = (UCHAR)(Context->State[i]);
    }
}

/**
 * @brief Compute a hash of the input data (SHA-256).
 *
 * @param Data Input data buffer.
 * @param DataSize Size of input data.
 * @param Hash Output buffer for 32-byte hash.
 * @return NTSTATUS STATUS_SUCCESS on success.
 */
static __forceinline 
NTSTATUS
ComputeHash(
    _In_reads_bytes_(DataSize) PVOID Data, 
    _In_ SIZE_T DataSize,
    _Out_writes_(32) PUCHAR Hash
)
{
    if (Data == NULL || Hash == NULL)
    {
        return STATUS_INVALID_PARAMETER;
    }

    SHA256_CONTEXT ctx;
    Sha256Init(&ctx);
    Sha256Update(&ctx, (UCHAR*)Data, DataSize);
    Sha256Final(&ctx, Hash);

    return STATUS_SUCCESS;
}