2021-10-08 06:12:54 -07:00
|
|
|
#[cfg(not(target_arch = "bpf"))]
|
|
|
|
use rand::{rngs::OsRng, Rng};
|
2021-10-07 15:21:31 -07:00
|
|
|
|
2021-10-08 06:12:54 -07:00
|
|
|
use aes::cipher::{BlockDecrypt, BlockEncrypt, NewBlockCipher};
|
|
|
|
use aes::{Aes128, Block};
|
|
|
|
|
|
|
|
use arrayref::array_ref;
|
2021-10-07 15:21:31 -07:00
|
|
|
|
2021-10-11 12:16:29 -07:00
|
|
|
use zeroize::Zeroize;
|
|
|
|
|
2021-10-07 15:21:31 -07:00
|
|
|
pub struct AES;
|
|
|
|
impl AES {
|
2021-10-11 12:16:29 -07:00
|
|
|
#[cfg(not(target_arch = "bpf"))]
|
|
|
|
#[allow(clippy::new_ret_no_self)]
|
2021-10-07 15:21:31 -07:00
|
|
|
pub fn new() -> AESKey {
|
2021-10-08 06:12:54 -07:00
|
|
|
let random_bytes = OsRng.gen::<[u8; 16]>();
|
|
|
|
AESKey(random_bytes)
|
2021-10-07 15:21:31 -07:00
|
|
|
}
|
|
|
|
|
2021-10-11 12:16:29 -07:00
|
|
|
#[cfg(not(target_arch = "bpf"))]
|
2021-10-07 15:21:31 -07:00
|
|
|
pub fn encrypt(sk: &AESKey, amount: u64) -> AESCiphertext {
|
2021-10-08 06:12:54 -07:00
|
|
|
let amount_bytes = amount.to_le_bytes();
|
|
|
|
|
|
|
|
let mut aes_block: Block = [0_u8; 16].into();
|
|
|
|
aes_block[..8].copy_from_slice(&amount_bytes);
|
|
|
|
|
|
|
|
Aes128::new(&sk.0.into()).encrypt_block(&mut aes_block);
|
|
|
|
AESCiphertext(aes_block.into())
|
2021-10-07 15:21:31 -07:00
|
|
|
}
|
|
|
|
|
2021-10-11 12:16:29 -07:00
|
|
|
#[cfg(not(target_arch = "bpf"))]
|
2021-10-07 15:21:31 -07:00
|
|
|
pub fn decrypt(sk: &AESKey, ct: &AESCiphertext) -> u64 {
|
2021-10-08 06:12:54 -07:00
|
|
|
let mut aes_block: Block = ct.0.into();
|
|
|
|
Aes128::new(&sk.0.into()).decrypt_block(&mut aes_block);
|
|
|
|
|
|
|
|
let amount_bytes = array_ref![aes_block[..8], 0, 8];
|
|
|
|
u64::from_le_bytes(*amount_bytes)
|
2021-10-07 15:21:31 -07:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2021-10-11 12:16:29 -07:00
|
|
|
#[derive(Debug, Zeroize)]
|
2021-10-08 06:12:54 -07:00
|
|
|
pub struct AESKey([u8; 16]);
|
2021-10-07 15:21:31 -07:00
|
|
|
impl AESKey {
|
|
|
|
pub fn encrypt(&self, amount: u64) -> AESCiphertext {
|
|
|
|
AES::encrypt(self, amount)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2021-10-08 06:12:54 -07:00
|
|
|
#[derive(Debug)]
|
|
|
|
pub struct AESCiphertext([u8; 16]);
|
2021-10-07 15:21:31 -07:00
|
|
|
impl AESCiphertext {
|
|
|
|
pub fn decrypt(&self, sk: &AESKey) -> u64 {
|
|
|
|
AES::decrypt(sk, self)
|
|
|
|
}
|
|
|
|
}
|
2021-10-08 06:12:54 -07:00
|
|
|
|
|
|
|
#[cfg(test)]
|
|
|
|
mod tests {
|
|
|
|
use super::*;
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
fn test_aes_encrypt_decrypt_correctness() {
|
|
|
|
let sk = AES::new();
|
|
|
|
let amount = 55;
|
|
|
|
|
|
|
|
let ct = sk.encrypt(amount);
|
|
|
|
let decrypted_amount = ct.decrypt(&sk);
|
|
|
|
|
|
|
|
assert_eq!(amount, decrypted_amount);
|
|
|
|
}
|
|
|
|
}
|