A Rust implementation of the ChaCha20-Poly1305 algorithm in RFC 7539 https://www.rfc-editor.org/rfc/rfc7539

chacha20.rs

// A Rust implementation of the ChaCha20-Poly1305 algorithm in RFC 7539 https://www.rfc-editor.org/rfc/rfc7539

fn from_le_bytes(x: &[u8]) -> u32 {
    u32::from_le_bytes([x[0], x[1], x[2], x[3]])
}

fn pad16(data: &[u8]) -> Vec<u8> {
    return vec![0; 16 - (data.len() % 16)];
}

fn quarter_round(a: usize, b: usize, c: usize, d: usize, block: &mut [u32; 16]) {
    block[a] = block[a].wrapping_add(block[b]);
    block[d] ^= block[a];
    block[d] = block[d].rotate_left(16);

    block[c] = block[c].wrapping_add(block[d]);
    block[b] ^= block[c];
    block[b] = block[b].rotate_left(12);

    block[a] = block[a].wrapping_add(block[b]);
    block[d] ^= block[a];
    block[d] = block[d].rotate_left(8);

    block[c] = block[c].wrapping_add(block[d]);
    block[b] ^= block[c];
    block[b] = block[b].rotate_left(7);
}

fn double_round(mut block: [u32; 16]) -> [u32; 16] {
    quarter_round(0, 4, 8, 12, &mut block);
    quarter_round(1, 5, 9, 13, &mut block);
    quarter_round(2, 6, 10, 14, &mut block);
    quarter_round(3, 7, 11, 15, &mut block);

    quarter_round(0, 5, 10, 15, &mut block);
    quarter_round(1, 6, 11, 12, &mut block);
    quarter_round(2, 7, 8, 13, &mut block);
    quarter_round(3, 4, 9, 14, &mut block);

    block
}

fn poly1305_mac(msg: &Vec<u8>, key: &Vec<u8>) -> Vec<u8> {
    let r = u128::from_le_bytes([
        key[0], key[1], key[2], key[3], key[4], key[5], key[6], key[7], key[8], key[9], key[10],
        key[11], key[12], key[13], key[14], key[15],
    ]) & 0x0ffffffc0ffffffc0ffffffc0fffffff;

    let s = u128::from_le_bytes([
        key[16], key[17], key[18], key[19], key[20], key[21], key[22], key[23], key[24], key[25],
        key[26], key[27], key[28], key[29], key[30], key[31],
    ]);

    let p: u128 = (1 << 127) - 5;
    let mut accumulator: u128 = 0;

    for i in 1..=((msg.len() + 15) / 16) {
        let start = (i - 1) * 16;
        let end = start + 16;
        let block = if end <= msg.len() {
            &msg[start..end]
        } else {
            let mut padded = [0u8; 16];
            padded[..(msg.len() - start)].copy_from_slice(&msg[start..]);
            padded[msg.len() - start] = 0x01;
            return padded.to_vec();
        };
        let n = u128::from_le_bytes([
            block[0], block[1], block[2], block[3], block[4], block[5], block[6], block[7],
            block[8], block[9], block[10], block[11], block[12], block[13], block[14], block[15],
        ]);
        accumulator = accumulator.wrapping_add(n);
        accumulator = accumulator.wrapping_mul(r);
        accumulator %= p;
    }

    accumulator += s;

    accumulator.to_le_bytes().to_vec()
}

pub struct ChaCha20 {
    key: Vec<u8>,
}

impl ChaCha20 {
    fn new(key: Vec<u8>) -> ChaCha20 {
        ChaCha20 { key }
    }

    fn keystream(&self, nonce: &[u8], counter: usize) -> Vec<u8> {
        let mut state: [u32; 16] = [
            0x61707865,
            0x3320646e,
            0x79622d32,
            0x6b206574,
            from_le_bytes(&self.key[0..4]),
            from_le_bytes(&self.key[4..8]),
            from_le_bytes(&self.key[8..12]),
            from_le_bytes(&self.key[12..16]),
            from_le_bytes(&self.key[16..20]),
            from_le_bytes(&self.key[20..24]),
            from_le_bytes(&self.key[24..28]),
            from_le_bytes(&self.key[28..32]),
            counter as u32,
            from_le_bytes(&nonce[0..4]),
            from_le_bytes(&nonce[4..8]),
            from_le_bytes(&nonce[8..12]),
        ];

        let mut working_state = state.clone();
        for _ in 0..10 {
            working_state = double_round(working_state);
        }

        for i in 0..16 {
            state[i] = state[i].wrapping_add(working_state[i]);
        }

        let mut result: Vec<u8> = Vec::new();
        for i in 0..16 {
            result.extend_from_slice(&state[i].to_le_bytes());
        }

        result
    }

    fn encrypt_stream(&self, msg: &[u8], nonce: &[u8]) -> Vec<u8> {
        let mut ciphertext: Vec<u8> = Vec::new();

        for (counter, chunk) in msg.chunks(64).enumerate() {
            let keystream = self.keystream(nonce, counter);

            for (i, byte) in chunk.iter().enumerate() {
                ciphertext.push(byte ^ keystream[i]);
            }
        }

        ciphertext
    }

    fn encrypt(&self, msg: &[u8], nonce: &[u8], aead: &[u8]) -> Vec<u8> {
        let otk = self.keystream(nonce, 0);
        let ciphertext = self.encrypt_stream(msg, nonce);
        let mut mac_data: Vec<u8> = Vec::new();
        mac_data.extend_from_slice(&aead);
        mac_data.append(&mut pad16(&ciphertext));
        mac_data.append(&mut ciphertext.clone());
        mac_data.append(&mut pad16(&ciphertext));

        let tag = poly1305_mac(&mac_data, &otk);

        [ciphertext, tag].concat()
    }

    fn decrypt(&self, msg: &[u8], nonce: &[u8], aead: &[u8]) -> Vec<u8> {
        let otk = self.keystream(nonce, 0);
        let ciphertext = &msg[..(msg.len() - 16)];
        let tag = &msg[(msg.len() - 16)..];
        let mut mac_data: Vec<u8> = Vec::new();
        mac_data.extend_from_slice(&aead);
        mac_data.append(&mut pad16(&ciphertext));
        mac_data.append(&mut ciphertext.to_vec());
        mac_data.append(&mut pad16(&ciphertext));
        let expected_tag = poly1305_mac(&mac_data, &otk);
        if expected_tag != tag {
            panic!("Invalid tag");
        }
        self.encrypt_stream(ciphertext, nonce)
    }
}