// Copyright (c) 2024 blackshirt.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.
//
module chacha20

import math.bits
import encoding.binary

// max_64bit_counter is a 64-bit maximum internal counter of original ChaCha20 variant.
const max_64bit_counter = max_u64
// max_32bit_counter is a 32-bit maximum internal counter of standard IETF ChaCha20 variant.
const max_32bit_counter = u64(max_u32)

// default chacha20 quarter round number
const default_qround_nr = 10

// Stream is an internal structure where main ChaCha20 algorithm operates on.
@[noinit]
struct Stream {
mut:
    // underlying stream's key
    key [8]u32
    // underlying stream's nonce with internal counter
    nonce [4]u32

    // The mode (variant) of this ChaCha20 stream
    // Standard IETF variant or original (from DJ Bernstein) variant, set on creation.
    mode CipherMode = .standard
    // Flag that tells whether this stream was an extended XChaCha20 standard variant.
    // only make sense when mode == .standard
    extended bool
    // Flag tells whether this stream has reached the counter limit
    overflow bool

    // counter-independent precomputed values
    precomp bool
    // vfmt off
    p1  u32 p5  u32 p9  u32 p13 u32
    p2  u32 p6  u32 p10 u32 p14 u32
    p3  u32 p7  u32 p11 u32 p15 u32
    // vfmt on
}

// new_stream_with_options creates a new chacha20 stream with provided options.
// The supported nonce size is 8, 12 or 24 bytes.
@[direct_array_access; inline]
fn new_stream_with_options(key []u8, nonce []u8, opt Options) !Stream {
    if key.len != key_size {
        return error('Bad key size provided')
    }
    // setup for default value
    mut mode := CipherMode.standard
    mut extended := false

    // Based on the nonce.len and option supplied, it determines the variant (mode) and
    // extended form of the new chacha20 stream intended to create.
    match nonce.len {
        nonce_size {}
        x_nonce_size {
            extended = true
            if opt.use_64bit_counter {
                mode = .original
            }
        }
        orig_nonce_size {
            mode = .original
        }
        else {
            return error('new_stream_with_options: unsupported nonce size')
        }
    }

    // if this an extended chacha20 construct, derives a new key and nonce
    new_key, new_nonce := if extended {
        xkey, xnonce := derive_xchacha20_key_nonce(key, nonce, opt.use_64bit_counter)!
        xkey, xnonce
    } else {
        // otherwise, use provided key and nonce
        key, nonce
    }
    // Build a new stream and setup the key
    mut b := Stream{
        mode:     mode
        extended: extended
    }
    // store the key
    b.key[0] = binary.little_endian_u32(new_key[0..4])
    b.key[1] = binary.little_endian_u32(new_key[4..8])
    b.key[2] = binary.little_endian_u32(new_key[8..12])
    b.key[3] = binary.little_endian_u32(new_key[12..16])
    b.key[4] = binary.little_endian_u32(new_key[16..20])
    b.key[5] = binary.little_endian_u32(new_key[20..24])
    b.key[6] = binary.little_endian_u32(new_key[24..28])
    b.key[7] = binary.little_endian_u32(new_key[28..32])

    // store the nonce
    if b.mode == .standard {
        // in standard IETF variant, first nonce was used as internal counter
        b.nonce[0] = 0
        b.nonce[1] = binary.little_endian_u32(new_nonce[0..4])
        b.nonce[2] = binary.little_endian_u32(new_nonce[4..8])
        b.nonce[3] = binary.little_endian_u32(new_nonce[8..12])
    } else {
        // in the original variant, two's of first counter servers as 64-bit counter value
        b.nonce[0] = 0
        b.nonce[1] = 0

        b.nonce[2] = binary.little_endian_u32(new_nonce[0..4])
        b.nonce[3] = binary.little_endian_u32(new_nonce[4..8])
    }
    return b
}

// reset resets internal stream
@[unsafe]
fn (mut s Stream) reset() {
    // we dont reset s.mode and s.extended
    unsafe {
        _ := vmemset(&s.key, 0, 32)
        _ := vmemset(&s.nonce, 0, 16)
    }
    s.precomp = false
    s.p1, s.p5, s.p9, s.p13 = u32(0), u32(0), u32(0), u32(0)
    s.p2, s.p6, s.p10, s.p14 = u32(0), u32(0), u32(0), u32(0)
    s.p3, s.p7, s.p11, s.p15 = u32(0), u32(0), u32(0), u32(0)
}

// new_curr_state creates a new State from current stream
@[direct_array_access]
fn (s Stream) new_curr_state() State {
    // initializes ChaCha20 state
    //      0:cccccccc   1:cccccccc   2:cccccccc   3:cccccccc
    //      4:kkkkkkkk   5:kkkkkkkk   6:kkkkkkkk   7:kkkkkkkk
    //      8:kkkkkkkk   9:kkkkkkkk  10:kkkkkkkk  11:kkkkkkkk
    //     12:bbbbbbbb  13:nnnnnnnn  14:nnnnnnnn  15:nnnnnnnn
    //
    // where c=constant k=key b=blockcounter n=nonce
    mut state := State{}
    // load chacha20 constant into state
    state[0] = cc0
    state[1] = cc1
    state[2] = cc2
    state[3] = cc3
    // load key into state
    for i, k in s.key {
        state[i + 4] = k
    }
    // load nonce into state
    for j, v in s.nonce {
        state[j + 12] = v
    }
    return state
}

// keystream_full process with full size of src being processed
@[direct_array_access]
fn (mut s Stream) keystream_full(mut dst []u8, src []u8) ! {
    if s.overflow {
        return error('chacha20: keystream_full counter has reached the limit')
    }
    // number of block to be processed
    nr_blocks := src.len / block_size
    // check for counter overflow
    if s.check_ctr(u64(nr_blocks)) {
        s.overflow = true
        return error('chacha20: internal counter overflow')
    }
    mut idx := 0
    // process for full block_size-d msg
    for i := 0; i < nr_blocks; i++ {
        // for every block_sized message, we generates 64-bytes block key stream
        // and then xor-ing this block with generated key stream
        block := unsafe { src[i * block_size..(i + 1) * block_size] }
        ks := s.keystream()!
        for j, b in ks {
            dst[idx + j] = block[j] ^ b
        }
        // updates position
        idx += block_size
    }

    // process for remaining partial block
    if src.len % block_size != 0 {
        last_block := unsafe { src[nr_blocks * block_size..] }
        // generates one 64-bytes keystream block, and xor-ing bytes
        // in last_block with the key stream
        ks := s.keystream()!
        for i, b in last_block {
            dst[idx + i] = b ^ ks[i]
        }
        idx += last_block.len
    }
}

// keystream generates and retursns a 64-bytes block of key stream and increases internal counter.
@[direct_array_access]
fn (mut s Stream) keystream() ![]u8 {
    // initializes current state and working state
    mut awal := s.new_curr_state()
    mut ws := awal.clone()

    // precomputes cache counter-independent  values
    if s.mode == .standard && !s.precomp {
        s.precomp(awal)
    }
    // remaining first column round
    if s.mode == .standard {
        mut fcr := Quartet{awal[0], awal[4], awal[8], awal[12]}
        qround_on_quartet(mut fcr)

        // First diagonal round.
        qround_on_state_with_quartet(mut ws, fcr.e0, s.p5, s.p10, s.p15, 0, 5, 10, 15)
        qround_on_state_with_quartet(mut ws, s.p1, s.p6, s.p11, fcr.e3, 1, 6, 11, 12)
        qround_on_state_with_quartet(mut ws, s.p2, s.p7, fcr.e2, s.p13, 2, 7, 8, 13)
        qround_on_state_with_quartet(mut ws, s.p3, fcr.e1, s.p9, s.p14, 3, 4, 9, 14)
    }

    // The remaining quarter rounds
    //
    // For standard variant, the first column-round was already precomputed,
    // For original variant, its use full quarter round number.
    //
    // perform chacha20 quarter round n-times
    n := if s.mode == .standard { 9 } else { default_qround_nr }
    ws.qround(n)

    // Adding the working state values with inital state values.
    // We dont performs xor-ing here, its done on xor_key_stream and or keystream_full.
    for i, _ in ws {
        ws[i] += awal[i]
    }
    // increases stream internal counter
    s.inc_ctr(mut awal)!

    // serializes current working state in little-endian form
    mut block := []u8{len: block_size}
    for i, v in ws {
        block[i * 4] = u8(v)
        block[i * 4 + 1] = u8(v >> 8)
        block[i * 4 + 2] = u8(v >> 16)
        block[i * 4 + 3] = u8(v >> 24)
    }
    return block
}

// precomp performs quarter round on counter-independent quartet values on running state st.
@[direct_array_access; inline]
fn (mut s Stream) precomp(st State) {
    mut pcr1 := Quartet{st[1], st[5], st[9], st[13]}
    mut pcr2 := Quartet{st[2], st[6], st[10], st[14]}
    mut pcr3 := Quartet{st[3], st[7], st[11], st[15]}

    qround_on_quartet(mut pcr1)
    qround_on_quartet(mut pcr2)
    qround_on_quartet(mut pcr3)

    s.p1 = pcr1.e0
    s.p5 = pcr1.e1
    s.p9 = pcr1.e2
    s.p13 = pcr1.e3

    s.p2 = pcr2.e0
    s.p6 = pcr2.e1
    s.p10 = pcr2.e2
    s.p14 = pcr2.e3

    s.p3 = pcr3.e0
    s.p7 = pcr3.e1
    s.p11 = pcr3.e2
    s.p15 = pcr3.e3

    s.precomp = true
}

// Handling of Stream's internal counter
//

// ctr returns a current Stream's counter as u64 value.
@[direct_array_access; inline]
fn (b Stream) ctr() u64 {
    match b.mode {
        // In the original mode, counter was 64-bit size
        // stored on b.nonce[0], and b.nonce[1]
        .original {
            return u64(b.nonce[1]) << 32 | u64(b.nonce[0])
        }
        .standard {
            // in standard mode, counter was 32-bit value, stored on b.nonce[0]
            return u64(b.nonce[0])
        }
    }
}

// set_ctr sets stream internal counter
@[direct_array_access; inline]
fn (mut b Stream) set_ctr(ctr u64) {
    match b.mode {
        .original {
            b.nonce[0] = u32(ctr)
            b.nonce[1] = u32(ctr >> 32)
        }
        .standard {
            // check for ctr value that may exceed the counter limit
            if ctr > max_32bit_counter {
                panic('set_ctr: counter value exceed the limit ')
            }
            b.nonce[0] = u32(ctr)
        }
    }
}

// check_ctr checks for counter overflow when added by value.
// It returns true on counter overflow.
@[inline]
fn (b Stream) check_ctr(value u64) bool {
    ctr := b.ctr()
    sum := ctr + value
    max := b.max_ctr()
    if sum < ctr || sum < value || sum > max {
        return true
    }
    return false
}

// inc_ctr increases stream counter by one from the current state st
@[direct_array_access]
fn (mut s Stream) inc_ctr(mut st State) ! {
    // updates internal counter
    if s.mode == .original {
        st[12] += 1
        // first counter reset ?
        if st[12] == 0 {
            // increase second counter, if reset, mark as an overflow and return error
            st[13] += 1
            if st[13] == 0 {
                s.overflow = true
                return error('chacha20.keystream: 64-bit counter reached')
            }
        }
        // store the counter
        s.nonce[0] = st[12]
        s.nonce[1] = st[13]
    } else {
        st[12] += 1
        if st[12] == 0 {
            s.overflow = true
            return error('chacha20.keystream: overflow 32-bit counter')
        }
        s.nonce[0] = st[12]
    }
}

// max_ctr returns maximum counter value of this stream variant
@[inline]
fn (b Stream) max_ctr() u64 {
    match b.mode {
        .original { return max_64bit_counter }
        .standard { return max_32bit_counter }
    }
}

// State represents the running 64-bytes of chacha20 stream,
pub type State = [16]u32

// clone returns a new copy of this state.
@[direct_array_access; inline]
pub fn (s State) clone() State {
    mut sc := State{}
    for i, v in s {
        sc[i] = v
    }
    return sc
}

// reset resets internal values of this state.
@[direct_array_access; inline]
pub fn (mut s State) reset() {
    for i, _ in s {
        s[i] = u32(0)
    }
}

// qround performs quarter round on the working state ws with round number specified in nr.
// Its responsibility the user to provide the correct round number.
@[direct_array_access]
pub fn (mut ws State) qround(nr int) {
    for i := 0; i < nr; i++ {
        // Column-round
        //  0 |  1 |  2 |  3
        //  4 |  5 |  6 |  7
        //  8 |  9 | 10 | 11
        // 12 | 13 | 14 | 15
        qround_on_state(mut ws, 0, 4, 8, 12) // 0
        qround_on_state(mut ws, 1, 5, 9, 13) // 1
        qround_on_state(mut ws, 2, 6, 10, 14) // 2
        qround_on_state(mut ws, 3, 7, 11, 15) // 3

        // Diagonal round.
        //   0 \  1 \  2 \  3
        //   5 \  6 \  7 \  4
        //  10 \ 11 \  8 \  9
        //  15 \ 12 \ 13 \ 14
        qround_on_state(mut ws, 0, 5, 10, 15)
        qround_on_state(mut ws, 1, 6, 11, 12)
        qround_on_state(mut ws, 2, 7, 8, 13)
        qround_on_state(mut ws, 3, 4, 9, 14)
    }
}

// qround_on_state_with_quartet run qround_on_state by previously set up state values in offset
// (a,b,c,d) with values from quartet (q0, q1, q2, q3)
@[direct_array_access]
fn qround_on_state_with_quartet(mut s State, q0 u32, q1 u32, q2 u32, q3 u32, a int, b int, c int, d int) {
    s[a] = q0
    s[b] = q1
    s[c] = q2
    s[d] = q3
    qround_on_state(mut s, a, b, c, d)
}

// qround_on_state performs chacha20 quarter round on states with quartet index a, b, c, d.
@[direct_array_access]
fn qround_on_state(mut s State, a int, b int, c int, d int) {
    // a += b;  d ^= a;  d <<<= 16;
    s[a] += s[b]
    s[d] ^= s[a]
    s[d] = bits.rotate_left_32(s[d], 16)

    // c += d;  b ^= c;  b <<<= 12;
    s[c] += s[d]
    s[b] ^= s[c]
    s[b] = bits.rotate_left_32(s[b], 12)

    // a += b;  d ^= a;  d <<<=  8;
    s[a] += s[b]
    s[d] ^= s[a]
    s[d] = bits.rotate_left_32(s[d], 8)

    // c += d;  b ^= c;  b <<<=  7;
    s[c] += s[d]
    s[b] ^= s[c]
    s[b] = bits.rotate_left_32(s[b], 7)
}

// quartet of u32 values.
struct Quartet {
mut:
    e0 u32
    e1 u32
    e2 u32
    e3 u32
}

// qround_on_quartet runs chacha20 quarter round run on Quartet q.
fn qround_on_quartet(mut q Quartet) {
    // a += b;  d ^= a;  d <<<= 16;
    q.e0 += q.e1
    q.e3 ^= q.e0
    q.e3 = bits.rotate_left_32(q.e3, 16)

    // c += d;  b ^= c;  b <<<= 12;
    q.e2 += q.e3
    q.e1 ^= q.e2
    q.e1 = bits.rotate_left_32(q.e1, 12)

    // a += b;  d ^= a;  d <<<  8;
    q.e0 += q.e1
    q.e3 ^= q.e0
    q.e3 = bits.rotate_left_32(q.e3, 8)

    // c += d;  b ^= c;  b <<<=  7;
    q.e2 += q.e3
    q.e1 ^= q.e2
    q.e1 = bits.rotate_left_32(q.e1, 7)
}