sha2.hpp

/* SPDX-License-Identifier: CC0-1.0 */
#include <array>
#include <cstddef>
#include <cstdint>

#include <algorithm>
#include <span>
#include <bit>

namespace sha2
{
    using std::size_t;
    using std::array;
    using std::span;
    using std::byte;
    using std::rotr;

    template <
        size_t   DigestBits,
        size_t   BlockBits,
        typename WordType,
        typename BitSizeType,
        const auto& RoundConstants,
        const array<WordType, 8> InitialState
    >
    class sha2
    {
    public:
        static constexpr size_t digest_bits   = DigestBits;
        static constexpr size_t digest_bytes  = DigestBits/8;
        static constexpr size_t block_bits    = BlockBits;
        static constexpr size_t block_bytes   = BlockBits/8;
        static constexpr size_t rounds        = RoundConstants.size();
        static constexpr auto initial_hash    = InitialState;
        static constexpr auto round_constants = RoundConstants;

        using word_t    = WordType;
        using bitsize_t = BitSizeType;
        using hash_t    = array<word_t, 8>;
        using digest_t  = array<byte, digest_bytes>;
        constexpr sha2(void) = default;

        template<std::ranges::contiguous_range R>
        constexpr sha2(R&& data)
            requires std::same_as<std::ranges::range_value_t<R>, byte>
        {
            update(data);
            finish();
        }

        template<std::ranges::contiguous_range R>
        void update(R&& data)
            requires std::same_as<std::ranges::range_value_t<R>, byte>
        {
            auto bytes = std::as_bytes(span{data});
            size_t rem = bytes.size();
            size_t pos = 0;

            while (rem > 0)
            {
                size_t free   = blockbuf.size() - blockpos;
                size_t tocopy = std::min(free, rem);
                std::copy_n(bytes.data()+pos, tocopy, blockbuf.begin()+blockpos);
                blockpos += tocopy;

                if (blockpos == blockbuf.size())
                {
                    compress_block(blockbuf);
                    blockpos = 0;
                }

                pos += tocopy;
                rem -= tocopy;
            }
            total_bits += bytes.size()*8;
        }

        digest_t finish(void)
        {
            /* always in-range, since a full block update() would have reset blockpos to 0 */
            blockbuf[blockpos++] = byte(0x80);

            /* zero-pad/compress rest of block if not enough room for bit size */
            size_t bits_pos = block_bytes - sizeof(bitsize_t);
            if (blockpos >= bits_pos)
            {
                std::fill(blockbuf.begin() + blockpos, blockbuf.end(), byte(0));
                compress_block(blockbuf);
                blockpos = 0;
            }

            /* write bit count 'l' */
            auto bitbytes    = std::as_bytes(span{&total_bits, 1});
            auto bitdest  = blockbuf.begin() + bits_pos;

            /* zero pad up to bitsize position */
            std::fill(blockbuf.begin() + blockpos, bitdest, byte(0));
            /* fill bitsize (which is stored native-endian so may need swapping */
            be_copy(bitbytes.begin(), bitbytes.end(), bitdest);

            compress_block(blockbuf);

            return digest();
        }

        digest_t digest(void) const
        {
            digest_t digest;
            auto it = digest.begin();
            for (size_t word_i = 0; word_i < digest.size() / sizeof(word_t); word_i++)
            {
                /* copy native-endian state into digest as big-endian word */
                auto bytes = std::as_bytes(span{&state[word_i], 1});
                be_copy(bytes.begin(), bytes.end(), it);
                it += sizeof(word_t);
            }
            return digest;
        }

    private:
        hash_t state = initial_hash;
        bitsize_t total_bits = 0;
        size_t blockpos = 0;
        alignas(word_t) array<byte, block_bytes> blockbuf;

        void compress_block(span<const byte, block_bytes> block)
        {
            array<word_t, 16> W; /* sliding window of W to care about */

            hash_t work = state;
            for (size_t t = 0; t < rounds; t++)
            {
                constexpr size_t mask = W.size() - 1; /* limit indices into W */
                size_t Widx = t & mask;
                if (t < 16)
                {
                    /* convert big-endian block word to native-endian word */
                    auto target = std::as_writable_bytes(span{&W[Widx], 1});
                    auto start  = block.begin() + t*sizeof(word_t);
                    auto end    = start + sizeof(word_t);
                    be_copy(start, end, target.begin());
                }
                else
                {
                    word_t Wt_2  = W[(t-2)  & mask];
                    word_t Wt_7  = W[(t-7)  & mask];
                    word_t Wt_15 = W[(t-15) & mask];
                    word_t Wt_16 = W[(t-16) & mask];
                    W[Widx] = sigma1(Wt_2) + Wt_7 + sigma0(Wt_15) + Wt_16;
                }
                word_t Wt = W[Widx];
                word_t Kt = round_constants[t];
                auto& [a,b,c,d,e,f,g,h] = work;

                word_t T1 = h + big_sigma1(e) + ch(e, f, g) + Kt + Wt;
                word_t T2 = big_sigma0(a) + maj(a, b, c);
                h = g;
                g = f;
                f = e;
                e = d + T1;
                d = c;
                c = b;
                b = a;
                a = T1 + T2;
            }

            for (size_t i = 0; i < state.size(); i++)
            {
                state[i] += work[i];
            }
        }

        static word_t ch(word_t x, word_t y, word_t z)
        {
            return (x & y) ^ (~x & z);                  // FIPS.180-4: (Equation 4.2, Equation 4.8)
        }

        static word_t maj(word_t x, word_t y, word_t z)
        {
          return (x & y) ^ (x & z) ^ (y & z);           // FIPS.180-4: (Equation 4.3, Equation 4.9)
        }

        static word_t big_sigma0(word_t x)
        {
            return std::is_same_v<word_t, uint32_t>
                ? rotr(x, 2) ^ rotr(x,13) ^ rotr(x,22)  // FIPS.180-4: (Equation 4.4)
                : rotr(x,28) ^ rotr(x,34) ^ rotr(x,39); // FIPS.180-4: (Equation 4.10)
        }

        static word_t big_sigma1(word_t x)
        {
            return std::is_same_v<word_t, uint32_t>
                ? rotr(x, 6) ^ rotr(x,11) ^ rotr(x,25)  // FIPS.180-4: (Equation 4.5)
                : rotr(x,14) ^ rotr(x,18) ^ rotr(x,41); // FIPS.180-4: (Equation 4.11)
        }

        static word_t sigma0(word_t x)
        {
            return std::is_same_v<word_t, uint32_t>
                ? rotr(x,7) ^ rotr(x,18) ^ (x >> 3)     // FIPS.180-4: (Equation 4.6)
                : rotr(x,1) ^ rotr(x, 8) ^ (x >> 7);    // FIPS.180-4: (Equation 4.12)
        }

        static word_t sigma1(word_t x)
        {
            return std::is_same_v<word_t, uint32_t>
                ? rotr(x,17) ^ rotr(x,19) ^ (x >> 10)   // FIPS.180-4: (Equation 4.7)
                : rotr(x,19) ^ rotr(x,61) ^ (x >> 6);   // FIPS.180-4: (Equation 4.13)
        }

        template <typename InputIt, typename OutputIt>
        void be_copy(InputIt first, InputIt last, OutputIt dest) const
        {
            /* copy a range of bytes, reversing them on little-endian architectures */
            if constexpr (std::endian::big == std::endian::native)
            {
                std::copy(first, last, dest);
            }
            else
            {
                std::reverse_copy(first, last, dest);
            }
        }
    };

    static constexpr array<uint32_t, 8> sha256_H0 = {
        0x6A09E667, 0xBB67AE85, 0x3C6EF372, 0xA54FF53A,
        0x510E527F, 0x9B05688C, 0x1F83D9AB, 0x5BE0CD19,
    };

    static constexpr array<uint32_t, 8> sha224_H0 = {
        0xC1059ED8, 0x367CD507, 0x3070DD17, 0xF70E5939,
        0xFFC00B31, 0x68581511, 0x64F98FA7, 0xBEFA4FA4,
    };

    static constexpr array<uint32_t, 64> sha256_K = {
        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,
    };

    static constexpr array<uint64_t, 8> sha512t224_H0 = {
        0x8C3D37C819544DA2, 0x73E1996689DCD4D6, 0x1DFAB7AE32FF9C82, 0x679DD514582F9FCF,
        0x0F6D2B697BD44DA8, 0x77E36F7304C48942, 0x3F9D85A86A1D36C8, 0x1112E6AD91D692A1,
    };

    static constexpr array<uint64_t, 8> sha512t256_H0 = {
        0x22312194FC2BF72C, 0x9F555FA3C84C64C2, 0x2393B86B6F53B151, 0x963877195940EABD,
        0x96283EE2A88EFFE3, 0xBE5E1E2553863992, 0x2B0199FC2C85B8AA, 0x0EB72DDC81C52CA2,
    };

    static constexpr array<uint64_t, 8> sha384_H0 = {
        0xCBBB9D5DC1059ED8, 0x629A292A367CD507, 0x9159015A3070DD17, 0x152FECD8F70E5939,
        0x67332667FFC00B31, 0x8EB44A8768581511, 0xDB0C2E0D64F98FA7, 0x47B5481DBEFA4FA4,
    };

    static constexpr array<uint64_t, 8> sha512_H0 = {
        0x6A09E667F3BCC908, 0xBB67AE8584CAA73B, 0x3C6EF372FE94F82B, 0xA54FF53A5F1D36F1,
        0x510E527FADE682D1, 0x9B05688C2B3E6C1F, 0x1F83D9ABFB41BD6B, 0x5BE0CD19137E2179,
    };

    static constexpr array<uint64_t, 80> sha512_K = {
        0x428A2F98D728AE22, 0x7137449123EF65CD, 0xB5C0FBCFEC4D3B2F, 0xE9B5DBA58189DBBC,
        0x3956C25BF348B538, 0x59F111F1B605D019, 0x923F82A4AF194F9B, 0xAB1C5ED5DA6D8118,
        0xD807AA98A3030242, 0x12835B0145706FBE, 0x243185BE4EE4B28C, 0x550C7DC3D5FFB4E2,
        0x72BE5D74F27B896F, 0x80DEB1FE3B1696B1, 0x9BDC06A725C71235, 0xC19BF174CF692694,
        0xE49B69C19EF14AD2, 0xEFBE4786384F25E3, 0x0FC19DC68B8CD5B5, 0x240CA1CC77AC9C65,
        0x2DE92C6F592B0275, 0x4A7484AA6EA6E483, 0x5CB0A9DCBD41FBD4, 0x76F988DA831153B5,
        0x983E5152EE66DFAB, 0xA831C66D2DB43210, 0xB00327C898FB213F, 0xBF597FC7BEEF0EE4,
        0xC6E00BF33DA88FC2, 0xD5A79147930AA725, 0x06CA6351E003826F, 0x142929670A0E6E70,
        0x27B70A8546D22FFC, 0x2E1B21385C26C926, 0x4D2C6DFC5AC42AED, 0x53380D139D95B3DF,
        0x650A73548BAF63DE, 0x766A0ABB3C77B2A8, 0x81C2C92E47EDAEE6, 0x92722C851482353B,
        0xA2BFE8A14CF10364, 0xA81A664BBC423001, 0xC24B8B70D0F89791, 0xC76C51A30654BE30,
        0xD192E819D6EF5218, 0xD69906245565A910, 0xF40E35855771202A, 0x106AA07032BBD1B8,
        0x19A4C116B8D2D0C8, 0x1E376C085141AB53, 0x2748774CDF8EEB99, 0x34B0BCB5E19B48A8,
        0x391C0CB3C5C95A63, 0x4ED8AA4AE3418ACB, 0x5B9CCA4F7763E373, 0x682E6FF3D6B2B8A3,
        0x748F82EE5DEFB2FC, 0x78A5636F43172F60, 0x84C87814A1F0AB72, 0x8CC702081A6439EC,
        0x90BEFFFA23631E28, 0xA4506CEBDE82BDE9, 0xBEF9A3F7B2C67915, 0xC67178F2E372532B,
        0xCA273ECEEA26619C, 0xD186B8C721C0C207, 0xEADA7DD6CDE0EB1E, 0xF57D4F7FEE6ED178,
        0x06F067AA72176FBA, 0x0A637DC5A2C898A6, 0x113F9804BEF90DAE, 0x1B710B35131C471B,
        0x28DB77F523047D84, 0x32CAAB7B40C72493, 0x3C9EBE0A15C9BEBC, 0x431D67C49C100D4C,
        0x4CC5D4BECB3E42B6, 0x597F299CFC657E2A, 0x5FCB6FAB3AD6FAEC, 0x6C44198C4A475817,
    };

    /* Hash functions based on SHA-256, using different initial hash and digest bits */
    template <const auto& InitialState, size_t TruncateBits = 256>
    using sha256x = sha2<
        TruncateBits,
        512,
        uint32_t,
        uint64_t,
        sha256_K,
        InitialState
    >;

    /* Hash functions based on SHA-512, using different initial hash and digest bits */
    template <const auto& InitialState, size_t TruncateBits = 512>
    using sha512x = sha2<
        TruncateBits,
        1024,
        uint64_t,
        __uint128_t,
        sha512_K,
        InitialState
    >;

    using sha224     = sha256x<sha224_H0, 224>;
    using sha256     = sha256x<sha256_H0>;
    using sha512t224 = sha512x<sha512t224_H0, 224>;
    using sha512t256 = sha512x<sha512t256_H0, 256>;
    using sha384     = sha512x<sha384_H0, 384>;
    using sha512     = sha512x<sha512_H0>;
};