use crate::sm3::BLOCK_SIZE as SM3_BLOCK;
use crate::sm4::cipher::Sm4Cipher;
use crate::tlcp::key_schedule::TlcpRole;
use alloc::vec::Vec;
use rand_core::TryCryptoRng;
use subtle::{Choice, ConditionallySelectable, ConstantTimeEq, ConstantTimeGreater};
use zeroize::{Zeroize, ZeroizeOnDrop};
pub const TLCP_RECORD_VERSION: [u8; 2] = [0x01, 0x01];
pub const CBC_KEY_BLOCK_LEN: usize = 128;
#[cfg(feature = "sm4-aead")]
pub const GCM_KEY_BLOCK_LEN: usize = 40;
const MAC_LEN: usize = 32;
const HEADER_LEN: usize = 13;
const MAX_PLAINTEXT: usize = 1 << 14;
#[derive(Clone, ZeroizeOnDrop)]
pub struct RecordKeysCbc {
pub(crate) mac_key: [u8; 32],
pub(crate) enc_key: [u8; 16],
}
#[cfg(feature = "sm4-aead")]
#[derive(Clone, ZeroizeOnDrop)]
pub struct RecordKeysGcm {
pub(crate) enc_key: [u8; 16],
pub(crate) salt: [u8; 4],
}
impl RecordKeysCbc {
#[must_use]
pub fn client_half(key_block: &[u8; 128]) -> Self {
let mut mac_key = [0u8; 32];
let mut enc_key = [0u8; 16];
mac_key.copy_from_slice(&key_block[0..32]);
enc_key.copy_from_slice(&key_block[64..80]);
Self { mac_key, enc_key }
}
#[must_use]
pub fn server_half(key_block: &[u8; 128]) -> Self {
let mut mac_key = [0u8; 32];
let mut enc_key = [0u8; 16];
mac_key.copy_from_slice(&key_block[32..64]);
enc_key.copy_from_slice(&key_block[80..96]);
Self { mac_key, enc_key }
}
#[must_use]
pub fn from_key_block(role: TlcpRole, key_block: &[u8; 128]) -> Self {
match role {
TlcpRole::Client => Self::client_half(key_block),
TlcpRole::Server => Self::server_half(key_block),
}
}
}
#[cfg(feature = "sm4-aead")]
impl RecordKeysGcm {
#[must_use]
pub fn client_half(key_block: &[u8; 40]) -> Self {
let mut enc_key = [0u8; 16];
let mut salt = [0u8; 4];
enc_key.copy_from_slice(&key_block[0..16]);
salt.copy_from_slice(&key_block[32..36]);
Self { enc_key, salt }
}
#[must_use]
pub fn server_half(key_block: &[u8; 40]) -> Self {
let mut enc_key = [0u8; 16];
let mut salt = [0u8; 4];
enc_key.copy_from_slice(&key_block[16..32]);
salt.copy_from_slice(&key_block[36..40]);
Self { enc_key, salt }
}
#[must_use]
pub fn from_key_block(role: TlcpRole, key_block: &[u8; 40]) -> Self {
match role {
TlcpRole::Client => Self::client_half(key_block),
TlcpRole::Server => Self::server_half(key_block),
}
}
}
fn record_header(seq: u64, content_type: u8, version: [u8; 2], len: u16) -> [u8; HEADER_LEN] {
let mut h = [0u8; HEADER_LEN];
h[0..8].copy_from_slice(&seq.to_be_bytes());
h[8] = content_type;
h[9..11].copy_from_slice(&version);
h[11..13].copy_from_slice(&len.to_be_bytes());
h
}
#[cfg(feature = "sm4-aead")]
fn gcm_nonce(salt: [u8; 4], explicit: [u8; 8]) -> [u8; 12] {
let mut nonce = [0u8; 12];
nonce[0..4].copy_from_slice(&salt);
nonce[4..12].copy_from_slice(&explicit);
nonce
}
#[cfg(feature = "sm4-aead")]
#[must_use]
pub fn protect_gcm(
keys: &RecordKeysGcm,
seq: u64,
content_type: u8,
version: [u8; 2],
plaintext: &[u8],
) -> Option<Vec<u8>> {
if plaintext.len() > MAX_PLAINTEXT {
return None;
}
let len = u16::try_from(plaintext.len()).ok()?;
let explicit = seq.to_be_bytes();
let nonce = gcm_nonce(keys.salt, explicit);
let aad = record_header(seq, content_type, version, len);
let (ct, tag) = crate::sm4::mode_gcm::encrypt(&keys.enc_key, &nonce, &aad, plaintext)?;
let mut out = Vec::with_capacity(8 + ct.len() + 16);
out.extend_from_slice(&explicit);
out.extend_from_slice(&ct);
out.extend_from_slice(&tag);
Some(out)
}
#[cfg(feature = "sm4-aead")]
#[must_use]
pub fn deprotect_gcm(
keys: &RecordKeysGcm,
seq: u64,
content_type: u8,
version: [u8; 2],
record: &[u8],
) -> Option<Vec<u8>> {
if record.len() < 8 + 16 {
return None;
}
let ct_len = record.len() - 8 - 16;
if ct_len > MAX_PLAINTEXT {
return None;
}
let len = u16::try_from(ct_len).ok()?;
let mut explicit = [0u8; 8];
explicit.copy_from_slice(&record[0..8]);
let nonce = gcm_nonce(keys.salt, explicit);
let ct = &record[8..8 + ct_len];
let mut tag = [0u8; 16];
tag.copy_from_slice(&record[8 + ct_len..]);
let aad = record_header(seq, content_type, version, len);
crate::sm4::mode_gcm::decrypt(&keys.enc_key, &nonce, &aad, ct, &tag)
}
#[inline]
fn inner_blocks(plaintext_len: usize) -> usize {
let n = SM3_BLOCK + HEADER_LEN + plaintext_len;
n / SM3_BLOCK + 1 + usize::from(n % SM3_BLOCK >= SM3_BLOCK - 8)
}
pub(crate) fn mac_equalized(
mac_key: &[u8; 32],
header: &[u8; 13],
plaintext: &[u8],
max_plaintext_len: usize,
) -> [u8; 32] {
let mut h = crate::hmac::HmacSm3::new(mac_key);
h.update(header);
h.update(plaintext);
let tag = h.finalize();
let actual = inner_blocks(plaintext.len());
let target = inner_blocks(max_plaintext_len);
let dummy = target.saturating_sub(actual);
let mut scratch = [0u32; 8];
let block = [0u8; 64];
let n = core::hint::black_box(dummy);
for _ in 0..n {
core::hint::black_box(&mut scratch);
crate::sm3::compress(&mut scratch, core::hint::black_box(&block));
}
let _ = core::hint::black_box(&scratch);
tag
}
#[allow(clippy::cast_possible_truncation, clippy::cast_lossless)]
fn check_tls_padding_ct(body: &[u8]) -> (Choice, usize) {
let n = body.len();
if n < MAC_LEN + 16 {
return (Choice::from(0u8), 0);
}
let padlen = body[n - 1]; let pad_region = padlen as u16 + 1; let window = core::cmp::min(256usize, n);
let mut bad: u8 = 0;
for i in 0..window {
let pos = (i as u16) + 1; let byte = body[n - 1 - i];
let in_pad = !pos.ct_gt(&pad_region); let diff = byte ^ padlen;
bad |= u8::conditional_select(&0u8, &diff, in_pad);
}
let region_ok = !pad_region.ct_gt(&((n - MAC_LEN) as u16));
let ok = region_ok & bad.ct_eq(&0u8);
let raw = (n as u64)
.wrapping_sub(MAC_LEN as u64)
.wrapping_sub(pad_region as u64);
let plaintext_len = u64::conditional_select(&0u64, &raw, ok) as usize;
(ok, plaintext_len)
}
#[allow(clippy::cast_lossless)]
fn extract_mac_ct(body: &[u8], plaintext_len: usize) -> [u8; 32] {
let n = body.len();
let mut out = [0u8; 32];
for i in 0..=(n - MAC_LEN) {
let hit = (i as u64).ct_eq(&(plaintext_len as u64));
for j in 0..MAC_LEN {
out[j] = u8::conditional_select(&out[j], &body[i + j], hit);
}
}
out
}
#[allow(clippy::missing_panics_doc)]
#[must_use]
pub fn protect_cbc<R: TryCryptoRng>(
keys: &RecordKeysCbc,
seq: u64,
content_type: u8,
version: [u8; 2],
plaintext: &[u8],
rng: &mut R,
) -> Option<Vec<u8>> {
if plaintext.len() > MAX_PLAINTEXT {
return None;
}
let len = u16::try_from(plaintext.len()).ok()?;
let hdr = record_header(seq, content_type, version, len);
let mut h = crate::hmac::HmacSm3::new(&keys.mac_key);
h.update(&hdr);
h.update(plaintext);
let mac = h.finalize();
let mut iv = [0u8; 16];
rng.try_fill_bytes(&mut iv).ok()?;
let mut out = Vec::with_capacity(16 + plaintext.len() + MAC_LEN + 16);
out.extend_from_slice(&iv);
out.extend_from_slice(plaintext);
out.extend_from_slice(&mac);
let padlen = 15 - ((out.len() - 16) % 16);
#[allow(clippy::cast_possible_truncation)]
let padbyte = padlen as u8;
for _ in 0..=padlen {
out.push(padbyte);
}
let cipher = Sm4Cipher::new(&keys.enc_key);
let mut prev = iv;
for chunk in out[16..].chunks_exact_mut(16) {
let block: &mut [u8; 16] = chunk.try_into().expect("16-byte chunk");
for j in 0..16 {
block[j] ^= prev[j];
}
cipher.encrypt_block(block);
prev = *block;
}
Some(out)
}
#[allow(
clippy::cast_possible_truncation,
clippy::cast_lossless,
clippy::missing_panics_doc
)]
#[must_use]
pub fn deprotect_cbc(
keys: &RecordKeysCbc,
seq: u64,
content_type: u8,
version: [u8; 2],
record: &[u8],
) -> Option<Vec<u8>> {
if record.len() < 16 + MAC_LEN + 16 {
return None; }
let body_len = record.len() - 16;
if body_len % 16 != 0 {
return None;
}
if body_len > MAX_PLAINTEXT + MAC_LEN + 16 {
return None; }
let mut iv = [0u8; 16];
iv.copy_from_slice(&record[0..16]);
let mut body = record[16..].to_vec();
let cipher = Sm4Cipher::new(&keys.enc_key);
let mut prev = iv;
for chunk in body.chunks_exact_mut(16) {
let block: &mut [u8; 16] = chunk.try_into().expect("16-byte chunk");
let saved = *block;
cipher.decrypt_block(block);
for j in 0..16 {
block[j] ^= prev[j];
}
prev = saved;
}
let (pad_ok, plaintext_len) = check_tls_padding_ct(&body);
let max_plaintext_len = body_len - MAC_LEN - 1;
let eff_len =
u64::conditional_select(&(max_plaintext_len as u64), &(plaintext_len as u64), pad_ok)
as usize;
let len16 = eff_len as u16;
let hdr = record_header(seq, content_type, version, len16);
let computed = mac_equalized(&keys.mac_key, &hdr, &body[..eff_len], max_plaintext_len);
let received = extract_mac_ct(&body, eff_len);
let mac_ok = computed.ct_eq(&received);
let valid = pad_ok & mac_ok;
if bool::from(valid) {
body.truncate(eff_len);
Some(body)
} else {
body.zeroize();
None
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::cast_possible_truncation)]
use super::*;
use rand_core::TryRng;
struct CountRng(u8);
impl TryRng for CountRng {
type Error = core::convert::Infallible;
fn try_next_u32(&mut self) -> core::result::Result<u32, Self::Error> {
Ok(0)
}
fn try_next_u64(&mut self) -> core::result::Result<u64, Self::Error> {
Ok(0)
}
fn try_fill_bytes(&mut self, dst: &mut [u8]) -> core::result::Result<(), Self::Error> {
for b in dst.iter_mut() {
*b = self.0;
self.0 = self.0.wrapping_add(1);
}
Ok(())
}
}
impl TryCryptoRng for CountRng {}
fn cbc_keys() -> RecordKeysCbc {
let mut kb = [0u8; 128];
for (i, b) in kb.iter_mut().enumerate() {
*b = (i as u8).wrapping_mul(3).wrapping_add(1);
}
RecordKeysCbc::client_half(&kb)
}
#[test]
fn cbc_keys_carve_offsets() {
let mut kb = [0u8; 128];
for (i, b) in kb.iter_mut().enumerate() {
*b = i as u8;
}
let c = RecordKeysCbc::client_half(&kb);
let s = RecordKeysCbc::server_half(&kb);
assert_eq!(c.mac_key, kb[0..32]);
assert_eq!(c.enc_key, kb[64..80]);
assert_eq!(s.mac_key, kb[32..64]);
assert_eq!(s.enc_key, kb[80..96]);
assert_eq!(
RecordKeysCbc::from_key_block(TlcpRole::Client, &kb).enc_key,
c.enc_key
);
assert_eq!(
RecordKeysCbc::from_key_block(TlcpRole::Server, &kb).mac_key,
s.mac_key
);
}
#[test]
fn inner_blocks_matches_sm3_spill_rule() {
for p in [0usize, 1, 42, 43, 44, 55, 56, 57, 100, 256, 1000] {
let n = 64 + 13 + p;
let expected = (n + 9).div_ceil(64);
assert_eq!(inner_blocks(p), expected, "p={p}");
}
}
#[test]
fn equalized_mac_matches_hmac_and_count_is_constant() {
let mac_key = [0x2bu8; 32];
for plen in [0usize, 1, 13, 42, 43, 44, 55, 56, 57, 100, 256, 1000] {
let hdr = record_header(9, 0x17, TLCP_RECORD_VERSION, plen as u16);
let pt: Vec<u8> = (0..plen).map(|i| i as u8).collect();
let mut h = crate::hmac::HmacSm3::new(&mac_key);
h.update(&hdr);
h.update(&pt);
let want = h.finalize();
let max = plen + 300;
let got = mac_equalized(&mac_key, &hdr, &pt, max);
assert_eq!(
got, want,
"equalized MAC must equal ordinary HMAC, plen={plen}"
);
assert!(
inner_blocks(plen) + (inner_blocks(max) - inner_blocks(plen)) == inner_blocks(max),
"equalization invariant, plen={plen}"
);
}
}
#[test]
fn cbc_round_trip_boundary_lengths() {
let keys = cbc_keys();
for len in [
0usize,
1,
15,
16,
17,
31,
32,
47,
48,
100,
1000,
MAX_PLAINTEXT,
] {
let pt: Vec<u8> = (0..len).map(|i| (i as u8).wrapping_mul(7)).collect();
let mut rng = CountRng(0xA0);
let rec = protect_cbc(&keys, 42, 0x17, TLCP_RECORD_VERSION, &pt, &mut rng).unwrap();
assert_eq!((rec.len() - 16) % 16, 0, "block-aligned body, len={len}");
let got = deprotect_cbc(&keys, 42, 0x17, TLCP_RECORD_VERSION, &rec);
assert_eq!(got.as_deref(), Some(&pt[..]), "round-trip len={len}");
}
}
#[test]
fn cbc_rejects_tamper_and_wrong_context() {
let keys = cbc_keys();
let pt = b"the quick brown fox jumps over the lazy tlcp record";
let mut rng = CountRng(0x11);
let rec = protect_cbc(&keys, 7, 0x17, TLCP_RECORD_VERSION, pt, &mut rng).unwrap();
assert!(deprotect_cbc(&keys, 8, 0x17, TLCP_RECORD_VERSION, &rec).is_none());
assert!(deprotect_cbc(&keys, 7, 0x16, TLCP_RECORD_VERSION, &rec).is_none());
for i in 0..rec.len() {
let mut t = rec.clone();
t[i] ^= 0x01;
assert!(
deprotect_cbc(&keys, 7, 0x17, TLCP_RECORD_VERSION, &t).is_none(),
"tamper at byte {i} must reject"
);
}
}
#[test]
fn cbc_rejects_malformed_lengths() {
let keys = cbc_keys();
assert!(deprotect_cbc(&keys, 0, 0x17, TLCP_RECORD_VERSION, &[0u8; 63]).is_none());
assert!(deprotect_cbc(&keys, 0, 0x17, TLCP_RECORD_VERSION, &[0u8; 16 + 48 + 1]).is_none());
let huge = alloc::vec![0u8; 16 + MAX_PLAINTEXT + MAC_LEN + 16 + 16];
assert!(deprotect_cbc(&keys, 0, 0x17, TLCP_RECORD_VERSION, &huge).is_none());
}
#[test]
fn pad_check_accepts_and_rejects() {
let mut body = alloc::vec![0u8; 48];
for b in body.iter_mut().skip(32) {
*b = 15;
}
let (ok, pl) = check_tls_padding_ct(&body);
assert!(bool::from(ok));
assert_eq!(pl, 0);
body[40] = 14;
assert!(!bool::from(check_tls_padding_ct(&body).0));
let mut body2 = alloc::vec![0xffu8; 48];
body2[47] = 0xff;
assert!(!bool::from(check_tls_padding_ct(&body2).0));
}
#[test]
fn extract_mac_ct_matches_slice() {
let mut body = alloc::vec![0u8; 100];
for (i, b) in body.iter_mut().enumerate() {
*b = i as u8;
}
for off in [0usize, 1, 33, 68] {
let got = extract_mac_ct(&body, off);
assert_eq!(&got[..], &body[off..off + 32], "offset {off}");
}
}
#[cfg(feature = "sm4-aead")]
#[test]
fn gcm_round_trip_and_context() {
let mut kb = [0u8; 40];
for (i, b) in kb.iter_mut().enumerate() {
*b = (i as u8).wrapping_mul(5).wrapping_add(2);
}
let keys = RecordKeysGcm::client_half(&kb);
for len in [0usize, 1, 16, 100, MAX_PLAINTEXT] {
let pt: Vec<u8> = (0..len).map(|i| i as u8).collect();
let rec = protect_gcm(&keys, 5, 0x17, TLCP_RECORD_VERSION, &pt).unwrap();
assert_eq!(rec.len(), 8 + pt.len() + 16);
assert_eq!(&rec[0..8], &5u64.to_be_bytes());
assert_eq!(
deprotect_gcm(&keys, 5, 0x17, TLCP_RECORD_VERSION, &rec).as_deref(),
Some(&pt[..]),
"gcm round-trip len={len}"
);
}
let rec = protect_gcm(&keys, 5, 0x17, TLCP_RECORD_VERSION, b"x").unwrap();
assert!(deprotect_gcm(&keys, 6, 0x17, TLCP_RECORD_VERSION, &rec).is_none());
}
#[cfg(feature = "sm4-aead")]
#[test]
fn gcm_keys_carve_offsets() {
let mut kb = [0u8; 40];
for (i, b) in kb.iter_mut().enumerate() {
*b = i as u8;
}
let c = RecordKeysGcm::client_half(&kb);
let s = RecordKeysGcm::server_half(&kb);
assert_eq!(c.enc_key, kb[0..16]);
assert_eq!(c.salt, kb[32..36]);
assert_eq!(s.enc_key, kb[16..32]);
assert_eq!(s.salt, kb[36..40]);
}
}