module chacha20

import rand
import encoding.hex

// Test for Stream counter handling.
// See the discussion at [here](https://discord.com/channels/592103645835821068/592114487759470596/1417900997090607215)
fn test_stream_counter_handling() ! {
    // creates a original mode of the cipher with 64-bit counter
    mut ctx := new_cipher(rand.bytes(32)!, rand.bytes(8)!)!
    // set the cipher's counter near the maximum of 64-bit counter
    ctr := max_u64 - 2
    ctx.set_counter(ctr)

    // by setting internal counter into near of max 64-bit counter,
    // it need a message with minimum length of  2*block_size bytes to reach the limit.
    // let's build this message with 2 * block_size bytes in size
    msg0 := []u8{len: 2 * block_size}
    mut dst := []u8{len: msg0.len}
    ctx.xor_key_stream(mut dst, msg0)
    // at this step, the counter has reached the maximum_64bit_counter, but still not overflow
    assert ctx.Stream.overflow == false
    assert ctx.Stream.ctr() == max_64bit_counter

    // after above process, the counter should reach the maximum limit
    // we use keystream_full to test this counter handling, because
    // xor_key_stream would panic on counter reset
    msg1 := []u8{len: block_size}
    ctx.Stream.keystream_full(mut dst[..block_size], msg1) or {
        assert ctx.Stream.overflow == true
        assert err == error('chacha20: internal counter overflow')
        return
    }
}

fn test_qround_on_state() {
    mut s := State{}
    s[0] = 0x11111111
    s[1] = 0x01020304
    s[2] = 0x9b8d6f43
    s[3] = 0x01234567

    qround_on_state(mut s, 0, 1, 2, 3)
    assert s[0] == 0xea2a92f4
    assert s[1] == 0xcb1cf8ce
    assert s[2] == 0x4581472e
    assert s[3] == 0x5881c4bb
}

fn test_state_of_chacha20_block_simple() ! {
    key := '000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f'
    key_bytes := hex.decode(key)!

    nonce := '000000090000004a00000000'
    nonce_bytes := hex.decode(nonce)!

    mut stream := new_stream_with_options(key_bytes, nonce_bytes)!

    mut block := []u8{len: block_size}
    stream.set_ctr(1)
    stream.keystream_full(mut block, block)!

    expected_raw_bytes := '10f1e7e4d13b5915500fdd1fa32071c4c7d1f4c733c068030422aa9ac3d46c4ed2826446079faa0914c2d705d98b02a2b5129cd1de164eb9cbd083e8a2503c4e'
    exp_bytes := hex.decode(expected_raw_bytes)!

    assert block == exp_bytes
}

fn test_keystream_encryption() ! {
    for val in blocks_testcases {
        key := hex.decode(val.key)!
        nonce := hex.decode(val.nonce)!

        mut stream := new_stream_with_options(key, nonce)!
        stream.set_ctr(val.counter)

        mut block := []u8{len: block_size}
        stream.keystream_full(mut block, block)!
        exp_bytes := hex.decode(val.output)!

        assert block == exp_bytes
    }
}

struct BlockCase {
    key     string
    nonce   string
    counter u32
    output  string
}

const blocks_testcases = [
    // section 2.3.4 https://datatracker.ietf.org/doc/html/rfc8439#section-2.3.2
    BlockCase{
        key:     '000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f'
        nonce:   '000000090000004a00000000'
        counter: u32(1)
        output:  '10f1e7e4d13b5915500fdd1fa32071c4c7d1f4c733c068030422aa9ac3d46c4ed2826446079faa0914c2d705d98b02a2b5129cd1de164eb9cbd083e8a2503c4e'
    },
    // https://datatracker.ietf.org/doc/html/rfc8439#appendix-A.1.1
    BlockCase{
        key:     '0000000000000000000000000000000000000000000000000000000000000000'
        nonce:   '000000000000000000000000'
        counter: u32(0)
        output:  '76b8e0ada0f13d90405d6ae55386bd28bdd219b8a08ded1aa836efcc8b770dc7da41597c5157488d7724e03fb8d84a376a43b8f41518a11cc387b669b2ee6586'
    },
    // #appendix-A.1.2
    BlockCase{
        key:     '0000000000000000000000000000000000000000000000000000000000000000'
        nonce:   '000000000000000000000000'
        counter: u32(1)
        output:  '9f07e7be5551387a98ba977c732d080dcb0f29a048e3656912c6533e32ee7aed29b721769ce64e43d57133b074d839d531ed1f28510afb45ace10a1f4b794d6f'
    },
]