/*
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
 *
 * Copyright 2018 Danny Robson <danny@nerdcruft.net>
 */

#include "blake2.hpp"

#include <cruft/util/bitwise.hpp>
#include <cruft/util/debug.hpp>
#include <cruft/util/endian.hpp>
#include <cruft/util/types.hpp>

using cruft::crypto::hash::blake2;


///////////////////////////////////////////////////////////////////////////////
// blake2b: uint64_t
struct traits {
    static constexpr int word_bits = 64;
    using word_t = typename cruft::bits_type<word_bits>::uint;
    static constexpr int F_rounds = 12;
    static constexpr int block_bytes = 128; // bb
    static constexpr int max_hash_bytes = 64; // nn
    static constexpr int max_key_bytes = 64; // kk
    //static constexpr int max_input_bytes = 2**128; // ll
    static constexpr int rotations[4] = { 32, 24, 16, 63 };

    static constexpr
    std::array<word_t,8> iv {
        0x6A09E667F3BCC908, 0xBB67AE8584CAA73B,
        0x3C6EF372FE94F82B, 0xA54FF53A5F1D36F1,
        0x510E527FADE682D1, 0x9B05688C2B3E6C1F,
        0x1F83D9ABFB41BD6B, 0x5BE0CD19137E2179
    };
};


using word_t = traits::word_t;

static constexpr
int
    sigma[12][16] {

    {  0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, },
    { 14, 10,  4,  8,  9, 15, 13,  6,  1, 12,  0,  2, 11,  7,  5,  3, },
    { 11,  8, 12,  0,  5,  2, 15, 13, 10, 14,  3,  6,  7,  1,  9,  4, },
    {  7,  9,  3,  1, 13, 12, 11, 14,  2,  6,  5, 10,  4,  0, 15,  8, },
    {  9,  0,  5,  7,  2,  4, 10, 15, 14,  1, 11, 12,  6,  8,  3, 13, },
    {  2, 12,  6, 10,  0, 11,  8,  3,  4, 13,  7,  5, 15, 14,  1,  9, },
    { 12,  5,  1, 15, 14, 13,  4, 10,  0,  7,  6,  3,  9,  2,  8, 11, },
    { 13, 11,  7, 14, 12,  1,  3,  9,  5,  0, 15,  4,  8,  6,  2, 10, },
    {  6, 15, 14,  9, 11,  3,  0,  8, 12,  2, 13,  7,  1,  4, 10,  5, },
    { 10,  2,  8,  4,  7,  6,  1,  5, 15, 11,  9, 14,  3, 12, 13,  0, },
    {  0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, },
    { 14, 10,  4,  8,  9, 15, 13,  6,  1, 12,  0,  2, 11,  7,  5,  3, },
};


///////////////////////////////////////////////////////////////////////////////
// mixing function
static constexpr
void
G (std::array<word_t,16> &v, int a, int b, int c, int d, word_t x, word_t y)
{
    v[a] = v[a] + v[b] + x;
    v[d] = cruft::rotater (v[d] ^ v[a], traits::rotations[0]);

    v[c] = v[c] + v[d];
    v[b] = cruft::rotater (v[b] ^ v[c], traits::rotations[1]);

    v[a] = v[a] + v[b] + y;
    v[d] = cruft::rotater (v[d] ^ v[a], traits::rotations[2]);

    v[c] = v[c] + v[d];
    v[b] = cruft::rotater (v[b] ^ v[c], traits::rotations[3]);
}


//-----------------------------------------------------------------------------
// compression function
std::array<word_t,8>
F (std::array<word_t,8> h, const word_t m[16], uint64_t t, bool f)
{
    std::array<word_t,16> v {
        h[0], h[1], h[2], h[3],
        h[4], h[5], h[6], h[7],

        traits::iv[0], traits::iv[1],
        traits::iv[2], traits::iv[3],
        traits::iv[4], traits::iv[5],
        traits::iv[6], traits::iv[7],
    };

    v[12] ^= t;
    v[13] ^= 0;

    if (f)
        v[14] ^= ~0;

    for (int i = 0; i < traits::F_rounds; ++i) {
        const auto *s = sigma[i];

        G (v, 0, 4,  8, 12, m[s[ 0]], m[s[ 1]]);
        G (v, 1, 5,  9, 13, m[s[ 2]], m[s[ 3]]);
        G (v, 2, 6, 10, 14, m[s[ 4]], m[s[ 5]]);
        G (v, 3, 7, 11, 15, m[s[ 6]], m[s[ 7]]);

        G (v, 0, 5, 10, 15, m[s[ 8]], m[s[ 9]]);
        G (v, 1, 6, 11, 12, m[s[10]], m[s[11]]);
        G (v, 2, 7,  8, 13, m[s[12]], m[s[13]]);
        G (v, 3, 4,  9, 14, m[s[14]], m[s[15]]);
    }

    for (int i = 0; i < 8; ++i)
        h[i] ^= v[i] ^ v[i + 8];

    return h;
}


///////////////////////////////////////////////////////////////////////////////
blake2::blake2 () noexcept:
    blake2 (cruft::view<const uint8_t*>{nullptr})
{ ; }


//-----------------------------------------------------------------------------
blake2::blake2 (cruft::view<const uint8_t *> key)
{
    // don't give the user flexibility to provide too much key
    if (key.size () > ::traits::max_key_bytes)
        throw std::invalid_argument ("key is too large");

    std::fill (m_salt.begin (), m_salt.end (), 0);
    memcpy (m_salt.data (), key.data (), key.size ());
    m_keylen = key.size ();
}


//-----------------------------------------------------------------------------
blake2::digest_t
blake2::operator() (cruft::view<const uint8_t *> data) const noexcept
{
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wcast-align"
    auto h = ::traits::iv;
    h[0] ^= 0x01010000 ^ (m_keylen << 8) ^ sizeof (digest_t);

    if (m_keylen)
        h = F (h, reinterpret_cast<const word_t*> (m_salt.data ()), ::traits::block_bytes, data.empty ());

    // special case for the empty key and empty data
    if (!m_keylen && data.empty ()) {
        std::array<word_t,16> zeroes {};
        h = F (h, zeroes.data (), 0, true);
    }

    uint64_t counter = m_keylen?::traits::block_bytes:0;

    auto cursor = data.begin ();
    while (cursor + ::traits::block_bytes < data.end ()) {
        counter += ::traits::block_bytes;
        h = F (h, reinterpret_cast<const word_t*> (cursor), counter, false);
        cursor += ::traits::block_bytes;
    }

    if (cursor != data.cend ()) {
        std::array<uint64_t,16> tail {};
        memcpy (tail.data(), data.data (), data.cend () - cursor);
        counter += data.end () - cursor;
        h = F (h, tail.data (), counter, true);
    }

    digest_t d;
    memcpy (&d, h.data (), sizeof (d));
    return d;
#pragma GCC diagnostic pop
}