bittensor_core/timelock/
mod.rs

1//! Drand quicknet timelock encryption (absorbed from `bittensor-drand`).
2//!
3//! Lives in the monorepo so the epoch-schedule simulation always matches the
4//! chain: `tle` and `w3f-bls` inherit the workspace pins, so ciphertexts
5//! never drift from what pallet-drand can decrypt.
6
7// Client-side code: arithmetic on locally validated values is the norm here,
8// and this crate never runs inside the runtime.
9#![allow(clippy::arithmetic_side_effects)]
10
11pub mod constants;
12pub mod epoch_schedule;
13#[cfg(test)]
14mod epoch_schedule_vectors;
15
16use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
17use codec::{Decode, Encode};
18use rand_core::{OsRng, RngCore};
19use serde::Deserialize;
20use sha2::Digest;
21// std::time::SystemTime::now() panics on wasm32-unknown-unknown; web-time is
22// a Date.now()-backed drop-in with the same types on that target.
23#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
24use std::time::{SystemTime, UNIX_EPOCH};
25use subtensor_macros::freeze_struct;
26use tle::{
27    curves::drand::TinyBLS381,
28    ibe::fullident::Identity,
29    stream_ciphers::AESGCMStreamCipherProvider,
30    tlock::{tld, tle, TLECiphertext},
31};
32use w3f_bls::EngineBLS;
33#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
34use web_time::{SystemTime, UNIX_EPOCH};
35
36use crate::error::CoreError;
37#[cfg(feature = "host")]
38use constants::{DRAND_ENDPOINTS, QUICKNET_CHAIN_HASH};
39use constants::{DRAND_PERIOD, DRAND_PUBLIC_KEY, GENESIS_TIME, SECURITY_BLOCK_OFFSET};
40
41fn tl_err(msg: impl Into<String>) -> CoreError {
42    CoreError::Crypto(msg.into())
43}
44
45fn now_unix() -> Result<std::time::Duration, CoreError> {
46    SystemTime::now()
47        .duration_since(UNIX_EPOCH)
48        .map_err(|e| tl_err(format!("SystemTime error: {e:?}")))
49}
50
51#[freeze_struct("705b0f6dde3ed6e")]
52#[derive(Encode, Decode, Debug, PartialEq)]
53pub struct WeightsTlockPayload {
54    pub hotkey: Vec<u8>,
55    pub uids: Vec<u16>,
56    pub values: Vec<u16>,
57    pub version_key: u64,
58}
59
60#[freeze_struct("b96b617eb03c11a6")]
61#[derive(Encode, Decode)]
62pub struct UserData {
63    pub encrypted_data: Vec<u8>,
64    pub reveal_round: u64,
65}
66
67fn decode_user_data(encrypted_data: &[u8]) -> Result<UserData, CoreError> {
68    let mut input = encrypted_data;
69    let user_data = UserData::decode(&mut input)
70        .map_err(|e| tl_err(format!("Error deserializing data: {e:?}")))?;
71    if !input.is_empty() {
72        return Err(tl_err("Error deserializing data: trailing bytes"));
73    }
74    Ok(user_data)
75}
76
77#[derive(Deserialize)]
78pub struct DrandResponse {
79    pub round: u64,
80    pub signature: String,
81}
82
83/// The drand reveal round a `UserData` envelope was encrypted to. Portable
84/// counterpart to [`decrypt`]: a host without network access in the core
85/// (e.g. a browser shell) extracts the round here, fetches the signature
86/// itself, and finishes with [`decrypt_with_signature`].
87pub fn reveal_round(encrypted_data: &[u8]) -> Result<u64, CoreError> {
88    let user_data = decode_user_data(encrypted_data)?;
89    Ok(user_data.reveal_round)
90}
91
92/// Inner compressed TLE ciphertext from a `UserData` envelope.
93///
94/// `Commitments.set_commitment` stores this in `TimelockEncrypted.encrypted`
95/// and the reveal round as a separate field. `ciphertext` / `hex()` on the
96/// Python `Timelocked` object is the portable envelope, not this inner blob.
97pub fn inner_ciphertext(encrypted_data: &[u8]) -> Result<Vec<u8>, CoreError> {
98    let user_data = decode_user_data(encrypted_data)?;
99    Ok(user_data.encrypted_data)
100}
101
102/// Timelock-encrypt `serialized_data` to the drand quicknet round
103/// `reveal_round`, returning the compressed TLE ciphertext.
104pub fn encrypt_and_compress(
105    serialized_data: &[u8],
106    reveal_round: u64,
107) -> Result<Vec<u8>, CoreError> {
108    let pub_key_bytes = hex::decode(DRAND_PUBLIC_KEY)
109        .map_err(|e| tl_err(format!("Decoding public key failed: {e:?}")))?;
110    let pub_key =
111        <TinyBLS381 as EngineBLS>::PublicKeyGroup::deserialize_compressed(&*pub_key_bytes)
112            .map_err(|e| tl_err(format!("Deserializing public key failed: {e:?}")))?;
113
114    // Create identity from reveal_round
115    let message = {
116        let mut hasher = sha2::Sha256::new();
117        hasher.update(reveal_round.to_be_bytes());
118        hasher.finalize().to_vec()
119    };
120    let identity = Identity::new(b"", vec![message]);
121
122    // Encrypt payload
123    let mut esk = [0u8; 32];
124    OsRng.fill_bytes(&mut esk);
125    let ct = tle::<TinyBLS381, AESGCMStreamCipherProvider, OsRng>(
126        pub_key,
127        esk,
128        serialized_data,
129        identity,
130        OsRng,
131    )
132    .map_err(|e| tl_err(format!("Encryption failed: {e:?}")))?;
133
134    // Compress ciphertext
135    let mut ct_bytes: Vec<u8> = Vec::new();
136    ct.serialize_compressed(&mut ct_bytes)
137        .map_err(|e| tl_err(format!("Ciphertext serialization failed: {e:?}")))?;
138
139    Ok(ct_bytes)
140}
141
142/// Decrypt a compressed TLE ciphertext with the round's BLS signature.
143pub fn decrypt_and_decompress(
144    encrypted_data: &[u8],
145    signature_bytes: &[u8],
146) -> Result<Vec<u8>, CoreError> {
147    let ciphertext = TLECiphertext::<TinyBLS381>::deserialize_compressed(encrypted_data)
148        .map_err(|e| tl_err(format!("Error deserializing ciphertext: {e:?}")))?;
149
150    let sign = <TinyBLS381 as EngineBLS>::SignatureGroup::deserialize_compressed(signature_bytes)
151        .map_err(|e| tl_err(format!("Signature deserialization error: {e:?}")))?;
152
153    tld::<TinyBLS381, AESGCMStreamCipherProvider>(ciphertext, sign)
154        .map_err(|e| tl_err(format!("Error decrypting ciphertext: {e:?}")))
155}
156
157/// Stateful epoch model commit (v2): simulate the chain's `block_step`
158/// pipeline to find the exact reveal block, then encrypt the weights payload
159/// to the drand round covering it. Returns `(ciphertext, reveal_round)`.
160pub fn generate_commit_v2(
161    uids: Vec<u16>,
162    values: Vec<u16>,
163    version_key: u64,
164    state: epoch_schedule::EpochScheduleState,
165    subnet_reveal_period_epochs: u64,
166    block_time: f64,
167    hotkey: Vec<u8>,
168) -> Result<(Vec<u8>, u64), CoreError> {
169    let first_reveal_blk =
170        epoch_schedule::predict_first_reveal_block(&state, subnet_reveal_period_epochs)
171            .map_err(|e| tl_err(e.to_string()))?;
172
173    let target_ingest_blk = first_reveal_blk.saturating_add(SECURITY_BLOCK_OFFSET);
174    let blocks_until_ingest = target_ingest_blk.saturating_sub(state.current_block);
175    let secs_until_ingest = blocks_until_ingest as f64 * block_time;
176
177    let now_secs = now_unix()?.as_secs_f64();
178
179    let target_secs = now_secs + secs_until_ingest;
180    let mut reveal_round =
181        ((target_secs - GENESIS_TIME as f64) / DRAND_PERIOD as f64).floor() as u64;
182    if reveal_round < 1 {
183        reveal_round = 1;
184    }
185
186    let payload = WeightsTlockPayload {
187        hotkey,
188        uids,
189        values,
190        version_key,
191    };
192
193    let ct_bytes = encrypt_and_compress(&payload.encode(), reveal_round)?;
194    Ok((ct_bytes, reveal_round))
195}
196
197/// Encrypt a string commitment to the round `blocks_until_reveal` blocks from
198/// now. Returns `(ciphertext, reveal_round)`.
199pub fn encrypt_commitment(
200    data: &str,
201    blocks_until_reveal: u64,
202    block_time: f64,
203) -> Result<(Vec<u8>, u64), CoreError> {
204    let serialized_data = data.encode();
205
206    let now = now_unix()?.as_secs();
207    let reveal_round = ((now - GENESIS_TIME)
208        + (blocks_until_reveal as f64 * block_time).round() as u64)
209        / DRAND_PERIOD;
210
211    let ct_bytes = encrypt_and_compress(&serialized_data, reveal_round)?;
212    Ok((ct_bytes, reveal_round))
213}
214
215/// Encrypt binary data to the round `n_blocks` blocks from now, wrapped in
216/// the `UserData` envelope that carries the reveal round.
217pub fn encrypt_n_blocks(
218    data: &[u8],
219    n_blocks: u64,
220    block_time: f64,
221) -> Result<(Vec<u8>, u64), CoreError> {
222    let now = now_unix()?.as_secs_f64();
223    let reveal_timestamp = (n_blocks as f64 * block_time + now).ceil() as u64 - GENESIS_TIME;
224    let reveal_round = reveal_timestamp / DRAND_PERIOD;
225    encrypt_at_round(data, reveal_round)
226}
227
228/// Encrypt binary data to a specific drand reveal round, wrapped in the
229/// `UserData` envelope that carries the reveal round.
230pub fn encrypt_at_round(data: &[u8], reveal_round: u64) -> Result<(Vec<u8>, u64), CoreError> {
231    let encrypted_data = encrypt_and_compress(data, reveal_round)?;
232    let encrypted_with_reveal_round = UserData {
233        encrypted_data,
234        reveal_round,
235    }
236    .encode();
237    Ok((encrypted_with_reveal_round, reveal_round))
238}
239
240/// Fetch drand round info (blocking; tries each public endpoint in order).
241#[cfg(feature = "host")]
242pub fn get_round_info(round: Option<u64>) -> Result<DrandResponse, CoreError> {
243    let mut last_error = None;
244
245    for endpoint in DRAND_ENDPOINTS.iter() {
246        let url = match round {
247            Some(r) => format!("{}/{}/public/{}", endpoint, QUICKNET_CHAIN_HASH, r),
248            None => format!("{}/{}/public/latest", endpoint, QUICKNET_CHAIN_HASH),
249        };
250
251        let response = match reqwest::blocking::get(&url) {
252            Ok(resp) => resp,
253            Err(e) => {
254                last_error = Some(format!("Connection error to {}: {}", endpoint, e));
255                continue;
256            }
257        };
258
259        match response.json::<DrandResponse>() {
260            Ok(parsed) => return Ok(parsed),
261            Err(e) => {
262                last_error = Some(format!("Parsing error from {}: {}", endpoint, e));
263                continue;
264            }
265        }
266    }
267
268    Err(tl_err(last_error.unwrap_or_else(|| {
269        "Failed to get data from all Drand endpoints".to_string()
270    })))
271}
272
273/// The BLS signature (hex) for a reveal round, or `Ok(None)` on fetch errors
274/// when `no_errors` is set.
275#[cfg(feature = "host")]
276pub fn get_reveal_round_signature(
277    reveal_round: Option<u64>,
278    no_errors: bool,
279) -> Result<Option<String>, CoreError> {
280    match get_round_info(reveal_round) {
281        Ok(r) => Ok(Some(r.signature)),
282        Err(e) if no_errors => {
283            let _ = e;
284            Ok(None)
285        }
286        Err(e) => Err(tl_err(format!(
287            "Failed to get Drand round {reveal_round:?}: {e}"
288        ))),
289    }
290}
291
292/// Decrypt a `UserData` envelope end-to-end: extract the reveal round, fetch
293/// the round's BLS signature from drand, and decrypt. With `no_errors`, a
294/// malformed envelope or unavailable round yields `Ok(None)`; decryption
295/// failures against a fetched signature always error (the data is present
296/// but wrong, which the caller should see).
297#[cfg(feature = "host")]
298pub fn decrypt(encrypted_data: &[u8], no_errors: bool) -> Result<Option<Vec<u8>>, CoreError> {
299    let user_data = match decode_user_data(encrypted_data) {
300        Ok(data) => data,
301        Err(_) if no_errors => return Ok(None),
302        Err(e) => return Err(e),
303    };
304
305    let Some(signature_hex) = get_reveal_round_signature(Some(user_data.reveal_round), no_errors)?
306    else {
307        return if no_errors {
308            Ok(None)
309        } else {
310            Err(tl_err("Signature not available"))
311        };
312    };
313
314    decrypt_with_signature_hex(&user_data.encrypted_data, &signature_hex).map(Some)
315}
316
317/// Decrypt a `UserData` envelope with an already-fetched signature. Useful
318/// when decrypting multiple ciphertexts for the same round.
319pub fn decrypt_with_signature(
320    encrypted_data: &[u8],
321    signature_hex: &str,
322) -> Result<Vec<u8>, CoreError> {
323    let user_data = decode_user_data(encrypted_data)?;
324    decrypt_with_signature_hex(&user_data.encrypted_data, signature_hex)
325}
326
327fn decrypt_with_signature_hex(encrypted: &[u8], signature_hex: &str) -> Result<Vec<u8>, CoreError> {
328    let signature_bytes = hex::decode(signature_hex)
329        .map_err(|e| tl_err(format!("Invalid hex in signature: {e:?}")))?;
330    decrypt_and_decompress(encrypted, &signature_bytes)
331}
332
333#[cfg(test)]
334mod tests {
335    #![allow(clippy::unwrap_used, clippy::expect_used)]
336
337    use super::*;
338
339    #[test]
340    fn test_encrypt_and_decrypt_static_key() {
341        let message = b"hello, bittensor!";
342        let reveal_round = 17200000;
343
344        let encrypted =
345            encrypt_and_compress(message, reveal_round).expect("Encryption should succeed");
346
347        let signature_hex = get_reveal_round_signature(Some(reveal_round), false)
348            .expect("Should get signature")
349            .expect("Signature should not be None");
350
351        let signature_bytes = hex::decode(&signature_hex).expect("Hex decoding failed");
352
353        let decrypted = decrypt_and_decompress(&encrypted, &signature_bytes)
354            .expect("Decryption should succeed");
355
356        assert_eq!(message.to_vec(), decrypted);
357    }
358
359    #[test]
360    fn test_get_round_info_and_signature() {
361        let round = 17200000;
362        let info = get_round_info(Some(round)).expect("Drand round should be available");
363
364        assert_eq!(info.round, round);
365        assert!(!info.signature.is_empty());
366
367        let sig = get_reveal_round_signature(Some(round), false).unwrap();
368        assert!(sig.is_some());
369    }
370
371    #[test]
372    fn test_encrypt_commitment_format() {
373        let data = "example string";
374        let (encrypted, round) = encrypt_commitment(data, 10, 12.0).expect("Encryption failed");
375        assert!(!encrypted.is_empty());
376        assert!(round > 0);
377    }
378
379    #[test]
380    fn inner_ciphertext_unwraps_user_data_envelope() {
381        let expected = b"chain submit bytes".to_vec();
382        let mut envelope = UserData {
383            encrypted_data: expected.clone(),
384            reveal_round: 30_000_000,
385        }
386        .encode();
387
388        assert_eq!(inner_ciphertext(&envelope).unwrap(), expected);
389
390        envelope.push(0);
391        assert!(
392            inner_ciphertext(&envelope).is_err(),
393            "trailing bytes must be rejected"
394        );
395    }
396
397    #[test]
398    fn test_generate_commit_v2_structure() {
399        let state = epoch_schedule::EpochScheduleState {
400            last_epoch_block: 100,
401            pending_epoch_at: 0,
402            subnet_epoch_index: 0,
403            tempo: 50,
404            blocks_since_last_step: 0,
405            current_block: 120,
406        };
407        let uids = vec![1, 2, 3];
408        let values = vec![100, 200, 300];
409        let hotkey = vec![1, 2, 3];
410
411        let (encrypted, reveal_round) = generate_commit_v2(
412            uids.clone(),
413            values.clone(),
414            42,
415            state,
416            1,
417            12.0,
418            hotkey.clone(),
419        )
420        .expect("Commit generation failed");
421
422        assert!(!encrypted.is_empty());
423        assert!(reveal_round > 0);
424    }
425}
bittensor_core/timelock/mod.rs

bittensor_core/timelock/
mod.rs
