1
// Copyright 2019-2025 PureStake Inc.
2
// This file is part of Moonbeam.
3

            
4
// Moonbeam is free software: you can redistribute it and/or modify
5
// it under the terms of the GNU General Public License as published by
6
// the Free Software Foundation, either version 3 of the License, or
7
// (at your option) any later version.
8

            
9
// Moonbeam is distributed in the hope that it will be useful,
10
// but WITHOUT ANY WARRANTY; without even the implied warranty of
11
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12
// GNU General Public License for more details.
13

            
14
// You should have received a copy of the GNU General Public License
15
// along with Moonbeam.  If not, see <http://www.gnu.org/licenses/>.
16

            
17
//! VRF client primitives for client-side verification
18

            
19
use std::sync::Arc;
20

            
21
use nimbus_primitives::{DigestsProvider, NimbusId};
22
use schnorrkel::PublicKey;
23
use session_keys_primitives::{make_vrf_transcript, PreDigest, VrfApi, VrfId};
24
use sp_application_crypto::{AppCrypto, ByteArray};
25
use sp_core::H256;
26
use sp_keystore::{Keystore, KeystorePtr};
27

            
28
/// Uses the runtime API to get the VRF inputs and sign them with the VRF key that
29
/// corresponds to the authoring NimbusId.
30
26906
pub fn vrf_pre_digest<B, C>(
31
26906
	client: &C,
32
26906
	keystore: &KeystorePtr,
33
26906
	nimbus_id: NimbusId,
34
26906
	parent: H256,
35
26906
) -> Option<sp_runtime::generic::DigestItem>
36
26906
where
37
26906
	B: sp_runtime::traits::Block<Hash = sp_core::H256>,
38
26906
	C: sp_api::ProvideRuntimeApi<B>,
39
26906
	C::Api: VrfApi<B>,
40
26906
{
41
26906
	let runtime_api = client.runtime_api();
42

            
43
	// first ? for runtime API, second ? for if last vrf output is not available
44
26906
	let last_vrf_output = runtime_api.get_last_vrf_output(parent).ok()??;
45
	// first ? for runtime API, second ? for not VRF key associated with NimbusId
46
26016
	let key: VrfId = runtime_api.vrf_key_lookup(parent, nimbus_id).ok()??;
47
26016
	let vrf_pre_digest = sign_vrf(last_vrf_output, key, &keystore)?;
48
26016
	Some(session_keys_primitives::digest::CompatibleDigestItem::vrf_pre_digest(vrf_pre_digest))
49
26906
}
50

            
51
/// Signs the VrfInput using the private key corresponding to the input `key` public key
52
/// to be found in the input keystore
53
26016
fn sign_vrf(last_vrf_output: H256, key: VrfId, keystore: &KeystorePtr) -> Option<PreDigest> {
54
26016
	let transcript = make_vrf_transcript(last_vrf_output);
55
26016
	let try_sign = Keystore::sr25519_vrf_sign(
56
26016
		&**keystore,
57
26016
		VrfId::ID,
58
26016
		key.as_ref(),
59
26016
		&transcript.clone().into_sign_data(),
60
26016
	);
61
26016
	if let Ok(Some(signature)) = try_sign {
62
26016
		let public = PublicKey::from_bytes(&key.to_raw_vec()).ok()?;
63
26016
		if signature
64
26016
			.pre_output
65
26016
			.0
66
26016
			.attach_input_hash(&public, transcript.0.clone())
67
26016
			.is_err()
68
		{
69
			// VRF signature cannot be validated using key and transcript
70
			return None;
71
26016
		}
72
26016
		Some(PreDigest {
73
26016
			vrf_output: signature.pre_output,
74
26016
			vrf_proof: signature.proof,
75
26016
		})
76
	} else {
77
		// VRF key not found in keystore or VRF signing failed
78
		None
79
	}
80
26016
}
81

            
82
pub struct VrfDigestsProvider<B, C> {
83
	client: Arc<C>,
84
	keystore: Arc<dyn Keystore>,
85
	_marker: std::marker::PhantomData<B>,
86
}
87

            
88
impl<B, C> VrfDigestsProvider<B, C> {
89
8
	pub fn new(client: Arc<C>, keystore: Arc<dyn Keystore>) -> Self {
90
8
		Self {
91
8
			client,
92
8
			keystore,
93
8
			_marker: Default::default(),
94
8
		}
95
8
	}
96
}
97

            
98
impl<B, C> DigestsProvider<NimbusId, H256> for VrfDigestsProvider<B, C>
99
where
100
	B: sp_runtime::traits::Block<Hash = sp_core::H256>,
101
	C: sp_api::ProvideRuntimeApi<B>,
102
	C::Api: VrfApi<B>,
103
{
104
	type Digests = Option<sp_runtime::generic::DigestItem>;
105

            
106
8
	fn provide_digests(&self, nimbus_id: NimbusId, parent: H256) -> Self::Digests {
107
8
		vrf_pre_digest::<B, C>(&self.client, &self.keystore, nimbus_id, parent)
108
8
	}
109
}