Small updates to the VDF + add benchmarks (#839)

* Small updates + add benchmarks * Optimize verification * clean up * test * align * rev * rev * simplify
MystenLabs · Oct 31, 2024 · d5b0059 · d5b0059
1 parent d838f9f
commit d5b0059
Show file tree

Hide file tree

Showing 7 changed files with 172 additions and 50 deletions.
diff --git a/fastcrypto-vdf/Cargo.toml b/fastcrypto-vdf/Cargo.toml
@@ -44,3 +44,8 @@ required-features = ["experimental"]
 name = "math"
 harness = false
 required-features = ["experimental"]
+
+[[bench]]
+name = "vdf_rsa"
+harness = false
+required-features = ["experimental"]
diff --git a/fastcrypto-vdf/benches/vdf_rsa.rs b/fastcrypto-vdf/benches/vdf_rsa.rs
diff --git a/fastcrypto-vdf/src/math/parameterized_group.rs b/fastcrypto-vdf/src/math/parameterized_group.rs
@@ -19,31 +19,25 @@ pub trait ParameterizedGroupElement:
 
     /// Returns true if this is an element of the group defined by `parameter`.
     fn is_in_group(&self, parameter: &Self::ParameterType) -> bool;
-}
 
-/// Compute self * scalar using a "Double-and-Add" algorithm for a positive scalar.
-pub(crate) fn multiply<G: ParameterizedGroupElement>(
-    input: &G,
-    scalar: &BigUint,
-    parameter: &G::ParameterType,
-) -> G {
-    (0..scalar.bits())
-        .rev()
-        .map(|i| scalar.bit(i))
-        .fold(G::zero(parameter), |acc, bit| {
-            let mut res = acc.double();
-            if bit {
-                res = res + input;
-            }
-            res
-        })
+    /// Compute self * scalar.
+    fn multiply(&self, scalar: &BigUint, parameter: &Self::ParameterType) -> Self {
+        // Generic double-and-add algorithm.
+        (0..scalar.bits())
+            .rev()
+            .map(|i| scalar.bit(i))
+            .fold(Self::zero(parameter), |acc, bit| match bit {
+                true => acc.double() + self,
+                false => acc.double(),
+            })
+    }
 }
 
 #[cfg(test)]
 mod tests {
     use crate::class_group::discriminant::Discriminant;
     use crate::class_group::QuadraticForm;
-    use crate::math::parameterized_group::{multiply, ParameterizedGroupElement};
+    use crate::math::parameterized_group::ParameterizedGroupElement;
     use num_bigint::BigUint;
     use num_traits::{One, Zero};
 
@@ -55,12 +49,12 @@ mod tests {
         // Edge cases
         assert_eq!(
             QuadraticForm::zero(&discriminant),
-            multiply(&input, &BigUint::zero(), &discriminant)
+            input.multiply(&BigUint::zero(), &discriminant)
         );
-        assert_eq!(input, multiply(&input, &BigUint::one(), &discriminant));
+        assert_eq!(input, input.multiply(&BigUint::one(), &discriminant));
 
         let exponent = 12345u64;
-        let output = multiply(&input, &BigUint::from(exponent), &discriminant);
+        let output = input.multiply(&BigUint::from(exponent), &discriminant);
 
         // Check alignment with repeated addition.
         let mut expected_output = input.clone();

diff --git a/fastcrypto-vdf/src/rsa_group/mod.rs b/fastcrypto-vdf/src/rsa_group/mod.rs
@@ -65,12 +65,19 @@ impl Doubling for RSAGroupElement<'_> {
 impl<'a> ParameterizedGroupElement for RSAGroupElement<'a> {
     type ParameterType = &'a RSAModulus;
 
-    fn zero(parameter: &Self::ParameterType) -> Self {
-        Self::new(BigUint::one(), parameter)
+    fn zero(modulus: &Self::ParameterType) -> Self {
+        Self::new(BigUint::one(), modulus)
     }
 
-    fn is_in_group(&self, parameter: &Self::ParameterType) -> bool {
-        &self.modulus == parameter
+    fn is_in_group(&self, modulus: &Self::ParameterType) -> bool {
+        &self.modulus == modulus
+    }
+
+    fn multiply(&self, scalar: &BigUint, modulus: &Self::ParameterType) -> Self {
+        let value = self
+            .modulus
+            .ensure_in_subgroup(self.value.modpow(scalar, &modulus.value));
+        Self { value, modulus }
     }
 }
 

diff --git a/fastcrypto-vdf/src/rsa_group/modulus.rs b/fastcrypto-vdf/src/rsa_group/modulus.rs
@@ -43,7 +43,12 @@ impl RSAModulus {
     /// modulus, the result is negated. This is to ensure that the result is in the subgroup
     /// <i>Z<sub>N</sub><sup>*</sup> / <±1></i>.
     pub(super) fn reduce(&self, value: BigUint) -> BigUint {
-        let value = value.mod_floor(&self.value);
+        self.ensure_in_subgroup(value.mod_floor(&self.value))
+    }
+
+    /// Assuming that `value < N`, this ensures that the given value is in the subgroup <i>Z<sub>N</sub><sup>*</sup> / <±1></i>.
+    /// Panics if `value` is greater than or equal to `N`.
+    pub(super) fn ensure_in_subgroup(&self, value: BigUint) -> BigUint {
         if value < self.half {
             value
         } else {

diff --git a/fastcrypto-vdf/src/vdf/pietrzak.rs b/fastcrypto-vdf/src/vdf/pietrzak.rs
@@ -1,7 +1,7 @@
 // Copyright (c) 2022, Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::math::parameterized_group::{multiply, ParameterizedGroupElement};
+use crate::math::parameterized_group::ParameterizedGroupElement;
 use crate::vdf::VDF;
 use fastcrypto::error::FastCryptoError::{InvalidInput, InvalidProof};
 use fastcrypto::error::FastCryptoResult;
@@ -39,6 +39,14 @@ impl<G: ParameterizedGroupElement> PietrzaksVDF<G> {
             iterations,
         }
     }
+
+    /// Compute the Fiat-Shamir challenge used in Pietrzak's VDF construction.
+    fn compute_challenge(&self, input: &G, output: &G, mu: &G) -> BigUint {
+        let seed = bcs::to_bytes(&(input, output, mu, self.iterations, &self.group_parameter))
+            .expect("Failed to serialize Fiat-Shamir input.");
+        let hash = Keccak256::digest(seed);
+        BigUint::from_bytes_be(&hash.digest[..DEFAULT_CHALLENGE_SIZE_IN_BYTES])
+    }
 }
 
 impl<G: ParameterizedGroupElement + Serialize> VDF for PietrzaksVDF<G>
@@ -75,9 +83,9 @@ where
             // TODO: Precompute some of the mu's to speed up the proof generation.
             let mu = x.clone().repeated_doubling(t);
 
-            let r = compute_challenge(self, &x, &y, &mu);
-            x = multiply(&x, &r, &self.group_parameter) + &mu;
-            y = multiply(&mu, &r, &self.group_parameter) + &y;
+            let r = self.compute_challenge(&x, &y, &mu);
+            x = x.multiply(&r, &self.group_parameter) + &mu;
+            y = mu.multiply(&r, &self.group_parameter) + &y;
 
             proof.push(mu);
         }
@@ -103,9 +111,9 @@ where
                 y = y.double();
             }
 
-            let r = compute_challenge(self, &x, &y, mu);
-            x = multiply(&x, &r, &self.group_parameter) + mu;
-            y = multiply(mu, &r, &self.group_parameter) + &y;
+            let r = self.compute_challenge(&x, &y, mu);
+            x = x.multiply(&r, &self.group_parameter) + mu;
+            y = mu.multiply(&r, &self.group_parameter) + &y;
         }
 
         // In case the proof is shorter than the full proof, we need to compute the remaining powers.
@@ -117,22 +125,6 @@ where
     }
 }
 
-/// Compute the Fiat-Shamir challenge used in Pietrzak's VDF construction.
-fn compute_challenge<G: ParameterizedGroupElement + Serialize>(
-    vdf: &PietrzaksVDF<G>,
-    input: &G,
-    output: &G,
-    mu: &G,
-) -> BigUint
-where
-    G::ParameterType: Serialize,
-{
-    let seed = bcs::to_bytes(&(input, output, mu, vdf.iterations, &vdf.group_parameter))
-        .expect("Failed to serialize Fiat-Shamir input.");
-    let hash = Keccak256::digest(seed);
-    BigUint::from_bytes_be(&hash.digest[..DEFAULT_CHALLENGE_SIZE_IN_BYTES])
-}
-
 /// Replace t with (t+1) >> 1 and return true iff the input was odd.
 #[inline]
 fn check_parity_and_iterate(t: &mut u64) -> bool {
@@ -210,4 +202,35 @@ mod tests {
         let other_input = input.clone() + &input;
         assert!(vdf.verify(&other_input, &output, &proof).is_err())
     }
+
+    #[test]
+    fn test_verification() {
+        let iterations = 100000u64;
+        let input = RSAGroupElement::new(7u32.into(), &AMAZON_MODULUS_2048_REF);
+
+        let vdf = PietrzaksVDF::<RSAGroupElement>::new(&AMAZON_MODULUS_2048_REF, iterations);
+
+        let output = RSAGroupElement::new(BigUint::from_str("3672355472927810456151679905868758794616388736583581652860721896616180315835575064423750778001349052163223791724800901623993790866712822248888484109754966734883928871058715188132045817948670396453684148159917970871906760229102881786328076485709065174288114937597861379962478300572943579355762549162779462210911294105646062939394772010450396339819711345421231813212879300843731930189807951083800491228370195479515271316034750609611188992819886615677507435486955879522713591072027715880883277172276263534234336298488388979051136904224916555055212735687426081668066225539132052845808964107807201586990131791093492371778").unwrap(), &AMAZON_MODULUS_2048_REF);
+
+        let proof = [
+            "8719512874474216460687150647680854955292132854750703068763484269093165999505251186898178062066004619862656663698343691868497980570426526019975578885850544581914946790597872355640934341428693413899572424721269873153938666911042969304059487718891954826837874698630809943131078862483827048179876513916989352726065292297373409109550091793480245914861995645289031849847427508439025383190674698027602276320860090412338328667933755723078371787557008469872399255533817861369300159517919410447709106340745759633527098227506660977309887993315579818822200330785614717258672305343338813672124794370455365197734851856001184976030",
+            "3524225440490227716581914005817764728689326138271987326546913156143949550779591139824077984832090908071998951770389914336769029333480915517009260387548811330930648406468414527213930607990485031826572327430297154420012068864491234371049362834108095572222288210563942740991735388228265830602511240768621994740770804049029617345757918078538264560135467586163671018157956690861770799025782115742787596459481739552535516586086172662006852110953782305125365361525831420213691071299126328246388152487905764478231871643596200406647908089848722382856021765764700103197413553278821980238755080812642569566460155874926864973989",
+            "3386256193912258600563321427374151160437651417107230398140990906215918611413870201888792916418033967177817103485384285178003637404550898822951512439384277984377627754055120144092077652950815810643804322665300251229793337536560764764926079235179412052806380151642397223785653772910005773851385366606281262555049596002588808028004178177665489626495310464581979168814142288652114023204247090962814602454684901841935354676441552550033218538383127851117213323460171046993929640395141844677358751651632616897741890084456824284415415554894221269513288351780048483154293635621337002324127284236944745716110705339322771478131",
+            "7408862030172429723032582356161847253198188005112855906630189356951227804097029509883935224836192563731148726408119632983128842196882987765881545177086653651757096748219643205448581241037817385331996915432821304066168098773880393211961153863820790217510052941214079784321117971743731404301810060365970610316645229069960509651124240339464525485721413226027035555493674864907371842778371294951677138477875018003840619969373647380251579931187688167577599725809029577381715849590344663076500835477018110706747067929066540698745588647653032284281207590297280558283161226254697184936417182114406611858688290199092351638794",
+            "9528369401155898421336374761623167878924922005826918572675787143246237849874471257635361123764028975937343311664804485861820128360684057779979271237362810683850151327885936654765160960323604793157027745850425197516858623567687919335408856735970509721863383148065881142539924459849066413081803234738981126021705399792445956657733940194160711501942192074012722676182134009812129315502787932498635715069416080651060699212375799959775531304461678248368699251392095272154244316553547920213132393799101764442827812390852292543047588306534171173257087058146331742057376971999785759427057334778797976189560957105013350467045",
+            "10294651820160651728456829898879555551196961487064613372692384294265418044982827373843041565455655933992371447144565864904575509809947198474416922744319637148661233405028052214366165153902932698351161834805653679280884257744536629934561559541496076130305993849014886971786115597309448966484871145453391677209950550556478590364500734905015127553627540473598051645020805146472132102067032709164655196448902087849028550166457630712049225557448992309930729811971031335828818305551184854572895352878770266426086789864481188000057417074198845278355152274308525042623428164257444044541498663412953148100768983136146188235997",
+            "4749045743059502299076848945196073558631005868948441651589845799068150488425734086270947139878995514214119726260212551705844889280663704807069073071616386852792450852711805042766312070760756053408646662762397826414400752299162726413950336169847086249885638830927861421398231965683427732018263765712940081366775804116587569845878336332182016655558080273929449303602571070901792076802687535487545915059452072480809069254144639332974988342423715575563084819806637919328178725980692796438546565820486572953340526856360623620668058599315200343597721637054937740264892162125653205249722406415399818735024248737314552063747",
+            "5272521386133313093121504847176339485318640410159988398538154570040609938127193342269575965623490130675739152927549782661334170941321238880903335951305493326431599940919086128527082527057864468821840814550477028653687293814925163394101119890565259089917449040659282833402394902253943586909550853333057424633788010291693428637999909856531013688038817805974292565379653351375127609695388247301423539646327950220115990830645404216008980887159586046198784837543462911777611082256896302074042733126733093536925128289699426032159423325697676408264099787619827716147575504711932293307925317181532838376975174334760490178044",
+            "4865631439244694792682051927120887855014714782764996443468316201280149186411764462638919743896836216210355086464033599674608823616620181993224084354774021985050767398675323086421509833372146310317782744044433481505781466337868841734616907727794600499486966898550073615373789361818554381065300597950458611369588509114065985025852518182253888929859936361242940138578344904177539257408094849163580060796966127477986366148739231408126796157723280304456556233704922267395211269025280937735721607598808102862937894271692476411316343801775630519501390182293503303960113390510808164020764496445889549138522442194944316766219",
+            "1716098664111733751302684239386687925838384836544658006905034083331006633019856429890234964042806996064770725585636147008586818637936193658888381443567507219576974926569248997425135873112782252084087049174247625853371226689113885574956468036468201035128573779522685641470930864834546438686526548212910642194495508777532033343509629089340426924730626446347003141889733967730229296355986139573616824252745173021579889146711264881818395175840508044230472692250294317790079274412733463156608459907600873015033659104883567469349509785038318299646299550636867596927464451110109337227918075217459535287195599185046432272198",
+            "10721991997697793738717595864735689654814408802179139713299139161130174918125420293856907422920539738936394204606043102026032040519309913895955997218242869889330722353056722164513913325053855254230359609203831468726499782865513152890774696492202216752929686242664105376380907060826609793521588890690791106895587503800467949772773581224066322672834898563688334924748703360423621150393100717027630221055080843102194314393166621862146840652586817200363151726362689280490364249590829799246039599450021900326960690007908969307676046114319938297872389335199854504572884094360772400490069773329668111595259339308363041063871",
+            "24993235836866849416196592570323660996972817499564456104037943507194360405061754866850348770077420997103335102890585709469767629001870643341451985396909081660774052156782419424390497656420915820543525425352103071684707145165891840374067169328373245883329511678687043434601293171271199312025423435502927076645446502230342559583292008660994947133886690268696014405994231692380704451596507223379197197420921324939035607218784585824666269355224282438899491842717180848598449014025167402088262687787425520237057062755968280168112911893559397514451243586114205429244216713755580836976353622394774491637290756440823473904",
+            "8002840044191636897394105768503100351367843810943426955203265943939407779211249672257698205017741561535958104095665612435684465282015904737730202451805337613184916942787797112106584192538882516595181912717026458237097485563704041509438199717078206419533982984474171725446518666727110378813664918107052405840880667070832701817128246672829292017261411874510037160227026367399797925223118665562035040322045995481973053566119964317652423365276712206105413327459764222767039305354559649984488929528403367306919207004930556248185192286629700484211799294085032088235292426962942458433119325773711776898892739174573133259798",
+            "9808298167953463044183224853747590991821290083082693475897746710109152766350649057730805604162385834890955802232144653208502844156265486182999277061377895494700822405534823371717810993077165144659899399251399946686003871122626066126316460159851320953097822716750103981636398570978788394229208496720512000415568014823831712088816947453635369935021640264506413372998966855334669695718142871353958255047552753123177745051748277174077245602488815883354127415057464229321675973890199622571893243720584859039940876060951672389949282570769956932017208645843036419518163184395801601473092002267026597399366492737904207164145",
+            "3419924672544574100524058325383025198496207713265005632195290353279586550544804988055757778145112066368161851153318915869887514873926726636081263451040999914232751332583833616365834650019287835558155301861869825336660948178770848366160254240651864317486675126510537151881923158433631746791994990810114742904187471698817430077934405459479950737131082312281050974043218768012562886561272629075590257984018309718254526858184923594606358659772686305061082946653881595244349982087597745499473392706220050012667178661314295322165215791098424265370584312926915581831293094295870670813932186950843650836176864709988630778898",
+            "3876660104277748734969970799207654963546831221524157837790450138448447235917188249028042741771865887141807354675325586321042930971396579383599429668029844089074571818157381649964860225310811444870630848131459324585689921588479979399255647304741995870412178808125039448834571923055165839896461006202917186695764779234307600813983399230786701019298692652551177614160338386117023117927173943068934500519930664358506164490744248358331236180729050425334573543140408803841919596166155205764351816894658586553341628917380296648889364729838239067182538618852656851850409265336348330079350182770839095482851902292535994651519"
+        ].into_iter().map(|s| RSAGroupElement::new(BigUint::from_str(s).unwrap(), &AMAZON_MODULUS_2048_REF)).collect::<Vec<_>>();
+
+        assert!(vdf.verify(&input, &output, &proof).is_ok());
+    }
 }
diff --git a/fastcrypto-vdf/src/vdf/wesolowski.rs b/fastcrypto-vdf/src/vdf/wesolowski.rs
@@ -81,7 +81,7 @@ impl<G: ParameterizedGroupElement, M: ScalarMultiplier<G, BigUint>> VDF for Weso
         // Algorithm from page 3 on https://crypto.stanford.edu/~dabo/pubs/papers/VDFsurvey.pdf
         let challenge = (self.challenge)(self, input, &output);
         let mut quotient_remainder = (BigUint::from(0u8), BigUint::from(2u8));
-        let mut proof = multiplier.mul(&quotient_remainder.0);
+        let mut proof = G::zero(&self.group_parameter);
         for _ in 1..self.iterations {
             quotient_remainder.1.shl_assign(1);
             quotient_remainder = quotient_remainder.1.div_mod_floor(&challenge);