Schnorr Non-interactive Zero-Knowledge Proof in Rust

zkp.rs

//! Unevaluated prototype of IETF RFC 8235 §3.3 with §4.

// [dependencies]
// elliptic-curve = "0.12.3"
// p256 = "0.11.1"
// sha2 = "0.10.6"

use std::cmp::min;
use std::ops::{Add, Mul};

use elliptic_curve::bigint::{Random, UInt};
use elliptic_curve::ops::Reduce;
use elliptic_curve::sec1::ToEncodedPoint;
use elliptic_curve::{Curve, Field};
use p256::ecdsa::signature::rand_core::OsRng;
use p256::{AffinePoint, NistP256, NonZeroScalar, Scalar, U256};
use sha2::{Digest, Sha256};

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn zkp() {
        let user_id = UserId::try_from("myUserName".as_bytes()).unwrap();
        let spec = ZeroKnowledgeContext::new(user_id, &[]);
        let witness = NonZeroScalar::from_uint_reduced(U256::from_be_slice(
            &"abcdefghijklmnopqrstuvwxyzabcdef".as_bytes(),
        ));
        let proof = spec.prove_knowledge_of(&witness);
        let public_key = spec.public_key(&witness);
        assert!(spec.verify(&proof, public_key));
    }
}

struct ZeroKnowledgeProof {
    challenge: Scalar,
    response: Scalar, // hidden witness
}

struct ZeroKnowledgeContext {
    generator: AffinePoint,
    user_id: UserId,
    other_info: Box<[u8]>,
}

struct UserId(Box<[u8]>);

/// BSI TR-03111 v2.10 § 4.1.1 Algorithm 2
fn choose() -> Scalar {
    let value = U256::random(OsRng);
    Scalar::from_uint_reduced(value.add_mod(&UInt::ONE, &NistP256::ORDER.min(UInt::ONE)))
}

impl ZeroKnowledgeContext {
    fn new_with_generator(generator: AffinePoint, user_id: UserId, other_info: &[u8]) -> Self {
        ZeroKnowledgeContext {
            generator,
            user_id,
            other_info: Box::from(other_info),
        }
    }

    fn new(user_id: UserId, other_info: &[u8]) -> Self {
        Self::new_with_generator(AffinePoint::GENERATOR, user_id, other_info)
    }

    pub fn public_key(self: &Self, secret: &NonZeroScalar) -> AffinePoint {
        self.generator.mul(secret.as_ref()).to_affine()
    }

    /// Uses an injective concatenation function to hash the input values.
    /// Assumes byte input is not longer than 2^(8*4) bytes
    fn challenge(self: &Self, commitment: AffinePoint, public_key: AffinePoint) -> Scalar {
        let mut digest = Sha256::new();
        for point in [self.generator, commitment, public_key] {
            let normalized = point.to_encoded_point(false);
            for c in [normalized.x(), normalized.y()] {
                let c = c.unwrap();
                assert_eq!(c.len(), 32);
                digest.update(c);
            }
        }
        for a in [&self.user_id.0, &self.other_info] {
            assert!(a.len() < 2 ^ (8 * 4));
            let prefix_content = &a.len().to_le_bytes()[..4];
            let mut buffer: [u8; 4] = [0; 4];
            buffer[..min(prefix_content.len(), 4)].copy_from_slice(&prefix_content);
            digest.update(buffer);
            digest.update(a);
        }
        let c: [u8; 32] = digest.finalize().try_into().unwrap();
        Scalar::from_uint_reduced(U256::from_be_slice(&c))
    }

    pub fn prove_knowledge_of(self: &Self, witness: &NonZeroScalar) -> ZeroKnowledgeProof {
        let randomness = choose();
        let commitment = self.generator.mul(randomness).to_affine();
        let challenge = self.challenge(commitment, self.public_key(witness));
        let response = randomness.sub(&challenge.mul(witness.as_ref()));
        let c_is_zero: bool = challenge.is_zero().into();
        let r_is_zero: bool = response.is_zero().into();
        if !c_is_zero && !r_is_zero {
            ZeroKnowledgeProof {
                challenge,
                response,
            }
        } else {
            self.prove_knowledge_of(witness)
        }
    }

    fn verify(self: &Self, proof: &ZeroKnowledgeProof, public_key: AffinePoint) -> bool {
        let commitment = self
            .generator
            .mul(proof.response)
            .add(public_key.mul(proof.challenge))
            .to_affine();
        let x = self.challenge(commitment, public_key);
        x == proof.challenge
    }
}

impl TryFrom<&[u8]> for UserId {
    type Error = ();

    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
        if value.is_empty() {
            Err(())
        } else {
            Ok(UserId(Box::from(value)))
        }
    }
}