jam-std-common 0.1.16

//! VRFs backed by [Bandersnatch](https://neuromancer.sk/std/bls/Bandersnatch),
//! an elliptic curve built over BLS12-381 scalar field.
//!
//! The primitive can operate both as a regular VRF or as an anonymized Ring VRF.

use super::concat;
use crate::{simple::TrancheIndex, Entropy, TicketAttempt, TicketId};
use ark_vrf::{
	reexports::ark_serialize::{CanonicalDeserialize, CanonicalSerialize},
	suites::bandersnatch::{self, AffinePoint, Secret as SecretImpl},
	Error,
};
use jam_types::AnyVec;
use scale::{Decode, Encode, MaxEncodedLen};

const SCALAR_SERIALIZED_SIZE: usize = 32;
#[cfg(all(test, feature = "full-test-suite"))]
const SEED_SERIALIZED_SIZE: usize = 32;
const POINT_SERIALIZED_SIZE: usize = 32;

pub const PUBLIC_SERIALIZED_SIZE: usize = POINT_SERIALIZED_SIZE;
pub const PREOUT_SERIALIZED_SIZE: usize = POINT_SERIALIZED_SIZE;

const ERR_MSG: &str = "object length is constant and checked by test; qed";

/// The raw secret seed, which can be used to reconstruct the secret [`Secret`].
type Seed = [u8; SCALAR_SERIALIZED_SIZE];

/// Bandersnatch public key.
#[derive(Clone, Copy, PartialEq, Eq, Hash, Encode, Decode, MaxEncodedLen, Default)]
pub struct Public(pub [u8; PUBLIC_SERIALIZED_SIZE]);

impl std::fmt::Debug for Public {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		write!(f, "{:?}", AnyVec(self.0.to_vec()))
	}
}

impl std::fmt::Display for Public {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		write!(f, "{}", AnyVec(self.0.to_vec()))
	}
}

/// Bandersnatch secret key.
#[derive(Clone)]
pub struct Secret(SecretImpl);

impl Secret {
	/// Construct a new secret from the given `rng`.
	pub fn new(rng: &mut impl rand::CryptoRng) -> Self {
		Self::from_seed(rand::Rng::random(rng))
	}

	/// Construct from seed.
	pub fn from_seed(seed: Seed) -> Self {
		Self(SecretImpl::from_seed(&seed))
	}

	/// Get public key component.
	pub fn public(&self) -> Public {
		let public = self.0.public();
		let mut raw = [0; PUBLIC_SERIALIZED_SIZE];
		public.serialize_compressed(raw.as_mut_slice()).expect(ERR_MSG);
		Public(raw)
	}

	/// Generate from system entropy.
	pub fn generate() -> Self {
		Self::new(&mut rand::rng())
	}

	/// Generate a trivial key from `id`.
	pub fn trivial(id: u32) -> Self {
		[id; 8].using_encoded(|data| Self::from_seed(data.try_into().expect(ERR_MSG)))
	}
}

/// A message signed by a node's Bandersnatch key.
pub enum Message {
	/// RingVRF ticket generation and regular block seal.
	TicketSeal(Entropy, TicketAttempt),
	/// Fallback block seal.
	FallbackSeal(Entropy),
	/// On-chain entropy generation.
	Entropy(TicketId),
	// Audit selection entropy for initial tranche
	AuditInitial(Entropy),
	// Audit selection entropy for subsequent tranches
	AuditSubsequent(Entropy, jam_types::WorkReportHash, TrancheIndex),
}

impl Message {
	/// Call `f` with the encoded message.
	pub fn using_encoded<R>(&self, f: impl FnOnce(&[u8]) -> R) -> R {
		match self {
			Self::TicketSeal(entropy, attempt) =>
				f(concat(&mut [0; 15 + 32 + 1], &[b"jam_ticket_seal", &entropy.0, &[*attempt]])),
			Self::FallbackSeal(entropy) =>
				f(concat(&mut [0; 17 + 32], &[b"jam_fallback_seal", &entropy.0])),
			Self::Entropy(ticket_id) =>
				f(concat(&mut [0; 11 + 32], &[b"jam_entropy", &ticket_id.0])),
			Self::AuditInitial(entropy_source) =>
				f(concat(&mut [0; 9 + 32], &[b"jam_audit", &entropy_source.0])),
			Self::AuditSubsequent(entropy_source, report, index) => f(concat(
				&mut [0; 9 + 64 + 1],
				&[b"jam_audit", &entropy_source.0, &report.0, &[*index]],
			)),
		}
	}
}

/// Generate VRF input point from raw bytes.
#[inline(always)]
fn vrf_input(message: &Message) -> bandersnatch::Input {
	message
		.using_encoded(|enc| bandersnatch::Input::new(enc).expect("Elligator2 H2C is infallible"))
}

/// Bandersnatch VRF types and operations.
pub mod vrf {
	use super::*;
	use ark_vrf::ietf::{Prover, Verifier};
	use scale::ConstEncodedLen;

	pub(crate) const PROOF_SERIALIZED_SIZE: usize = 64;
	pub const SIGNATURE_SERIALIZED_SIZE: usize = PREOUT_SERIALIZED_SIZE + PROOF_SERIALIZED_SIZE;

	/// VRF signature.
	#[derive(Clone, Debug, PartialEq, Eq, Encode, Decode, MaxEncodedLen)]
	pub struct Signature(pub [u8; SIGNATURE_SERIALIZED_SIZE]);

	impl ConstEncodedLen for Signature {}

	impl Signature {
		/// Generate VRF output bytes from the signature.
		pub fn vrf_output(&self) -> Entropy {
			bandersnatch::Output::deserialize_compressed_unchecked(
				&self.0[..PREOUT_SERIALIZED_SIZE],
			)
			.map(|p| {
				let mut raw = [0u8; 32];
				raw.copy_from_slice(&p.hash()[..32]);
				Entropy(raw)
			})
			.expect("TODO @davxy: this can fail")
		}

		/// Verify the signature.
		pub fn vrf_verify(
			&self,
			message: &Message,
			aux_data: &[u8],
			public: &Public,
		) -> Result<(), Error> {
			public.vrf_verify(message, aux_data, self)
		}
	}

	impl Secret {
		/// Creates a VRF signature of `message` and `aux_data`.
		///
		/// Auxiliary data (`aux_data`) is signed but doesn't contribute to the VRF output.
		pub fn vrf_sign(&self, message: &Message, aux_data: &[u8]) -> Signature {
			let input = vrf_input(message);
			let output = self.0.output(input);
			let proof = self.0.prove(input, output, aux_data);
			let mut raw = [0_u8; SIGNATURE_SERIALIZED_SIZE];
			output.serialize_compressed(&mut raw[..PREOUT_SERIALIZED_SIZE]).expect(ERR_MSG);
			proof.serialize_compressed(&mut raw[PREOUT_SERIALIZED_SIZE..]).expect(ERR_MSG);
			Signature(raw)
		}

		/// Generate VRF output from the given `message`.
		pub fn vrf_output(&self, message: &Message) -> Entropy {
			let input = vrf_input(message);
			let hash = self.0.output(input).hash();
			let mut raw = [0u8; 32];
			raw.copy_from_slice(&hash[..32]);
			Entropy(raw)
		}
	}

	impl Public {
		/// Verify a VRF signature.
		pub fn vrf_verify(
			&self,
			message: &Message,
			aux_data: &[u8],
			signature: &Signature,
		) -> Result<(), Error> {
			let public = bandersnatch::Public::deserialize_compressed_unchecked(self.0.as_ref())
				.map_err(|_| Error::InvalidData)?;
			let output = bandersnatch::Output::deserialize_compressed_unchecked(
				&signature.0[..PREOUT_SERIALIZED_SIZE],
			)
			.map_err(|_| Error::InvalidData)?;
			let proof = bandersnatch::IetfProof::deserialize_compressed_unchecked(
				&signature.0[PREOUT_SERIALIZED_SIZE..],
			)
			.map_err(|_| Error::InvalidData)?;
			let input = vrf_input(message);
			public.verify(input, output, aux_data, &proof)
		}
	}

	impl AsRef<[u8]> for Public {
		fn as_ref(&self) -> &[u8] {
			&self.0
		}
	}
}

pub mod ring_vrf {
	use super::*;
	use ark_vrf::ring::{Prover, Verifier};
	use bandersnatch::{RingCommitment as RingCommitmentImpl, RingProofParams};
	pub use bandersnatch::{RingProver, RingVerifier};

	pub const RING_KEYSET_SIZE: usize = 1023;

	pub const RING_COMMITMENT_SERIALIZED_SIZE: usize = 144;

	pub const RING_PROOF_SERIALIZED_SIZE: usize = 752;

	pub const RING_SIGNATURE_SERIALIZED_SIZE: usize =
		PREOUT_SERIALIZED_SIZE + RING_PROOF_SERIALIZED_SIZE;

	#[derive(Clone, PartialEq, Eq, Debug)]
	pub struct RingCommitment(RingCommitmentImpl);

	impl Encode for RingCommitment {
		fn encode(&self) -> Vec<u8> {
			let mut raw = [0; RING_COMMITMENT_SERIALIZED_SIZE];
			self.0.serialize_compressed(&mut raw[..]).expect(ERR_MSG);
			raw.encode()
		}
	}

	impl Decode for RingCommitment {
		fn decode<R: scale::Input>(i: &mut R) -> Result<Self, scale::Error> {
			let buf = <[u8; RING_COMMITMENT_SERIALIZED_SIZE]>::decode(i)?;
			let inner = RingCommitmentImpl::deserialize_compressed_unchecked(&mut &buf[..])
				.map_err(|_| scale::Error::from("Bad commitment data"))?;
			Ok(Self(inner))
		}
	}

	impl MaxEncodedLen for RingCommitment {
		fn max_encoded_len() -> usize {
			<[u8; RING_COMMITMENT_SERIALIZED_SIZE]>::max_encoded_len()
		}
	}

	/// Ring VRF signature.
	#[derive(Clone, Debug, PartialEq, Eq, Encode, Decode, MaxEncodedLen)]
	pub struct Signature(pub [u8; RING_SIGNATURE_SERIALIZED_SIZE]);

	impl Signature {
		/// Verify a ring-vrf signature.
		///
		/// The signature is verifiable if it has been produced by a member of the ring
		/// from which the [`RingVerifier`] has been constructed.
		pub fn ring_vrf_verify(
			&self,
			message: &Message,
			aux_data: &[u8],
			verifier: &RingVerifier,
		) -> Result<(), Error> {
			let output = bandersnatch::Output::deserialize_compressed_unchecked(
				&self.0[..PREOUT_SERIALIZED_SIZE],
			)
			.map_err(|_| Error::InvalidData)?;
			let proof = bandersnatch::RingProof::deserialize_compressed_unchecked(
				&self.0[PREOUT_SERIALIZED_SIZE..],
			)
			.map_err(|_| Error::InvalidData)?;
			let input = vrf_input(message);
			bandersnatch::Public::verify(input, output, aux_data, &proof, verifier)
		}

		/// Generate VRF output bytes from the signature.
		pub fn vrf_output(&self) -> Entropy {
			bandersnatch::Output::deserialize_compressed_unchecked(
				&self.0[..PREOUT_SERIALIZED_SIZE],
			)
			.map(|p| {
				let mut raw = [0u8; 32];
				raw.copy_from_slice(&p.hash()[..32]);
				Entropy(raw)
			})
			.expect("TODO @davxy: this can fail")
		}
	}

	impl Secret {
		/// Construct a Ring VRF signature.
		///
		/// Auxiliary data (`aux_data`) is signed but doesn't contribute to the VRF output.
		pub fn ring_vrf_sign(
			&self,
			message: &Message,
			aux_data: &[u8],
			prover: &RingProver,
		) -> Signature {
			let input = vrf_input(message);
			let output = self.0.output(input);
			let proof = self.0.prove(input, output, aux_data, prover);
			let mut raw = [0_u8; RING_SIGNATURE_SERIALIZED_SIZE];
			output.serialize_compressed(&mut raw[..PREOUT_SERIALIZED_SIZE]).expect(ERR_MSG);
			proof.serialize_compressed(&mut raw[PREOUT_SERIALIZED_SIZE..]).expect(ERR_MSG);
			Signature(raw)
		}
	}

	/// Ring context.
	#[derive(Clone)]
	pub struct RingContext;

	impl RingContext {
		/// Get a reference to the ring context instance.
		pub fn params() -> &'static RingProofParams {
			use std::sync::OnceLock;
			static PARAMS: OnceLock<RingProofParams> = OnceLock::new();
			PARAMS.get_or_init(|| {
				use bandersnatch::PcsParams;
				// Zcash SRS file derived from (https://zfnd.org/conclusion-of-the-powers-of-tau-ceremony).
				let buf = include_bytes!(concat!(
					env!("CARGO_MANIFEST_DIR"),
					"/data/zcash-srs-2-11-uncompressed.bin"
				));
				let pcs_params = PcsParams::deserialize_uncompressed_unchecked(&mut &buf[..])
					.expect("Error deserializing PcsParam");
				RingProofParams::from_pcs_params(jam_types::val_count() as usize, pcs_params)
					.expect("Error constructing RingProofParams from PcsParams")
			})
		}

		/// Get ring prover for the key at index `public_index` in the `public_keys` sequence.
		///
		/// Public keys sequence order matters.
		pub fn prover(public_keys: &[Public], public_index: usize) -> RingProver {
			let params = Self::params();
			let pks = public_to_affine(public_keys);
			let prover_key = params.prover_key(&pks);
			params.prover(prover_key, public_index)
		}

		/// Get ring verifier for the `public_keys` sequence.
		///
		/// Public keys sequence order matters.
		pub fn verifier(public_keys: &[Public]) -> RingVerifier {
			let params = Self::params();
			let pks = public_to_affine(public_keys);
			let verifier_key = params.verifier_key(&pks);
			params.verifier(verifier_key)
		}

		/// Get ring commitment which may be used for lazy ring verifier construction.
		///
		/// Public keys sequence order matters.
		pub fn commitment(public_keys: &[Public]) -> RingCommitment {
			let params = Self::params();
			let pks = public_to_affine(public_keys);
			let verifier_key = params.verifier_key(&pks);
			RingCommitment(verifier_key.commitment())
		}

		/// Construct a ring verifier using the ring commitment.
		pub fn verifier_from_commitment(commitment: RingCommitment) -> RingVerifier {
			let params = Self::params();
			let verifier_key = params.verifier_key_from_commitment(commitment.0);
			params.verifier(verifier_key)
		}
	}

	/// Maps a sequence of public keys to affine points on the Bandersnatch curve.
	///
	/// Any Public that cannot be mapped to a specific point will have the padding point used in its
	/// place.
	#[inline(always)]
	fn public_to_affine(pks: &[Public]) -> Vec<AffinePoint> {
		pks.iter()
			.map(|pk| {
				AffinePoint::deserialize_compressed_unchecked(&pk.0[..])
					.unwrap_or(RingProofParams::padding_point())
			})
			.collect()
	}
}

#[cfg(all(test, feature = "full-test-suite"))]
mod tests {
	#![allow(clippy::unwrap_used)]
	use super::{ring_vrf::*, vrf::*, *};
	use crate::NewNull;
	use ring_vrf::{RING_COMMITMENT_SERIALIZED_SIZE, RING_KEYSET_SIZE};

	const DEV_SEED: [u8; SEED_SERIALIZED_SIZE] = [0xcb; SEED_SERIALIZED_SIZE];

	#[allow(dead_code)]
	fn serialize<T: CanonicalSerialize>(obj: &T) -> Vec<u8> {
		let mut buf = Vec::new();
		obj.serialize_compressed(&mut buf).unwrap();
		buf
	}

	#[test]
	fn backend_assumptions_check() {
		use ark_vrf::reexports::ark_std;
		use bandersnatch::RingProofParams;
		const EXPECTED_DOMAIN_OVERHEAD: usize = 257;

		let ctx = RingProofParams::from_seed(RING_KEYSET_SIZE, [0; 32]);

		let expected_domain_size = 1 << ark_std::log2(RING_KEYSET_SIZE + EXPECTED_DOMAIN_OVERHEAD);
		assert_eq!(ctx.max_ring_size(), expected_domain_size - EXPECTED_DOMAIN_OVERHEAD);

		let pks: Vec<_> =
			(0..16).map(|i| SecretImpl::from_seed(&[i as u8; 32]).public().0).collect();

		let secret = SecretImpl::from_seed(&[0u8; 32]);

		let public = secret.public();
		assert_eq!(public.compressed_size(), PUBLIC_SERIALIZED_SIZE);

		let input = vrf_input(&Message::FallbackSeal(Default::default()));
		let output = secret.output(input);
		assert_eq!(output.compressed_size(), PREOUT_SERIALIZED_SIZE);

		let prover_key = ctx.prover_key(&pks);
		let prover = ctx.prover(prover_key, 0);

		let verifier_key = ctx.verifier_key(&pks);
		let commitment = verifier_key.commitment();
		assert_eq!(commitment.compressed_size(), RING_COMMITMENT_SERIALIZED_SIZE);

		let ietf_proof = {
			use ark_vrf::ietf::Prover;
			secret.prove(input, output, [])
		};
		assert_eq!(ietf_proof.compressed_size(), PROOF_SERIALIZED_SIZE);

		let ring_proof = {
			use ark_vrf::ring::Prover;
			secret.prove(input, output, [], &prover)
		};
		assert_eq!(ring_proof.compressed_size(), RING_PROOF_SERIALIZED_SIZE);

		// Well known Bandersnatch padding point defined within the Bandersnatch-EC-VRFs
		// specification: https://github.com/davxy/bandersnatch-vrfs-spec
		let padding_raw = [
			0x92, 0xca, 0x79, 0xe6, 0x1d, 0xd9, 0x0c, 0x15, 0x73, 0xa8, 0x69, 0x3f, 0x19, 0x9b,
			0xf6, 0xe1, 0xe8, 0x68, 0x35, 0xcc, 0x71, 0x5c, 0xdc, 0xf9, 0x3f, 0x5e, 0xf2, 0x22,
			0x56, 0x00, 0x23, 0xaa,
		];
		let padding_point = AffinePoint::deserialize_compressed(&padding_raw[..]).unwrap();
		assert_eq!(RingProofParams::padding_point(), padding_point);
	}

	fn encode_decode<T: Encode + Decode + PartialEq + std::fmt::Debug>(expected_len: usize) {
		let obj = T::new_null();
		let buf = obj.encode();
		assert_eq!(buf.len(), expected_len);
		let dec = T::decode(&mut buf.as_slice()).unwrap();
		assert_eq!(obj, dec);
	}

	#[test]
	fn codec_works() {
		encode_decode::<Public>(PUBLIC_SERIALIZED_SIZE);
		encode_decode::<vrf::Signature>(SIGNATURE_SERIALIZED_SIZE);
		encode_decode::<ring_vrf::Signature>(RING_SIGNATURE_SERIALIZED_SIZE);
	}

	#[test]
	fn vrf_sign_verify() {
		let secret = Secret::from_seed(DEV_SEED);
		let public = secret.public();

		let input = Message::FallbackSeal(Default::default());
		let ad = b"data";

		let signature = secret.vrf_sign(&input, ad);

		assert!(public.vrf_verify(&input, ad, &signature).is_ok());

		// Fail with bad additional data
		assert!(public.vrf_verify(&input, b"bad", &signature).is_err());

		// Fail with bad input
		assert!(public
			.vrf_verify(&Message::FallbackSeal([1; 32].into()), ad, &signature)
			.is_err());
	}

	#[test]
	fn vrf_output_hash_matches() {
		let secret = Secret::from_seed(DEV_SEED);

		let input = Message::FallbackSeal(Default::default());
		let signature = secret.vrf_sign(&input, b"data");

		let output1 = secret.vrf_output(&input);
		let output2 = signature.vrf_output();
		assert_eq!(output1, output2);
	}

	#[test]
	fn ring_vrf_sign_verify() {
		let mut pks: Vec<_> = (0..16).map(|i| Secret::from_seed([i as u8; 32]).public()).collect();
		assert!(pks.len() <= RING_KEYSET_SIZE);

		let secret = Secret::from_seed(DEV_SEED);

		let input = Message::FallbackSeal(Default::default());
		let ad = b"data";

		// Just pick one index to patch with the used public key
		let prover_index = 3;
		pks[prover_index] = secret.public();

		let prover = RingContext::prover(&pks, prover_index);
		let signature = secret.ring_vrf_sign(&input, ad, &prover);

		let verifier = RingContext::verifier(&pks);
		assert!(signature.ring_vrf_verify(&input, ad, &verifier).is_ok());

		// Fail with bad additional data
		assert!(signature.ring_vrf_verify(&input, b"bad", &verifier).is_err());

		// Fail with bad input
		assert!(signature
			.ring_vrf_verify(&Message::FallbackSeal([1; 32].into()), ad, &verifier)
			.is_err());
	}

	#[test]
	fn ring_vrf_sign_verify_with_out_of_ring_key() {
		let pks: Vec<_> = (0..16).map(|i| Secret::from_seed([i as u8; 32]).public()).collect();
		let secret = Secret::from_seed(DEV_SEED);

		let input = Message::FallbackSeal(Default::default());
		let ad = b"data";

		// Fair assuption that secret.public() != pks[0]
		let prover = RingContext::prover(&pks, 0);
		let signature = secret.ring_vrf_sign(&input, ad, &prover);

		let verifier = RingContext::verifier(&pks);
		assert!(signature.ring_vrf_verify(&input, ad, &verifier).is_err());
	}

	#[test]
	fn ring_vrf_make_bytes_matches() {
		let mut pks: Vec<_> = (0..16).map(|i| Secret::from_seed([i as u8; 32]).public()).collect();
		assert!(pks.len() <= RING_KEYSET_SIZE);

		let secret = Secret::from_seed(DEV_SEED);

		// Just pick one index to patch with the used public key
		let prover_index = 3;
		pks[prover_index] = secret.public();

		let input = Message::FallbackSeal(Default::default());

		let prover = RingContext::prover(&pks, prover_index);
		let signature = secret.ring_vrf_sign(&input, b"data", &prover);

		let output1 = secret.vrf_output(&input);
		let output2 = signature.vrf_output();
		assert_eq!(output1, output2);
	}
}