use crate::shortint::{Ciphertext, ServerKey};
use crate::transciphering::ciphers::aes::AesIv;
use crate::transciphering::{FheKeyStream, InsufficientKeystream, 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,
) -> Result<FheKeyStream, InsufficientKeystream> {
let end_counter = self
.counter
.checked_add(n_bits as u64)
.ok_or(InsufficientKeystream)?;
let skip_head = (self.counter % 128) as usize;
let start_block = self.counter / 128;
let n_blocks = end_counter.div_ceil(128) - start_block;
let blocks: Vec<[Ciphertext; 128]> = (0..n_blocks)
.into_par_iter()
.map(|i| self.keystream_block(sks, (start_block + i) as u128))
.collect();
self.counter = end_counter;
let flat: Vec<Ciphertext> = blocks
.into_iter()
.flatten()
.skip(skip_head)
.take(n_bits)
.collect();
Ok(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
}
}