use crate::hazmat::kalyna_gcm::{GcmError, Kalyna256_256Gcm};
use crate::hazmat::kupyna_kmac::Kupyna256Kmac;
use core::fmt;
use zeroize::Zeroize;
const TAG_LEN: usize = 16;
const REKEY_CONTEXT: &[u8] = b"DSTU-secretstream-rekey";
pub struct Key([u8; 32]);
impl Drop for Key {
fn drop(&mut self) {
self.0.zeroize();
}
}
impl Key {
#[cfg(any(feature = "std", feature = "getrandom"))]
pub fn generate() -> Result<Self, crate::randombytes::RandomError> {
let mut bytes = [0u8; 32];
crate::randombytes::randombytes_buf(&mut bytes)?;
Ok(Key(bytes))
}
#[must_use]
pub fn from_bytes(bytes: [u8; 32]) -> Self {
Key(bytes)
}
#[must_use]
pub fn as_bytes(&self) -> &[u8; 32] {
&self.0
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Tag {
Message,
Push,
Rekey,
Final,
}
impl Tag {
#[must_use]
pub fn to_byte(self) -> u8 {
match self {
Tag::Message => 0x00,
Tag::Push => 0x01,
Tag::Rekey => 0x02,
Tag::Final => 0x03,
}
}
#[must_use]
pub fn from_byte(byte: u8) -> Option<Self> {
match byte {
0x00 => Some(Tag::Message),
0x01 => Some(Tag::Push),
0x02 => Some(Tag::Rekey),
0x03 => Some(Tag::Final),
_ => None,
}
}
}
#[derive(Debug)]
pub enum SecretstreamError {
InvalidLength,
TagMismatch,
UnknownTag,
StreamFinalized,
#[cfg(any(feature = "std", feature = "getrandom"))]
Random(crate::randombytes::RandomError),
}
impl fmt::Display for SecretstreamError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
SecretstreamError::InvalidLength => write!(f, "buffer length mismatch"),
SecretstreamError::TagMismatch => write!(f, "authentication failed"),
SecretstreamError::UnknownTag => write!(f, "unrecognized chunk tag byte"),
SecretstreamError::StreamFinalized => {
write!(f, "stream already finalized, no more chunks accepted")
}
#[cfg(any(feature = "std", feature = "getrandom"))]
SecretstreamError::Random(e) => write!(f, "{e}"),
}
}
}
impl core::error::Error for SecretstreamError {}
#[cfg(any(feature = "std", feature = "getrandom"))]
impl From<crate::randombytes::RandomError> for SecretstreamError {
fn from(e: crate::randombytes::RandomError) -> Self {
SecretstreamError::Random(e)
}
}
fn rekey(subkey: &mut [u8; 32]) {
let Ok(new_subkey) = Kupyna256Kmac::mac(subkey, REKEY_CONTEXT) else {
unreachable!("subkey is always exactly 32 bytes, Kupyna256Kmac's own mac_len")
};
subkey.zeroize();
*subkey = new_subkey;
}
fn chunk_iv(counter: u64) -> [u8; 32] {
let mut iv = [0u8; 32];
iv[..8].copy_from_slice(&counter.to_le_bytes());
iv
}
fn chunk_aad(counter: u64, tag_byte: u8) -> [u8; 9] {
let mut aad = [0u8; 9];
aad[..8].copy_from_slice(&counter.to_le_bytes());
aad[8] = tag_byte;
aad
}
pub struct PushState {
subkey: [u8; 32],
counter: u64,
finalized: bool,
}
impl Drop for PushState {
fn drop(&mut self) {
self.subkey.zeroize();
}
}
impl PushState {
#[cfg(any(feature = "std", feature = "getrandom"))]
pub fn init(key: &Key) -> Result<(Self, [u8; 32]), SecretstreamError> {
let mut header = [0u8; 32];
crate::randombytes::randombytes_buf(&mut header)?;
let Ok(subkey) = Kupyna256Kmac::mac(key.as_bytes(), &header) else {
unreachable!("Key::as_bytes() is always exactly 32 bytes, Kupyna256Kmac's own mac_len")
};
Ok((
PushState {
subkey,
counter: 0,
finalized: false,
},
header,
))
}
#[must_use]
pub fn is_finalized(&self) -> bool {
self.finalized
}
pub fn push(
&mut self,
tag: Tag,
plaintext: &[u8],
ciphertext_out: &mut [u8],
) -> Result<[u8; TAG_LEN], SecretstreamError> {
if self.finalized {
return Err(SecretstreamError::StreamFinalized);
}
if ciphertext_out.len() != plaintext.len() {
return Err(SecretstreamError::InvalidLength);
}
let cipher = Kalyna256_256Gcm::new(&self.subkey);
let iv = chunk_iv(self.counter);
let aad = chunk_aad(self.counter, tag.to_byte());
let Ok(full_tag) = cipher.encrypt(&iv, &aad, plaintext, ciphertext_out) else {
unreachable!("ciphertext_out.len() == plaintext.len() checked above")
};
self.counter += 1;
match tag {
Tag::Rekey => rekey(&mut self.subkey),
Tag::Final => self.finalized = true,
Tag::Message | Tag::Push => {}
}
let mut out = [0u8; TAG_LEN];
out.copy_from_slice(&full_tag[..TAG_LEN]);
Ok(out)
}
}
pub struct PullState {
subkey: [u8; 32],
counter: u64,
finalized: bool,
}
impl Drop for PullState {
fn drop(&mut self) {
self.subkey.zeroize();
}
}
impl PullState {
#[must_use]
pub fn init(key: &Key, header: &[u8; 32]) -> Self {
let Ok(subkey) = Kupyna256Kmac::mac(key.as_bytes(), header) else {
unreachable!("Key::as_bytes() is always exactly 32 bytes, Kupyna256Kmac's own mac_len")
};
PullState {
subkey,
counter: 0,
finalized: false,
}
}
#[must_use]
pub fn is_finalized(&self) -> bool {
self.finalized
}
pub fn pull(
&mut self,
tag_byte: u8,
ciphertext: &[u8],
auth_tag: &[u8],
plaintext_out: &mut [u8],
) -> Result<Tag, SecretstreamError> {
if self.finalized {
return Err(SecretstreamError::StreamFinalized);
}
let Some(tag) = Tag::from_byte(tag_byte) else {
return Err(SecretstreamError::UnknownTag);
};
if plaintext_out.len() != ciphertext.len() {
return Err(SecretstreamError::InvalidLength);
}
let cipher = Kalyna256_256Gcm::new(&self.subkey);
let iv = chunk_iv(self.counter);
let aad = chunk_aad(self.counter, tag_byte);
match cipher.decrypt(&iv, &aad, ciphertext, auth_tag, plaintext_out) {
Ok(()) => {}
Err(GcmError::TagMismatch) => return Err(SecretstreamError::TagMismatch),
Err(GcmError::InvalidLength) => return Err(SecretstreamError::InvalidLength),
}
self.counter += 1;
match tag {
Tag::Rekey => rekey(&mut self.subkey),
Tag::Final => self.finalized = true,
Tag::Message | Tag::Push => {}
}
Ok(tag)
}
}