1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

use std::convert::TryInto;

use hacspec_lib::{bytes_to_le_u32s, Conversions};

type State = [u32; 16];
type StateIdx = usize;

type Constants = [u32; 4];

type Block = Vec<u8>;
pub type ChaChaIV = [u8; 12];
pub type ChaChaKey = [u8; 32];

fn chacha20_line(a: StateIdx, b: StateIdx, d: StateIdx, s: usize, m: State) -> State {
    let mut state = m;
    state[a] = state[a].wrapping_add(state[b]);
    state[d] ^= state[a];
    state[d] = state[d].rotate_left(s as u32);
    state
}

pub fn chacha20_quarter_round(
    a: StateIdx,
    b: StateIdx,
    c: StateIdx,
    d: StateIdx,
    state: State,
) -> State {
    let state = chacha20_line(a, b, d, 16, state);
    let state = chacha20_line(c, d, b, 12, state);
    let state = chacha20_line(a, b, d, 8, state);
    chacha20_line(c, d, b, 7, state)
}

fn chacha20_double_round(state: State) -> State {
    let state = chacha20_quarter_round(0, 4, 8, 12, state);
    let state = chacha20_quarter_round(1, 5, 9, 13, state);
    let state = chacha20_quarter_round(2, 6, 10, 14, state);
    let state = chacha20_quarter_round(3, 7, 11, 15, state);

    let state = chacha20_quarter_round(0, 5, 10, 15, state);
    let state = chacha20_quarter_round(1, 6, 11, 12, state);
    let state = chacha20_quarter_round(2, 7, 8, 13, state);
    chacha20_quarter_round(3, 4, 9, 14, state)
}

pub fn chacha20_rounds(state: State) -> State {
    let mut st = state;
    for _i in 0..10 {
        st = chacha20_double_round(st);
    }
    st
}

pub fn chacha20_core(ctr: u32, st0: State) -> State {
    let mut state = st0;
    state[12] += ctr;
    let k = chacha20_rounds(state);
    for i in 0..16 {
        state[i] = state[i].wrapping_add(k[i]);
    }
    state
}

pub fn chacha20_constants_init() -> Constants {
    let mut constants = [0u32; 4];
    constants[0] = 0x6170_7865u32;
    constants[1] = 0x3320_646eu32;
    constants[2] = 0x7962_2d32u32;
    constants[3] = 0x6b20_6574u32;
    constants
}

pub fn chacha20_init(key: ChaChaKey, iv: ChaChaIV, ctr: u32) -> State {
    let mut st = [0u32; 16];
    st[0..4].copy_from_slice(&chacha20_constants_init());
    st[4..12].copy_from_slice(&bytes_to_le_u32s(&key));
    st[12] = ctr;
    st[13..16].copy_from_slice(&bytes_to_le_u32s(&iv));

    st
}

pub fn chacha20_key_block(state: State) -> Block {
    let state = chacha20_core(0u32, state);

    state.to_le_bytes()
}

pub fn chacha20_key_block0(key: ChaChaKey, iv: ChaChaIV) -> Block {
    let state = chacha20_init(key, iv, 0u32);
    chacha20_key_block(state)
}

pub fn chacha20_encrypt_block(st0: State, ctr: u32, plain: &[u8]) -> Block {
    let st = chacha20_core(ctr, st0);
    let pl: State = bytes_to_le_u32s(plain).try_into().unwrap();
    let mut st_new = [0u32; 16];
    for i in 0..16 {
        st_new[i] = st[i] ^ pl[i];
    }
    let st = st_new;
    st.to_le_bytes()
}

pub fn chacha20_encrypt_last(st0: State, ctr: u32, plain: &[u8]) -> Vec<u8> {
    let mut b = [0u8; 64];
    b[0..plain.len()].copy_from_slice(plain);
    b = chacha20_encrypt_block(st0, ctr, &b).try_into().unwrap();
    b[0..plain.len()].to_owned()
}

pub fn chacha20_update(st0: State, m: &[u8]) -> Vec<u8> {
    let mut blocks_out = Vec::new();
    let n_blocks = m.len() / 64;
    let mchunks = m.chunks_exact(64);
    let last_block = mchunks.remainder();
    for (i, msg_block) in mchunks.enumerate() {
        let b = chacha20_encrypt_block(st0, i as u32, msg_block);
        blocks_out.extend_from_slice(&b);
    }
    if !last_block.is_empty() {
        let b = chacha20_encrypt_last(st0, n_blocks as u32, last_block);
        blocks_out.extend_from_slice(&b);
    }
    blocks_out
}

pub fn chacha20(key: ChaChaKey, iv: ChaChaIV, ctr: u32, m: &[u8]) -> Vec<u8> {
    let state = chacha20_init(key, iv, ctr);
    chacha20_update(state, m)
}