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

#![cfg(not(target_arch = "bpf"))]

use {
    curve25519_dalek::{ristretto::RistrettoPoint, scalar::Scalar, traits::Identity},
    serde::{Deserialize, Serialize},
    std::collections::HashMap,
};

#[allow(dead_code)]
const TWO15: u64 = 32768;
#[allow(dead_code)]
const TWO14: u64 = 16384; // 2^14
const TWO16: u64 = 65536; // 2^16
#[allow(dead_code)]
const TWO18: u64 = 262144; // 2^18

/// 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,
}

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

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

    let two16_scalar = Scalar::from(TWO16);
    let identity = RistrettoPoint::identity(); // 0 * G
    let generator = two16_scalar * generator; // 2^16 * G

    // iterator for 2^12*0G , 2^12*1G, 2^12*2G, ...
    let ristretto_iter = RistrettoIterator::new(identity, generator);
    ristretto_iter.zip(0..TWO16).for_each(|(elem, x_hi)| {
        let key = elem.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 {
    /// Solves the discrete log problem under the assumption that the solution
    /// is a 32-bit number.
    pub(crate) fn decode_u32(self) -> Option<u64> {
        self.decode_online(&DECODE_PRECOMPUTATION_FOR_G, TWO16)
    }

    pub fn decode_online(self, hashmap: &DecodePrecomputation, solution_bound: u64) -> Option<u64> {
        // iterator for 0G, -1G, -2G, ...
        let ristretto_iter = RistrettoIterator::new(self.target, -self.generator);

        let mut decoded = None;
        ristretto_iter
            .zip(0..solution_bound)
            .for_each(|(elem, x_lo)| {
                let key = elem.compress().to_bytes();
                if hashmap.0.contains_key(&key) {
                    let x_hi = hashmap.0[&key];
                    decoded = Some(x_lo + solution_bound * x_hi as u64);
                }
            });
        decoded
    }
}

/// HashableRistretto 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 curr: RistrettoPoint,
    pub step: RistrettoPoint,
}

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

impl Iterator for RistrettoIterator {
    type Item = RistrettoPoint;

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

#[cfg(test)]
mod tests {
    use {
        super::*, curve25519_dalek::constants::RISTRETTO_BASEPOINT_POINT as G, 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);

        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() {
        let amount: u64 = 65545;

        let instance = DiscreteLog {
            generator: G,
            target: Scalar::from(amount) * G,
        };

        // Very informal measurements for now
        let start_precomputation = Instant::now();
        let precomputed_hashmap = decode_u32_precomputation(G);
        let precomputation_secs = start_precomputation.elapsed().as_secs_f64();

        let start_online = Instant::now();
        let computed_amount = instance.decode_online(&precomputed_hashmap, TWO16).unwrap();
        let online_secs = start_online.elapsed().as_secs_f64();

        assert_eq!(amount, computed_amount);

        println!("16/16 Split precomputation: {:?} sec", precomputation_secs);
        println!("16/16 Split online computation: {:?} sec", online_secs);
    }
}