solana_zk_token_sdk/encryption/
discrete_log.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
//! The discrete log implementation for the twisted ElGamal decryption.
//!
//! The implementation uses the baby-step giant-step method, which consists of a precomputation
//! step and an online step. The precomputation step involves computing a hash table of a number
//! of Ristretto points that is independent of a discrete log instance. The online phase computes
//! the final discrete log solution using the discrete log instance and the pre-computed hash
//! table. More details on the baby-step giant-step algorithm and the implementation can be found
//! in the [spl documentation](https://spl.solana.com).
//!
//! The implementation is NOT intended to run in constant-time. There are some measures to prevent
//! straightforward timing attacks. For instance, it does not short-circuit the search when a
//! solution is found. However, the use of hashtables, batching, and threads make the
//! implementation inherently not constant-time. This may theoretically allow an adversary to gain
//! information on a discrete log solution depending on the execution time of the implementation.
//!

#![cfg(not(target_os = "solana"))]

#[cfg(not(target_arch = "wasm32"))]
use std::thread;
use {
    crate::RISTRETTO_POINT_LEN,
    curve25519_dalek::{
        constants::RISTRETTO_BASEPOINT_POINT as G,
        ristretto::RistrettoPoint,
        scalar::Scalar,
        traits::{Identity, IsIdentity},
    },
    itertools::Itertools,
    serde::{Deserialize, Serialize},
    std::{collections::HashMap, num::NonZeroUsize},
    thiserror::Error,
};

const TWO16: u64 = 65536; // 2^16
const TWO17: u64 = 131072; // 2^17

/// Maximum number of threads permitted for discrete log computation
#[cfg(not(target_arch = "wasm32"))]
const MAX_THREAD: usize = 65536;

#[derive(Error, Clone, Debug, Eq, PartialEq)]
pub enum DiscreteLogError {
    #[error("discrete log number of threads not power-of-two")]
    DiscreteLogThreads,
    #[error("discrete log batch size too large")]
    DiscreteLogBatchSize,
}

/// Type that captures a discrete log challenge.
///
/// The goal of discrete log is to find x such that x * generator = target.
#[derive(Serialize, Deserialize, Copy, Clone, Debug, Eq, PartialEq)]
pub struct DiscreteLog {
    /// Generator point for discrete log
    pub generator: RistrettoPoint,
    /// Target point for discrete log
    pub target: RistrettoPoint,
    /// Number of threads used for discrete log computation
    num_threads: Option<NonZeroUsize>,
    /// Range bound for discrete log search derived from the max value to search for and
    /// `num_threads`
    range_bound: NonZeroUsize,
    /// Ristretto point representing each step of the discrete log search
    step_point: RistrettoPoint,
    /// Ristretto point compression batch size
    compression_batch_size: NonZeroUsize,
}

#[derive(Serialize, Deserialize, Default)]
pub struct DecodePrecomputation(HashMap<[u8; RISTRETTO_POINT_LEN], u16>);

/// Builds a HashMap of 2^16 elements
#[allow(dead_code)]
fn decode_u32_precomputation(generator: RistrettoPoint) -> DecodePrecomputation {
    let mut hashmap = HashMap::new();

    let two17_scalar = Scalar::from(TWO17);
    let identity = RistrettoPoint::identity(); // 0 * G
    let generator = two17_scalar * generator; // 2^17 * G

    // iterator for 2^17*0G , 2^17*1G, 2^17*2G, ...
    let ristretto_iter = RistrettoIterator::new((identity, 0), (generator, 1));
    for (point, x_hi) in ristretto_iter.take(TWO16 as usize) {
        let key = point.compress().to_bytes();
        hashmap.insert(key, x_hi as u16);
    }

    DecodePrecomputation(hashmap)
}

lazy_static::lazy_static! {
    /// Pre-computed HashMap needed for decryption. The HashMap is independent of (works for) any key.
    pub static ref DECODE_PRECOMPUTATION_FOR_G: DecodePrecomputation = {
        static DECODE_PRECOMPUTATION_FOR_G_BINCODE: &[u8] =
            include_bytes!("decode_u32_precomputation_for_G.bincode");
        bincode::deserialize(DECODE_PRECOMPUTATION_FOR_G_BINCODE).unwrap_or_default()
    };
}

/// Solves the discrete log instance using a 16/16 bit offline/online split
impl DiscreteLog {
    /// Discrete log instance constructor.
    ///
    /// Default number of threads set to 1.
    pub fn new(generator: RistrettoPoint, target: RistrettoPoint) -> Self {
        Self {
            generator,
            target,
            num_threads: None,
            range_bound: (TWO16 as usize).try_into().unwrap(),
            step_point: G,
            compression_batch_size: 32.try_into().unwrap(),
        }
    }

    /// Adjusts number of threads in a discrete log instance.
    #[cfg(not(target_arch = "wasm32"))]
    pub fn num_threads(&mut self, num_threads: NonZeroUsize) -> Result<(), DiscreteLogError> {
        // number of threads must be a positive power-of-two integer
        if !num_threads.is_power_of_two() || num_threads.get() > MAX_THREAD {
            return Err(DiscreteLogError::DiscreteLogThreads);
        }

        self.num_threads = Some(num_threads);
        self.range_bound = (TWO16 as usize)
            .checked_div(num_threads.get())
            .and_then(|range_bound| range_bound.try_into().ok())
            .unwrap(); // `num_threads` cannot exceed `TWO16`, so `range_bound` always non-zero
        self.step_point = Scalar::from(num_threads.get() as u64) * G;

        Ok(())
    }

    /// Adjusts inversion batch size in a discrete log instance.
    pub fn set_compression_batch_size(
        &mut self,
        compression_batch_size: NonZeroUsize,
    ) -> Result<(), DiscreteLogError> {
        if compression_batch_size.get() >= TWO16 as usize {
            return Err(DiscreteLogError::DiscreteLogBatchSize);
        }
        self.compression_batch_size = compression_batch_size;

        Ok(())
    }

    /// Solves the discrete log problem under the assumption that the solution
    /// is a positive 32-bit number.
    pub fn decode_u32(self) -> Option<u64> {
        if let Some(num_threads) = self.num_threads {
            #[cfg(not(target_arch = "wasm32"))]
            {
                let mut starting_point = self.target;
                let handles = (0..num_threads.get())
                    .map(|i| {
                        let ristretto_iterator = RistrettoIterator::new(
                            (starting_point, i as u64),
                            (-(&self.step_point), num_threads.get() as u64),
                        );

                        let handle = thread::spawn(move || {
                            Self::decode_range(
                                ristretto_iterator,
                                self.range_bound,
                                self.compression_batch_size,
                            )
                        });

                        starting_point -= G;
                        handle
                    })
                    .collect::<Vec<_>>();

                handles
                    .into_iter()
                    .map_while(|h| h.join().ok())
                    .find(|x| x.is_some())
                    .flatten()
            }
            #[cfg(target_arch = "wasm32")]
            unreachable!() // `self.num_threads` always `None` on wasm target
        } else {
            let ristretto_iterator =
                RistrettoIterator::new((self.target, 0_u64), (-(&self.step_point), 1u64));

            Self::decode_range(
                ristretto_iterator,
                self.range_bound,
                self.compression_batch_size,
            )
        }
    }

    fn decode_range(
        ristretto_iterator: RistrettoIterator,
        range_bound: NonZeroUsize,
        compression_batch_size: NonZeroUsize,
    ) -> Option<u64> {
        let hashmap = &DECODE_PRECOMPUTATION_FOR_G;
        let mut decoded = None;

        for batch in &ristretto_iterator
            .take(range_bound.get())
            .chunks(compression_batch_size.get())
        {
            // batch compression currently errors if any point in the batch is the identity point
            let (batch_points, batch_indices): (Vec<_>, Vec<_>) = batch
                .filter(|(point, index)| {
                    if point.is_identity() {
                        decoded = Some(*index);
                        return false;
                    }
                    true
                })
                .unzip();

            let batch_compressed = RistrettoPoint::double_and_compress_batch(&batch_points);

            for (point, x_lo) in batch_compressed.iter().zip(batch_indices.iter()) {
                let key = point.to_bytes();
                if hashmap.0.contains_key(&key) {
                    let x_hi = hashmap.0[&key];
                    decoded = Some(x_lo + TWO16 * x_hi as u64);
                }
            }
        }

        decoded
    }
}

/// Hashable Ristretto iterator.
///
/// Given an initial point X and a stepping point P, the iterator iterates through
/// X + 0*P, X + 1*P, X + 2*P, X + 3*P, ...
struct RistrettoIterator {
    pub current: (RistrettoPoint, u64),
    pub step: (RistrettoPoint, u64),
}

impl RistrettoIterator {
    fn new(current: (RistrettoPoint, u64), step: (RistrettoPoint, u64)) -> Self {
        RistrettoIterator { current, step }
    }
}

impl Iterator for RistrettoIterator {
    type Item = (RistrettoPoint, u64);

    fn next(&mut self) -> Option<Self::Item> {
        let r = self.current;
        self.current = (self.current.0 + self.step.0, self.current.1 + self.step.1);
        Some(r)
    }
}

#[cfg(test)]
mod tests {
    use {super::*, std::time::Instant};

    #[test]
    #[allow(non_snake_case)]
    fn test_serialize_decode_u32_precomputation_for_G() {
        let decode_u32_precomputation_for_G = decode_u32_precomputation(G);
        // let decode_u32_precomputation_for_G = decode_u32_precomputation(G);

        if decode_u32_precomputation_for_G.0 != DECODE_PRECOMPUTATION_FOR_G.0 {
            use std::{fs::File, io::Write, path::PathBuf};
            let mut f = File::create(PathBuf::from(
                "src/encryption/decode_u32_precomputation_for_G.bincode",
            ))
            .unwrap();
            f.write_all(&bincode::serialize(&decode_u32_precomputation_for_G).unwrap())
                .unwrap();
            panic!("Rebuild and run this test again");
        }
    }

    #[test]
    fn test_decode_correctness() {
        // general case
        let amount: u64 = 4294967295;

        let instance = DiscreteLog::new(G, Scalar::from(amount) * G);

        // Very informal measurements for now
        let start_computation = Instant::now();
        let decoded = instance.decode_u32();
        let computation_secs = start_computation.elapsed().as_secs_f64();

        assert_eq!(amount, decoded.unwrap());

        println!("single thread discrete log computation secs: {computation_secs:?} sec");
    }

    #[cfg(not(target_arch = "wasm32"))]
    #[test]
    fn test_decode_correctness_threaded() {
        // general case
        let amount: u64 = 55;

        let mut instance = DiscreteLog::new(G, Scalar::from(amount) * G);
        instance.num_threads(4.try_into().unwrap()).unwrap();

        // Very informal measurements for now
        let start_computation = Instant::now();
        let decoded = instance.decode_u32();
        let computation_secs = start_computation.elapsed().as_secs_f64();

        assert_eq!(amount, decoded.unwrap());

        println!("4 thread discrete log computation: {computation_secs:?} sec");

        // amount 0
        let amount: u64 = 0;

        let instance = DiscreteLog::new(G, Scalar::from(amount) * G);

        let decoded = instance.decode_u32();
        assert_eq!(amount, decoded.unwrap());

        // amount 1
        let amount: u64 = 1;

        let instance = DiscreteLog::new(G, Scalar::from(amount) * G);

        let decoded = instance.decode_u32();
        assert_eq!(amount, decoded.unwrap());

        // amount 2
        let amount: u64 = 2;

        let instance = DiscreteLog::new(G, Scalar::from(amount) * G);

        let decoded = instance.decode_u32();
        assert_eq!(amount, decoded.unwrap());

        // amount 3
        let amount: u64 = 3;

        let instance = DiscreteLog::new(G, Scalar::from(amount) * G);

        let decoded = instance.decode_u32();
        assert_eq!(amount, decoded.unwrap());

        // max amount
        let amount: u64 = (1_u64 << 32) - 1;

        let instance = DiscreteLog::new(G, Scalar::from(amount) * G);

        let decoded = instance.decode_u32();
        assert_eq!(amount, decoded.unwrap());
    }
}