Skip to content

Commit

Permalink
refactoring szk
Browse files Browse the repository at this point in the history
  • Loading branch information
Hannah Davis committed Oct 30, 2024
1 parent 5a46f32 commit 6c6e284
Show file tree
Hide file tree
Showing 2 changed files with 159 additions and 149 deletions.
247 changes: 142 additions & 105 deletions src/flp/szk.rs
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,11 @@ pub enum SzkError {
/// can be properly read upon receipt.
MissingVerifierLen(()),

#[error("Part of Szk query state not stored")]
/// Returned when a user fails to store the length of the verifier so it
/// can be properly read upon receipt.
InvalidState(String),

/// Returned if an FLP operation encountered an error.
#[error("Flp error: {0}")]
Flp(#[from] FlpError),
Expand Down Expand Up @@ -255,85 +260,66 @@ impl<F: FieldElement + Decode, const SEED_SIZE: usize> ParameterizedDecode<(bool
}
}

impl<F: FieldElement, const SEED_SIZE: usize> SzkQueryShare<F, SEED_SIZE> {
pub(crate) fn merge_verifiers(
mut leader_share: SzkQueryShare<F, SEED_SIZE>,
helper_share: SzkQueryShare<F, SEED_SIZE>,
) -> SzkVerifier<F, SEED_SIZE> {
for (x, y) in leader_share
.flp_verifier
.iter_mut()
.zip(helper_share.flp_verifier)
{
*x += y;
}
SzkVerifier {
flp_verifier: leader_share.flp_verifier,
leader_joint_rand_part_opt: leader_share.joint_rand_part_opt,
helper_joint_rand_part_opt: helper_share.joint_rand_part_opt,
}
}
}

/// Szk query state.
///
/// The state that needs to be stored by an Szk verifier between query() and decide().
pub type SzkQueryState<const SEED_SIZE: usize> = Option<Seed<SEED_SIZE>>;

/// Verifier type for the SZK proof.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct SzkVerifier<F: FieldElement, const SEED_SIZE: usize> {
flp_verifier: Vec<F>,
leader_joint_rand_part_opt: Option<Seed<SEED_SIZE>>,
helper_joint_rand_part_opt: Option<Seed<SEED_SIZE>>,
pub type SzkVerifier<F> = Vec<F>;

impl<F: FieldElement> Encode for SzkVerifier<F> {
fn encode(&self, bytes: &mut Vec<u8>) -> Result<(), CodecError> {
encode_fieldvec(&self, bytes)?;
Ok(())
}

fn encoded_len(&self) -> Option<usize> {
Some(self.len() * F::ENCODED_SIZE)
}
}

impl<F: FieldElement + Decode> ParameterizedDecode<usize> for SzkVerifier<F> {
fn decode_with_param(
verifier_len: &usize,
bytes: &mut Cursor<&[u8]>,
) -> Result<Self, CodecError> {
decode_fieldvec(*verifier_len, bytes)
}
}

impl<F: FieldElement, const SEED_SIZE: usize> Encode for SzkVerifier<F, SEED_SIZE> {
/// Joint share type for the SZK proof.
pub type SzkJointShare<const SEED_SIZE: usize> = Option<[Seed<SEED_SIZE>; 2]>;

impl<const SEED_SIZE: usize> Encode for SzkJointShare<SEED_SIZE> {
fn encode(&self, bytes: &mut Vec<u8>) -> Result<(), CodecError> {
encode_fieldvec(&self.flp_verifier, bytes)?;
if let Some(ref part) = self.leader_joint_rand_part_opt {
part.encode(bytes)?
};
if let Some(ref part) = self.helper_joint_rand_part_opt {
part.encode(bytes)?
if let Some([ref leader_part, ref helper_part]) = self {
leader_part.encode(bytes)?;
helper_part.encode(bytes)?;
};
Ok(())
}

fn encoded_len(&self) -> Option<usize> {
Some(
self.flp_verifier.len() * F::ENCODED_SIZE
+ match self.leader_joint_rand_part_opt {
Some(ref part) => part.encoded_len()?,
None => 0,
}
+ match self.helper_joint_rand_part_opt {
Some(ref part) => part.encoded_len()?,
None => 0,
},
)
Some(match self {
Some(ref part) => part[0].encoded_len()? * 2,
None => 0,
})
}
}

impl<F: FieldElement + Decode, const SEED_SIZE: usize> ParameterizedDecode<(bool, usize)>
for SzkVerifier<F, SEED_SIZE>
{
impl<const SEED_SIZE: usize> ParameterizedDecode<bool> for SzkJointShare<SEED_SIZE> {
fn decode_with_param(
(requires_joint_rand, verifier_len): &(bool, usize),
requires_joint_rand: &bool,
bytes: &mut Cursor<&[u8]>,
) -> Result<Self, CodecError> {
if *requires_joint_rand {
Ok(SzkVerifier {
flp_verifier: decode_fieldvec(*verifier_len, bytes)?,
leader_joint_rand_part_opt: Some(Seed::<SEED_SIZE>::decode(bytes)?),
helper_joint_rand_part_opt: Some(Seed::<SEED_SIZE>::decode(bytes)?),
})
Ok(Some([
Seed::<SEED_SIZE>::decode(bytes)?,
Seed::<SEED_SIZE>::decode(bytes)?,
]))
} else {
Ok(SzkVerifier {
flp_verifier: decode_fieldvec(*verifier_len, bytes)?,
leader_joint_rand_part_opt: None,
helper_joint_rand_part_opt: None,
})
Ok(None)
}
}
}
Expand Down Expand Up @@ -629,34 +615,56 @@ where
))
}

/// Returns true if the verifier message indicates that the input from which
/// it was generated is valid.
pub(crate) fn merge_verifiers(
&self,
mut leader_share: SzkQueryShare<T::Field, SEED_SIZE>,
helper_share: SzkQueryShare<T::Field, SEED_SIZE>,
) -> Result<SzkJointShare<SEED_SIZE>, SzkError> {
for (x, y) in leader_share
.flp_verifier
.iter_mut()
.zip(helper_share.flp_verifier)
{
*x += y;
}
if self.typ.decide(&leader_share.flp_verifier)? {
match (
leader_share.joint_rand_part_opt,
helper_share.joint_rand_part_opt,
) {
(Some(leader_part), Some(helper_part)) => Ok(Some([leader_part, helper_part])),
(None, None) => Ok(None),
_ => Err(SzkError::InvalidState(
"at least one of the joint randomness parts is missing".to_string(),
)),
}
} else {
Err(SzkError::Decide("failed to verify FLP proof".to_string()))
}
}
/// Returns true if the leader and helper derive identical joint randomness
/// seeds
pub fn decide(
&self,
verifier: SzkVerifier<T::Field, SEED_SIZE>,
query_state: SzkQueryState<SEED_SIZE>,
) -> Result<bool, SzkError> {
// Check if underlying FLP proof validates
let check_flp_proof = self.typ.decide(&verifier.flp_verifier)?;
if !check_flp_proof {
return Ok(false);
}
joint_share: &[Seed<SEED_SIZE>; 2],
) -> Result<(), SzkError> {
// Check that joint randomness was properly derived from both
// aggregators' parts
match (
query_state,
verifier.leader_joint_rand_part_opt,
verifier.helper_joint_rand_part_opt,
) {
(Some(joint_rand_seed), Some(leader_joint_rand_part), Some(helper_joint_rand_part)) => {
match (query_state, joint_share) {
(Some(joint_rand_seed), [leader_joint_rand_part, helper_joint_rand_part]) => {
let expected_joint_rand_seed =
self.derive_joint_rand_seed(&leader_joint_rand_part, &helper_joint_rand_part);
Ok(joint_rand_seed == expected_joint_rand_seed)
self.derive_joint_rand_seed(leader_joint_rand_part, helper_joint_rand_part);
if joint_rand_seed == expected_joint_rand_seed {
Ok(())
} else {
Err(SzkError::Decide(
"Aggregators failed to compute identical joint randomness seeds"
.to_string(),
))
}
}
(None, None, None) => Ok(true),
(_, _, _) => Err(SzkError::Decide(
"at least one of the input seeds is missing".to_string(),
)),
(None, _) => Ok(()),
}
}
}
Expand Down Expand Up @@ -744,29 +752,40 @@ mod tests {
.query(&helper_input_share, h_proof_share, &verify_key, &nonce)
.unwrap();

let verifier = SzkQueryShare::merge_verifiers(l_query_share.clone(), h_query_share.clone());
if let Ok(leader_decision) = szk_typ.decide(verifier.clone(), l_query_state.clone()) {
assert_eq!(
leader_decision, valid,
"Leader incorrectly determined validity",
);
} else {
panic!("Leader failed during decision");
};
if let Ok(helper_decision) = szk_typ.decide(verifier.clone(), h_query_state.clone()) {
assert_eq!(
helper_decision, valid,
"Helper incorrectly determined validity",
);
} else {
panic!("Helper failed during decision");
let joint_share_result =
szk_typ.merge_verifiers(l_query_share.clone(), h_query_share.clone());
match joint_share_result {
Ok(Some(ref joint_share)) => {
let leader_decision = match szk_typ.decide(l_query_state.clone(), joint_share) {
Ok(_) => true,
Err(_) => false,
};
assert_eq!(
leader_decision, valid,
"Leader incorrectly determined validity",
);
let helper_decision = match szk_typ.decide(h_query_state.clone(), joint_share) {
Ok(_) => true,
Err(_) => false,
};
assert_eq!(
helper_decision, valid,
"Helper incorrectly determined validity",
);
}
Ok(None) => assert_eq!(true, valid, "Aggregator incorrectly determined validity"),
Err(_) => {
assert_eq!(false, valid, "Aggregator incorrectly determined validity");
}
};

//test mutated jr seed
if szk_typ.requires_joint_rand() {
let joint_rand_seed_opt = Some(Seed::<16>::generate().unwrap());
if let Ok(leader_decision) = szk_typ.decide(verifier, joint_rand_seed_opt.clone()) {
assert!(!leader_decision, "Leader accepted wrong jr seed");
if let Ok(Some(ref joint_share)) = joint_share_result {
if let Ok(()) = szk_typ.decide(joint_rand_seed_opt.clone(), joint_share) {
panic!("Leader accepted wrong jr seed");
};
};
};

Expand All @@ -778,9 +797,15 @@ mod tests {
);
}

let verifier = SzkQueryShare::merge_verifiers(mutated_query_share, h_query_share.clone());

let leader_decision = szk_typ.decide(verifier, l_query_state.clone()).unwrap();
let joint_share_res = szk_typ.merge_verifiers(mutated_query_share, h_query_share.clone());
let leader_decision = match joint_share_res {
Ok(Some(ref joint_share)) => match szk_typ.decide(l_query_state.clone(), joint_share) {
Ok(_) => true,
Err(_) => false,
},
Ok(None) => true,
Err(_) => false,
};
assert!(!leader_decision, "Leader validated after proof mutation");

// test mutated input share
Expand All @@ -791,11 +816,17 @@ mod tests {
.query(&mutated_input, l_proof_share.clone(), &verify_key, &nonce)
.unwrap();

let verifier = SzkQueryShare::merge_verifiers(mutated_query_share, h_query_share.clone());
let joint_share_res = szk_typ.merge_verifiers(mutated_query_share, h_query_share.clone());

if let Ok(leader_decision) = szk_typ.decide(verifier, mutated_query_state) {
assert!(!leader_decision, "Leader validated after input mutation");
let leader_decision = match joint_share_res {
Ok(Some(ref joint_share)) => match szk_typ.decide(mutated_query_state, joint_share) {
Ok(_) => true,
Err(_) => false,
},
Ok(None) => true,
Err(_) => false,
};
assert!(!leader_decision, "Leader validated after input mutation");

// test mutated proof share
let (mut mutated_proof, leader_blind_and_helper_joint_rand_part_opt) = match l_proof_share {
Expand All @@ -822,11 +853,17 @@ mod tests {
&nonce,
)
.unwrap();
let verifier = SzkQueryShare::merge_verifiers(l_query_share, h_query_share.clone());
let joint_share_res = szk_typ.merge_verifiers(l_query_share, h_query_share.clone());

if let Ok(leader_decision) = szk_typ.decide(verifier, l_query_state) {
assert!(!leader_decision, "Leader validated after proof mutation");
let leader_decision = match joint_share_res {
Ok(Some(ref joint_share)) => match szk_typ.decide(l_query_state.clone(), joint_share) {
Ok(_) => true,
Err(_) => false,
},
Ok(None) => true,
Err(_) => false,
};
assert!(!leader_decision, "Leader validated after proof mutation");
}

#[test]
Expand Down
Loading

0 comments on commit 6c6e284

Please sign in to comment.