use std::io::{Read, Write};
use chacha20poly1305::{aead::KeyInit, ChaCha20Poly1305, Key};
use hkdf::Hkdf;
use sha2::Sha256;
use zeroize::Zeroize;
use crate::decrypt::decrypt_v3_chunks;
use crate::encrypt::encrypt_chunks;
use crate::error::PqfileError;
use crate::format::{
commitment_for_tlock, version_layout, PqfHeader, PqfHeaderTlock, BASE_NONCE_LEN, CHUNK_SIZE,
NONCE_LEN, TLOCK_CHAIN_HASH_LEN, VERSION_AUTH_BIT, VERSION_TLOCK,
};
use crate::secret::LockedSecret;
pub struct TlockChain {
pub hash: [u8; TLOCK_CHAIN_HASH_LEN],
pub public_key: &'static [u8],
pub default_relay: &'static str,
}
pub fn quicknet() -> TlockChain {
TlockChain {
hash: QUICKNET_HASH,
public_key: &QUICKNET_PUBLIC_KEY,
default_relay:
"https://api.drand.sh/52db9ba70e0cc0f6eaf7803dd07447a1f5477735fd3f661792ba94600c84e971",
}
}
const QUICKNET_HASH: [u8; TLOCK_CHAIN_HASH_LEN] = [
0x52, 0xdb, 0x9b, 0xa7, 0x0e, 0x0c, 0xc0, 0xf6, 0xea, 0xf7, 0x80, 0x3d, 0xd0, 0x74, 0x47, 0xa1,
0xf5, 0x47, 0x77, 0x35, 0xfd, 0x3f, 0x66, 0x17, 0x92, 0xba, 0x94, 0x60, 0x0c, 0x84, 0xe9, 0x71,
];
const QUICKNET_PUBLIC_KEY: [u8; 96] = [
0x83, 0xcf, 0x0f, 0x28, 0x96, 0xad, 0xee, 0x7e, 0xb8, 0xb5, 0xf0, 0x1f, 0xca, 0xd3, 0x91, 0x22,
0x12, 0xc4, 0x37, 0xe0, 0x07, 0x3e, 0x91, 0x1f, 0xb9, 0x00, 0x22, 0xd3, 0xe7, 0x60, 0x18, 0x3c,
0x8c, 0x4b, 0x45, 0x0b, 0x6a, 0x0a, 0x6c, 0x3a, 0xc6, 0xa5, 0x77, 0x6a, 0x2d, 0x10, 0x64, 0x51,
0x0d, 0x1f, 0xec, 0x75, 0x8c, 0x92, 0x1c, 0xc2, 0x2b, 0x0e, 0x17, 0xe6, 0x3a, 0xaf, 0x4b, 0xcb,
0x5e, 0xd6, 0x63, 0x04, 0xde, 0x9c, 0xf8, 0x09, 0xbd, 0x27, 0x4c, 0xa7, 0x3b, 0xab, 0x4a, 0xf5,
0xa6, 0xe9, 0xc7, 0x6a, 0x4b, 0xc0, 0x9e, 0x76, 0xea, 0xe8, 0x99, 0x1e, 0xf5, 0xec, 0xe4, 0x5a,
];
const TLOCK_HKDF_INFO: &[u8] = b"pqfile-tlock-v1";
fn derive_session_key(seed: &[u8; 16]) -> Result<LockedSecret<32>, PqfileError> {
let hk = Hkdf::<Sha256>::new(None, seed);
let mut okm = LockedSecret::<32>::zeroed();
hk.expand(TLOCK_HKDF_INFO, okm.as_mut())
.map_err(|_| PqfileError::EncryptionFailure)?;
Ok(okm)
}
pub fn encrypt_stream_tlock(
round: u64,
chain: Option<&TlockChain>,
original_size: u64,
reader: &mut dyn Read,
writer: &mut dyn Write,
) -> Result<(), PqfileError> {
let owned_chain;
let chain = match chain {
Some(c) => c,
None => {
owned_chain = quicknet();
&owned_chain
}
};
let mut seed = [0u8; 16];
loop {
getrandom::fill(&mut seed).map_err(|_| PqfileError::EncryptionFailure)?;
if seed[15] != 0 {
break;
}
}
let mut tlock_ct = Vec::new();
let encrypt_result = tlock::encrypt(&mut tlock_ct, seed.as_slice(), chain.public_key, round);
if encrypt_result.is_err() {
seed.zeroize();
return Err(PqfileError::EncryptionFailure);
}
let session_key = derive_session_key(&seed);
seed.zeroize();
let session_key = session_key?;
let mut nonce_bytes = [0u8; NONCE_LEN];
getrandom::fill(&mut nonce_bytes[..BASE_NONCE_LEN])
.map_err(|_| PqfileError::EncryptionFailure)?;
let header = PqfHeaderTlock {
chain_hash: chain.hash,
round,
tlock_ct,
nonce: nonce_bytes,
original_size,
};
let version = VERSION_TLOCK | VERSION_AUTH_BIT;
header.write(writer, version).map_err(PqfileError::Io)?;
let key_commitment = commitment_for_tlock(session_key.as_ref(), &header);
let base_nonce: &[u8; BASE_NONCE_LEN] = header.nonce[..BASE_NONCE_LEN]
.try_into()
.expect("BASE_NONCE_LEN <= NONCE_LEN");
let key: &Key = session_key.as_ref().try_into().expect("32-byte key");
let cipher = ChaCha20Poly1305::new(key);
encrypt_chunks(
&cipher,
base_nonce,
CHUNK_SIZE,
&key_commitment,
reader,
writer,
)
}
pub fn round_for_target_time(
when: &str,
chain: Option<&TlockChain>,
relay_url: Option<&str>,
) -> Result<u64, PqfileError> {
let owned_chain;
let chain = match chain {
Some(c) => c,
None => {
owned_chain = quicknet();
&owned_chain
}
};
let relay = relay_url.unwrap_or(chain.default_relay);
let client = drand_core::HttpClient::new(relay, None)
.map_err(|e| PqfileError::TlockBeaconFetchFailed(e.to_string()))?;
let info = client
.chain_info()
.map_err(|e| PqfileError::TlockBeaconFetchFailed(e.to_string()))?;
let time_info: drand_core::chain::ChainTimeInfo = info.into();
let parsed = drand_core::beacon::RandomnessBeaconTime::parse(&time_info, when)
.map_err(|e| PqfileError::TlockBeaconFetchFailed(e.to_string()))?;
Ok(parsed.round())
}
pub fn decrypt_stream_tlock(
relay_url: Option<&str>,
reader: &mut dyn Read,
writer: &mut dyn Write,
) -> Result<(), PqfileError> {
let version = PqfHeader::read_magic_version(reader)?;
if version_layout(version) != VERSION_TLOCK {
return Err(PqfileError::UnsupportedVersion(version));
}
let header = PqfHeaderTlock::read_body(reader)?;
let relay = match relay_url {
Some(u) => u,
None if header.chain_hash == QUICKNET_HASH => quicknet().default_relay,
None => {
return Err(PqfileError::TlockBeaconFetchFailed(
"file uses a non-default drand chain; pass an explicit relay URL".to_owned(),
))
}
};
let client = drand_core::HttpClient::new(relay, None)
.map_err(|e| PqfileError::TlockBeaconFetchFailed(e.to_string()))?;
let info = client
.chain_info()
.map_err(|e| PqfileError::TlockBeaconFetchFailed(e.to_string()))?;
if info.hash() != header.chain_hash {
let got: String = info.hash().iter().map(|b| format!("{b:02x}")).collect();
let want: String = header
.chain_hash
.iter()
.map(|b| format!("{b:02x}"))
.collect();
return Err(PqfileError::TlockBeaconFetchFailed(format!(
"relay {relay} serves chain {got}, but this file uses chain {want}"
)));
}
let beacon = match client.get(header.round) {
Ok(beacon) => beacon,
Err(drand_core::DrandError::Beacon(b))
if matches!(*b, drand_core::beacon::BeaconError::NotFound) =>
{
return Err(PqfileError::TlockRoundNotReached {
round: header.round,
})
}
Err(e) => {
let msg = e.to_string();
if msg.contains("http status: 404") || msg.contains("http status: 425") {
return Err(PqfileError::TlockRoundNotReached {
round: header.round,
});
}
return Err(PqfileError::TlockBeaconFetchFailed(msg));
}
};
let signature = beacon.signature();
let decrypt_result = std::panic::catch_unwind(|| {
let mut pt = Vec::new();
tlock::decrypt(&mut pt, header.tlock_ct.as_slice(), signature.as_slice()).map(|_| pt)
});
let mut plaintext = match decrypt_result {
Ok(Ok(pt)) if pt.len() == 16 => pt,
_ => return Err(PqfileError::TlockDecryptionFailed),
};
let mut seed = [0u8; 16];
seed.copy_from_slice(&plaintext);
plaintext.zeroize();
let session_key = derive_session_key(&seed);
seed.zeroize();
let session_key = session_key?;
let key_commitment = commitment_for_tlock(session_key.as_ref(), &header);
let key: &Key = session_key.as_ref().try_into().expect("32-byte key");
let cipher = ChaCha20Poly1305::new(key);
decrypt_v3_chunks(
&cipher,
&header.nonce,
CHUNK_SIZE,
&key_commitment,
reader,
writer,
)
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Cursor;
const ROUND_1000_SIGNATURE_HEX: &str =
"b44679b9a59af2ec876b1a6b1ad52ea9b1615fc3982b19576350f93447cb1125e342b73a8dd2bacbe47e4b6b63ed5e39";
fn decode_hex(s: &str) -> Vec<u8> {
(0..s.len())
.step_by(2)
.map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap())
.collect()
}
#[test]
fn tlock_ibe_roundtrip_against_real_quicknet_round() {
let signature = decode_hex(ROUND_1000_SIGNATURE_HEX);
let mut seed = [7u8; 16];
seed[15] = 1;
let mut ct = Vec::new();
tlock::encrypt(&mut ct, seed.as_slice(), &QUICKNET_PUBLIC_KEY, 1000).unwrap();
let mut pt = Vec::new();
tlock::decrypt(&mut pt, ct.as_slice(), signature.as_slice()).unwrap();
assert_eq!(pt, seed);
}
#[test]
fn full_stream_roundtrip_with_pinned_signature() {
let signature = decode_hex(ROUND_1000_SIGNATURE_HEX);
let plaintext = b"time-locked message, 5 bytes over a chunk boundary!";
let mut ciphertext = Vec::new();
encrypt_stream_tlock(
1000,
None,
plaintext.len() as u64,
&mut Cursor::new(plaintext.as_slice()),
&mut ciphertext,
)
.unwrap();
let mut reader = Cursor::new(ciphertext.as_slice());
let version = PqfHeader::read_magic_version(&mut reader).unwrap();
assert_eq!(version_layout(version), VERSION_TLOCK);
let header = PqfHeaderTlock::read_body(&mut reader).unwrap();
assert_eq!(header.round, 1000);
assert_eq!(header.chain_hash, QUICKNET_HASH);
let mut pt = Vec::new();
tlock::decrypt(&mut pt, header.tlock_ct.as_slice(), signature.as_slice()).unwrap();
let mut seed = [0u8; 16];
seed.copy_from_slice(&pt);
let session_key = derive_session_key(&seed).unwrap();
let key_commitment = commitment_for_tlock(session_key.as_ref(), &header);
let key: &Key = session_key.as_ref().try_into().unwrap();
let cipher = ChaCha20Poly1305::new(key);
let mut out = Vec::new();
decrypt_v3_chunks(
&cipher,
&header.nonce,
CHUNK_SIZE,
&key_commitment,
&mut reader,
&mut out,
)
.unwrap();
assert_eq!(out, plaintext);
}
#[test]
fn wrong_signature_is_a_clean_error_not_a_panic() {
let mut seed = [9u8; 16];
seed[15] = 1;
let mut ct = Vec::new();
tlock::encrypt(&mut ct, seed.as_slice(), &QUICKNET_PUBLIC_KEY, 1000).unwrap();
let wrong_signature = decode_hex(ROUND_1000_SIGNATURE_HEX);
let mut corrupted_ct = ct.clone();
corrupted_ct[10] ^= 0xff;
let result = std::panic::catch_unwind(|| {
let mut pt = Vec::new();
tlock::decrypt(&mut pt, corrupted_ct.as_slice(), wrong_signature.as_slice()).map(|_| pt)
});
let _ = result;
}
#[test]
#[ignore = "hits the network"]
fn decrypt_stream_tlock_against_live_relay_historical_round() {
let plaintext = b"integration test payload, network-verified";
let mut ciphertext = Vec::new();
encrypt_stream_tlock(
1000,
None,
plaintext.len() as u64,
&mut Cursor::new(plaintext.as_slice()),
&mut ciphertext,
)
.unwrap();
let mut out = Vec::new();
decrypt_stream_tlock(None, &mut Cursor::new(ciphertext.as_slice()), &mut out).unwrap();
assert_eq!(out, plaintext);
}
#[test]
#[ignore = "hits the network"]
fn decrypt_stream_tlock_future_round_is_not_reached() {
let plaintext = b"not decryptable yet";
let mut ciphertext = Vec::new();
let future_round = 200_000_000;
encrypt_stream_tlock(
future_round,
None,
plaintext.len() as u64,
&mut Cursor::new(plaintext.as_slice()),
&mut ciphertext,
)
.unwrap();
let mut out = Vec::new();
let err = decrypt_stream_tlock(None, &mut Cursor::new(ciphertext.as_slice()), &mut out)
.unwrap_err();
assert!(matches!(
err,
PqfileError::TlockRoundNotReached { round } if round == future_round
));
}
#[test]
fn seed_resampling_avoids_trailing_zero_byte() {
for _ in 0..1000 {
let mut seed = [0u8; 16];
loop {
getrandom::fill(&mut seed).unwrap();
if seed[15] != 0 {
break;
}
}
assert_ne!(seed[15], 0);
}
}
}