1use blake3::Hasher;
39use serde::{Deserialize, Serialize};
40use thiserror::Error;
41
42#[derive(Error, Debug)]
44pub enum VdfError {
45 #[error("Invalid proof")]
46 InvalidProof,
47 #[error("Invalid parameters")]
48 InvalidParameters,
49 #[error("Iteration count must be positive")]
50 InvalidIterations,
51}
52
53pub type VdfResult<T> = Result<T, VdfError>;
54
55#[derive(Clone, Debug, Serialize, Deserialize)]
57pub struct VdfParams {
58 pub iterations: u64,
60}
61
62impl VdfParams {
63 pub fn new(iterations: u64) -> Self {
79 assert!(iterations > 0);
80 Self { iterations }
81 }
82
83 pub fn from_duration_ms(target_ms: u64) -> Self {
92 const ITERATIONS_PER_MS: u64 = 10_000;
93 Self {
94 iterations: target_ms * ITERATIONS_PER_MS,
95 }
96 }
97}
98
99#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
101pub struct VdfOutput {
102 pub value: Vec<u8>,
104}
105
106#[derive(Clone, Debug, Serialize, Deserialize)]
111pub struct VdfProof {
112 checkpoints: Vec<Vec<u8>>,
114 output: Vec<u8>,
116}
117
118pub fn vdf_compute(params: &VdfParams, input: &[u8]) -> (VdfOutput, VdfProof) {
137 let mut current = hash_input(input);
138 let mut checkpoints = Vec::new();
139
140 let checkpoint_interval = (params.iterations / 10).max(1);
143
144 for i in 0..params.iterations {
145 current = hash_step(¤t);
146
147 if i > 0 && (i + 1) % checkpoint_interval == 0 {
149 checkpoints.push(current.clone());
150 }
151 }
152
153 let output = VdfOutput {
154 value: current.clone(),
155 };
156
157 let proof = VdfProof {
158 checkpoints,
159 output: current,
160 };
161
162 (output, proof)
163}
164
165pub fn vdf_verify(
189 params: &VdfParams,
190 input: &[u8],
191 output: &VdfOutput,
192 proof: &VdfProof,
193) -> VdfResult<()> {
194 if output.value != proof.output {
196 return Err(VdfError::InvalidProof);
197 }
198
199 let mut current = hash_input(input);
201 let checkpoint_interval = (params.iterations / 10).max(1);
202 let mut checkpoint_idx = 0;
203
204 for i in 0..params.iterations {
205 current = hash_step(¤t);
206
207 if i > 0 && (i + 1) % checkpoint_interval == 0 {
209 if checkpoint_idx >= proof.checkpoints.len() {
210 return Err(VdfError::InvalidProof);
211 }
212
213 if current != proof.checkpoints[checkpoint_idx] {
214 return Err(VdfError::InvalidProof);
215 }
216
217 checkpoint_idx += 1;
218 }
219 }
220
221 if current != proof.output {
223 return Err(VdfError::InvalidProof);
224 }
225
226 Ok(())
227}
228
229pub fn vdf_randomness_beacon(seed: &[u8], iterations: u64) -> Vec<u8> {
239 let params = VdfParams::new(iterations);
240 let (output, _proof) = vdf_compute(¶ms, seed);
241 output.value
242}
243
244fn hash_input(input: &[u8]) -> Vec<u8> {
246 let mut hasher = Hasher::new();
247 hasher.update(input);
248 hasher.finalize().as_bytes().to_vec()
249}
250
251fn hash_step(data: &[u8]) -> Vec<u8> {
253 let mut hasher = Hasher::new();
254 hasher.update(data);
255 hasher.finalize().as_bytes().to_vec()
256}
257
258#[cfg(test)]
259mod tests {
260 use super::*;
261
262 #[test]
263 fn test_vdf_basic() {
264 let params = VdfParams::new(100);
265 let input = b"test_input";
266
267 let (output, proof) = vdf_compute(¶ms, input);
268
269 assert!(vdf_verify(¶ms, input, &output, &proof).is_ok());
270 }
271
272 #[test]
273 fn test_vdf_deterministic() {
274 let params = VdfParams::new(100);
275 let input = b"test_input";
276
277 let (output1, _) = vdf_compute(¶ms, input);
278 let (output2, _) = vdf_compute(¶ms, input);
279
280 assert_eq!(output1.value, output2.value);
281 }
282
283 #[test]
284 fn test_vdf_different_inputs() {
285 let params = VdfParams::new(100);
286
287 let (output1, _) = vdf_compute(¶ms, b"input1");
288 let (output2, _) = vdf_compute(¶ms, b"input2");
289
290 assert_ne!(output1.value, output2.value);
291 }
292
293 #[test]
294 fn test_vdf_different_iterations() {
295 let input = b"test_input";
296
297 let params1 = VdfParams::new(100);
298 let params2 = VdfParams::new(200);
299
300 let (output1, _) = vdf_compute(¶ms1, input);
301 let (output2, _) = vdf_compute(¶ms2, input);
302
303 assert_ne!(output1.value, output2.value);
304 }
305
306 #[test]
307 fn test_vdf_invalid_proof() {
308 let params = VdfParams::new(100);
309 let input = b"test_input";
310
311 let (output, mut proof) = vdf_compute(¶ms, input);
312
313 if !proof.checkpoints.is_empty() {
315 proof.checkpoints[0][0] ^= 1;
316 }
317
318 assert!(vdf_verify(¶ms, input, &output, &proof).is_err());
319 }
320
321 #[test]
322 fn test_vdf_serialization() {
323 let params = VdfParams::new(100);
324 let input = b"test_input";
325
326 let (output, proof) = vdf_compute(¶ms, input);
327
328 let params_bytes = crate::codec::encode(¶ms).unwrap();
330 let output_bytes = crate::codec::encode(&output).unwrap();
331 let proof_bytes = crate::codec::encode(&proof).unwrap();
332
333 let params2: VdfParams = crate::codec::decode(¶ms_bytes).unwrap();
335 let output2: VdfOutput = crate::codec::decode(&output_bytes).unwrap();
336 let proof2: VdfProof = crate::codec::decode(&proof_bytes).unwrap();
337
338 assert!(vdf_verify(¶ms2, input, &output2, &proof2).is_ok());
339 }
340
341 #[test]
342 fn test_vdf_from_duration() {
343 let params = VdfParams::from_duration_ms(10);
344 assert_eq!(params.iterations, 100_000);
345
346 let params2 = VdfParams::from_duration_ms(100);
347 assert_eq!(params2.iterations, 1_000_000);
348 }
349
350 #[test]
351 fn test_vdf_randomness_beacon() {
352 let seed = b"beacon_seed";
353 let output1 = vdf_randomness_beacon(seed, 1000);
354 let output2 = vdf_randomness_beacon(seed, 1000);
355
356 assert_eq!(output1, output2);
358
359 let output3 = vdf_randomness_beacon(b"different_seed", 1000);
361 assert_ne!(output1, output3);
362 }
363
364 #[test]
365 fn test_vdf_output_length() {
366 let params = VdfParams::new(100);
367 let input = b"test";
368
369 let (output, _) = vdf_compute(¶ms, input);
370
371 assert_eq!(output.value.len(), 32);
373 }
374
375 #[test]
376 fn test_vdf_large_iterations() {
377 let params = VdfParams::new(10_000);
379 let input = b"large_test";
380
381 let (output, proof) = vdf_compute(¶ms, input);
382
383 assert!(vdf_verify(¶ms, input, &output, &proof).is_ok());
384 assert!(!proof.checkpoints.is_empty());
385 }
386}