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

use aes::cipher::generic_array::GenericArray;
use aes::cipher::{KeyIvInit, StreamCipher};
use aes::Aes128;

const BLOCK_SIZE: usize = 16;

type Aes128Ctr = ctr::Ctr128BE<Aes128>;

pub const KEY_SIZE: usize = 16;

pub type Key = [u8; KEY_SIZE];

pub struct PRG(Aes128Ctr);

impl PRG {
    pub fn new(key: &[u8; KEY_SIZE]) -> Self {
        let key = GenericArray::from_slice(key);
        let nonce = GenericArray::from_slice(&[0u8; 16]);
        PRG(Aes128Ctr::new(key, nonce))
    }

    pub fn xor_bytes(&mut self, dst: &mut [u8]) {
        debug_assert_eq!(dst.len() % BLOCK_SIZE, 0);
        self.0.apply_keystream(dst);
    }

    pub fn gen(&mut self, dst: &mut [u8]) {
        {
            let (_prefix, aligned, _suffix) = unsafe { dst.align_to_mut::<u128>() };
            debug_assert_eq!(_prefix.len(), 0);
            debug_assert_eq!(_suffix.len(), 0);
            for word in aligned.iter_mut() {
                *word = 0;
            }
        }
        self.xor_bytes(dst);
    }
}