use crate::shortint::{Ciphertext, ServerKey};
use crate::transciphering::ciphers::aes::AesIv;
use crate::transciphering::{FheKeyStream, StreamCipherKind, Transcipherer};
use rayon::prelude::*;
use super::encrypt::encrypt_block;
use super::key::AesFheRoundKeys;
pub struct AesFheState {
key: AesFheRoundKeys,
iv: AesIv,
counter: u64,
}
impl AesFheState {
pub fn new(key: AesFheRoundKeys, iv: impl Into<AesIv>) -> Self {
Self {
key,
iv: iv.into(),
counter: 0,
}
}
fn keystream_block(&self, sks: &ServerKey, block_index: u128) -> [Ciphertext; 128] {
let counter_value = self.iv.to_u128().wrapping_add(block_index);
encrypt_block(sks, counter_value, &self.key)
}
}
impl Transcipherer for AesFheState {
fn kind(&self) -> StreamCipherKind {
StreamCipherKind::Aes
}
fn next_keystream_bits(&mut self, sks: &ServerKey, n_bits: usize) -> FheKeyStream {
let skip_head = (self.counter % 128) as usize;
let start_block = self.counter / 128;
let n_blocks = (skip_head + n_bits).div_ceil(128);
let blocks: Vec<[Ciphertext; 128]> = (0..n_blocks as u64)
.into_par_iter()
.map(|i| self.keystream_block(sks, (start_block + i) as u128))
.collect();
self.counter = self
.counter
.checked_add(n_bits as u64)
.expect("AesFheStream: keystream bit counter overflowed u64");
let flat: Vec<Ciphertext> = blocks
.into_iter()
.flatten()
.skip(skip_head)
.take(n_bits)
.collect();
FheKeyStream::from_raw_parts(flat)
}
fn seek(&mut self, _sks: &ServerKey, target_counter: u64) {
self.counter = target_counter;
}
fn current_counter(&self) -> u64 {
self.counter
}
}