tlsn_core/transcript/encoding/
encoder.rs

1use crate::transcript::{Direction, Idx, Subsequence};
2use itybity::ToBits;
3use rand::{RngCore, SeedableRng};
4use rand_chacha::ChaCha12Rng;
5use serde::{Deserialize, Serialize};
6
7/// The size of the encoding for 1 bit, in bytes.
8const BIT_ENCODING_SIZE: usize = 16;
9/// The size of the encoding for 1 byte, in bytes.
10const BYTE_ENCODING_SIZE: usize = 128;
11
12/// Secret used by an encoder to generate encodings.
13#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)]
14pub struct EncoderSecret {
15    seed: [u8; 32],
16    delta: [u8; BIT_ENCODING_SIZE],
17}
18
19opaque_debug::implement!(EncoderSecret);
20
21impl EncoderSecret {
22    /// Creates a new secret.
23    ///
24    /// # Arguments
25    ///
26    /// * `seed` - The seed for the PRG.
27    /// * `delta` - Delta for deriving the one-encodings.
28    pub fn new(seed: [u8; 32], delta: [u8; 16]) -> Self {
29        Self { seed, delta }
30    }
31
32    /// Returns the seed.
33    pub fn seed(&self) -> &[u8; 32] {
34        &self.seed
35    }
36
37    /// Returns the delta.
38    pub fn delta(&self) -> &[u8; 16] {
39        &self.delta
40    }
41}
42
43/// Creates a new encoder.
44pub fn new_encoder(secret: &EncoderSecret) -> impl Encoder {
45    ChaChaEncoder::new(secret)
46}
47
48pub(crate) struct ChaChaEncoder {
49    seed: [u8; 32],
50    delta: [u8; 16],
51}
52
53impl ChaChaEncoder {
54    pub(crate) fn new(secret: &EncoderSecret) -> Self {
55        let seed = *secret.seed();
56        let delta = *secret.delta();
57
58        Self { seed, delta }
59    }
60
61    pub(crate) fn new_prg(&self, stream_id: u64) -> ChaCha12Rng {
62        let mut prg = ChaCha12Rng::from_seed(self.seed);
63        prg.set_stream(stream_id);
64        prg.set_word_pos(0);
65        prg
66    }
67}
68
69/// A transcript encoder.
70///
71/// This is an internal implementation detail that should not be exposed to the
72/// public API.
73pub trait Encoder {
74    /// Returns the zero encoding for the given index.
75    fn encode_idx(&self, direction: Direction, idx: &Idx) -> Vec<u8>;
76
77    /// Returns the encoding for the given subsequence of the transcript.
78    ///
79    /// # Arguments
80    ///
81    /// * `seq` - The subsequence to encode.
82    fn encode_subsequence(&self, direction: Direction, seq: &Subsequence) -> Vec<u8>;
83}
84
85impl Encoder for ChaChaEncoder {
86    fn encode_idx(&self, direction: Direction, idx: &Idx) -> Vec<u8> {
87        // ChaCha encoder works with 32-bit words. Each encoded bit is 128 bits long.
88        const WORDS_PER_BYTE: u128 = 8 * 128 / 32;
89
90        let stream_id: u64 = match direction {
91            Direction::Sent => 0,
92            Direction::Received => 1,
93        };
94
95        let mut prg = self.new_prg(stream_id);
96        let mut encoding: Vec<u8> = vec![0u8; idx.len() * BYTE_ENCODING_SIZE];
97
98        let mut pos = 0;
99        for range in idx.iter_ranges() {
100            let len = range.len() * BYTE_ENCODING_SIZE;
101            prg.set_word_pos(range.start as u128 * WORDS_PER_BYTE);
102            prg.fill_bytes(&mut encoding[pos..pos + len]);
103            pos += len;
104        }
105
106        encoding
107    }
108
109    fn encode_subsequence(&self, direction: Direction, seq: &Subsequence) -> Vec<u8> {
110        const ZERO: [u8; 16] = [0; BIT_ENCODING_SIZE];
111
112        let mut encoding = self.encode_idx(direction, seq.index());
113        for (pos, bit) in seq.data().iter_lsb0().enumerate() {
114            // Add the delta to the encoding whenever the encoded bit is 1,
115            // otherwise add a zero.
116            let summand = if bit { &self.delta } else { &ZERO };
117            encoding[pos * BIT_ENCODING_SIZE..(pos + 1) * BIT_ENCODING_SIZE]
118                .iter_mut()
119                .zip(summand)
120                .for_each(|(a, b)| *a ^= *b);
121        }
122
123        encoding
124    }
125}