1use super::{Epoch, AUTHORING_SCORE_LENGTH, AUTHORING_SCORE_VRF_CONTEXT};
22use codec::Encode;
23use sc_consensus_epochs::Epoch as EpochT;
24use sp_application_crypto::AppCrypto;
25use sp_consensus_babe::{
26 digests::{PreDigest, PrimaryPreDigest, SecondaryPlainPreDigest, SecondaryVRFPreDigest},
27 make_vrf_sign_data, AuthorityId, BabeAuthorityWeight, Randomness, Slot,
28};
29use sp_core::{
30 crypto::{ByteArray, Wraps},
31 U256,
32};
33use sp_keystore::KeystorePtr;
34
35pub(super) fn calculate_primary_threshold(
38 c: (u64, u64),
39 authorities: &[(AuthorityId, BabeAuthorityWeight)],
40 authority_index: usize,
41) -> u128 {
42 use num_bigint::BigUint;
43 use num_rational::BigRational;
44 use num_traits::{cast::ToPrimitive, identities::One};
45
46 if c.1 == 0 || authority_index >= authorities.len() {
51 return 0
52 }
53
54 let c = c.0 as f64 / c.1 as f64;
55
56 let theta = authorities[authority_index].1 as f64 /
57 authorities.iter().map(|(_, weight)| weight).sum::<u64>() as f64;
58
59 assert!(theta > 0.0, "authority with weight 0.");
60
61 let p = BigRational::from_float(1f64 - (1f64 - c).powf(theta)).expect(
67 "returns None when the given value is not finite; \
68 c is a configuration parameter defined in (0, 1]; \
69 theta must be > 0 if the given authority's weight is > 0; \
70 theta represents the validator's relative weight defined in (0, 1]; \
71 powf will always return values in (0, 1] given both the \
72 base and exponent are in that domain; \
73 qed.",
74 );
75
76 let numer = p.numer().to_biguint().expect(
77 "returns None when the given value is negative; \
78 p is defined as `1 - n` where n is defined in (0, 1]; \
79 p must be a value in [0, 1); \
80 qed.",
81 );
82
83 let denom = p.denom().to_biguint().expect(
84 "returns None when the given value is negative; \
85 p is defined as `1 - n` where n is defined in (0, 1]; \
86 p must be a value in [0, 1); \
87 qed.",
88 );
89
90 ((BigUint::one() << 128usize) * numer / denom).to_u128().expect(
91 "returns None if the underlying value cannot be represented with 128 bits; \
92 we start with 2^128 which is one more than can be represented with 128 bits; \
93 we multiple by p which is defined in [0, 1); \
94 the result must be lower than 2^128 by at least one and thus representable with 128 bits; \
95 qed.",
96 )
97}
98
99pub(super) fn secondary_slot_author(
103 slot: Slot,
104 authorities: &[(AuthorityId, BabeAuthorityWeight)],
105 randomness: Randomness,
106) -> Option<&AuthorityId> {
107 if authorities.is_empty() {
108 return None
109 }
110
111 let rand =
112 U256::from_big_endian(&(randomness, slot).using_encoded(sp_crypto_hashing::blake2_256));
113
114 let authorities_len = U256::from(authorities.len());
115 let idx = rand % authorities_len;
116
117 let expected_author = authorities.get(idx.as_u32() as usize).expect(
118 "authorities not empty; index constrained to list length; \
119 this is a valid index; qed",
120 );
121
122 Some(&expected_author.0)
123}
124
125fn claim_secondary_slot(
129 slot: Slot,
130 epoch: &Epoch,
131 keys: &[(AuthorityId, usize)],
132 keystore: &KeystorePtr,
133 author_secondary_vrf: bool,
134) -> Option<(PreDigest, AuthorityId)> {
135 if epoch.authorities.is_empty() {
136 return None
137 }
138
139 let mut epoch_index = epoch.epoch_index;
140 if epoch.end_slot() <= slot {
141 epoch_index = epoch.clone_for_slot(slot).epoch_index;
143 }
144
145 let expected_author = secondary_slot_author(slot, &epoch.authorities, epoch.randomness)?;
146
147 for (authority_id, authority_index) in keys {
148 if authority_id == expected_author {
149 let pre_digest = if author_secondary_vrf {
150 let data = make_vrf_sign_data(&epoch.randomness, slot, epoch_index);
151 let result =
152 keystore.sr25519_vrf_sign(AuthorityId::ID, authority_id.as_ref(), &data);
153 if let Ok(Some(vrf_signature)) = result {
154 Some(PreDigest::SecondaryVRF(SecondaryVRFPreDigest {
155 slot,
156 authority_index: *authority_index as u32,
157 vrf_signature,
158 }))
159 } else {
160 None
161 }
162 } else if keystore.has_keys(&[(authority_id.to_raw_vec(), AuthorityId::ID)]) {
163 Some(PreDigest::SecondaryPlain(SecondaryPlainPreDigest {
164 slot,
165 authority_index: *authority_index as u32,
166 }))
167 } else {
168 None
169 };
170
171 if let Some(pre_digest) = pre_digest {
172 return Some((pre_digest, authority_id.clone()))
173 }
174 }
175 }
176
177 None
178}
179
180pub fn claim_slot(
185 slot: Slot,
186 epoch: &Epoch,
187 keystore: &KeystorePtr,
188) -> Option<(PreDigest, AuthorityId)> {
189 let authorities = epoch
190 .authorities
191 .iter()
192 .enumerate()
193 .map(|(index, a)| (a.0.clone(), index))
194 .collect::<Vec<_>>();
195 claim_slot_using_keys(slot, epoch, keystore, &authorities)
196}
197
198pub fn claim_slot_using_keys(
201 slot: Slot,
202 epoch: &Epoch,
203 keystore: &KeystorePtr,
204 keys: &[(AuthorityId, usize)],
205) -> Option<(PreDigest, AuthorityId)> {
206 claim_primary_slot(slot, epoch, epoch.config.c, keystore, keys).or_else(|| {
207 if epoch.config.allowed_slots.is_secondary_plain_slots_allowed() ||
208 epoch.config.allowed_slots.is_secondary_vrf_slots_allowed()
209 {
210 claim_secondary_slot(
211 slot,
212 epoch,
213 keys,
214 keystore,
215 epoch.config.allowed_slots.is_secondary_vrf_slots_allowed(),
216 )
217 } else {
218 None
219 }
220 })
221}
222
223fn claim_primary_slot(
228 slot: Slot,
229 epoch: &Epoch,
230 c: (u64, u64),
231 keystore: &KeystorePtr,
232 keys: &[(AuthorityId, usize)],
233) -> Option<(PreDigest, AuthorityId)> {
234 let mut epoch_index = epoch.epoch_index;
235 if epoch.end_slot() <= slot {
236 epoch_index = epoch.clone_for_slot(slot).epoch_index;
238 }
239
240 let data = make_vrf_sign_data(&epoch.randomness, slot, epoch_index);
241
242 for (authority_id, authority_index) in keys {
243 let result = keystore.sr25519_vrf_sign(AuthorityId::ID, authority_id.as_ref(), &data);
244 if let Ok(Some(vrf_signature)) = result {
245 let threshold = calculate_primary_threshold(c, &epoch.authorities, *authority_index);
246
247 let can_claim = authority_id
248 .as_inner_ref()
249 .make_bytes::<AUTHORING_SCORE_LENGTH>(
250 AUTHORING_SCORE_VRF_CONTEXT,
251 &data.as_ref(),
252 &vrf_signature.pre_output,
253 )
254 .map(|bytes| u128::from_le_bytes(bytes) < threshold)
255 .unwrap_or_default();
256
257 if can_claim {
258 let pre_digest = PreDigest::Primary(PrimaryPreDigest {
259 slot,
260 authority_index: *authority_index as u32,
261 vrf_signature,
262 });
263
264 return Some((pre_digest, authority_id.clone()))
265 }
266 }
267 }
268
269 None
270}
271
272#[cfg(test)]
273mod tests {
274 use super::*;
275 use sp_consensus_babe::{
276 AllowedSlots, AuthorityId, BabeEpochConfiguration, Epoch, RANDOMNESS_LENGTH,
277 };
278 use sp_core::{crypto::Pair as _, sr25519::Pair};
279 use sp_keystore::testing::MemoryKeystore;
280
281 #[test]
282 fn claim_secondary_plain_slot_works() {
283 let keystore: KeystorePtr = MemoryKeystore::new().into();
284 let valid_public_key = keystore
285 .sr25519_generate_new(AuthorityId::ID, Some(sp_core::crypto::DEV_PHRASE))
286 .unwrap();
287
288 let authorities = vec![
289 (AuthorityId::from(Pair::generate().0.public()), 5),
290 (AuthorityId::from(Pair::generate().0.public()), 7),
291 ];
292
293 let mut epoch = Epoch {
294 epoch_index: 10,
295 start_slot: 0.into(),
296 duration: 20,
297 authorities: authorities.clone(),
298 randomness: Default::default(),
299 config: BabeEpochConfiguration {
300 c: (3, 10),
301 allowed_slots: AllowedSlots::PrimaryAndSecondaryPlainSlots,
302 },
303 }
304 .into();
305
306 assert!(claim_slot(10.into(), &epoch, &keystore).is_none());
307
308 epoch.authorities.push((valid_public_key.into(), 10));
309 assert_eq!(claim_slot(10.into(), &epoch, &keystore).unwrap().1, valid_public_key.into());
310 }
311
312 #[test]
313 fn secondary_slot_author_selection_works() {
314 let authorities = (0..1000)
315 .map(|i| (AuthorityId::from(Pair::generate().0.public()), i))
316 .collect::<Vec<_>>();
317
318 let randomness = [3; RANDOMNESS_LENGTH];
319
320 assert_eq!(
321 *secondary_slot_author(100.into(), &authorities, randomness).unwrap(),
322 authorities[167].0
323 );
324 }
325}