1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
// This file is part of Substrate.

// Copyright (C) Parity Technologies (UK) Ltd.
// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0

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

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

// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

//! BABE authority selection and slot claiming.

use super::{Epoch, AUTHORING_SCORE_LENGTH, AUTHORING_SCORE_VRF_CONTEXT};
use codec::Encode;
use sc_consensus_epochs::Epoch as EpochT;
use sp_application_crypto::AppCrypto;
use sp_consensus_babe::{
	digests::{PreDigest, PrimaryPreDigest, SecondaryPlainPreDigest, SecondaryVRFPreDigest},
	make_vrf_sign_data, AuthorityId, BabeAuthorityWeight, Randomness, Slot,
};
use sp_core::{
	blake2_256,
	crypto::{ByteArray, Wraps},
	U256,
};
use sp_keystore::KeystorePtr;

/// Calculates the primary selection threshold for a given authority, taking
/// into account `c` (`1 - c` represents the probability of a slot being empty).
pub(super) fn calculate_primary_threshold(
	c: (u64, u64),
	authorities: &[(AuthorityId, BabeAuthorityWeight)],
	authority_index: usize,
) -> u128 {
	use num_bigint::BigUint;
	use num_rational::BigRational;
	use num_traits::{cast::ToPrimitive, identities::One};

	// Prevent div by zero and out of bounds access.
	// While Babe's pallet implementation that ships with FRAME performs a sanity check over
	// configuration parameters, this is not sufficient to guarantee that `c.1` is non-zero
	// (i.e. third party implementations are possible).
	if c.1 == 0 || authority_index >= authorities.len() {
		return 0
	}

	let c = c.0 as f64 / c.1 as f64;

	let theta = authorities[authority_index].1 as f64 /
		authorities.iter().map(|(_, weight)| weight).sum::<u64>() as f64;

	assert!(theta > 0.0, "authority with weight 0.");

	// NOTE: in the equation `p = 1 - (1 - c)^theta` the value of `p` is always
	// capped by `c`. For all pratical purposes `c` should always be set to a
	// value < 0.5, as such in the computations below we should never be near
	// edge cases like `0.999999`.

	let p = BigRational::from_float(1f64 - (1f64 - c).powf(theta)).expect(
		"returns None when the given value is not finite; \
		 c is a configuration parameter defined in (0, 1]; \
		 theta must be > 0 if the given authority's weight is > 0; \
		 theta represents the validator's relative weight defined in (0, 1]; \
		 powf will always return values in (0, 1] given both the \
		 base and exponent are in that domain; \
		 qed.",
	);

	let numer = p.numer().to_biguint().expect(
		"returns None when the given value is negative; \
		 p is defined as `1 - n` where n is defined in (0, 1]; \
		 p must be a value in [0, 1); \
		 qed.",
	);

	let denom = p.denom().to_biguint().expect(
		"returns None when the given value is negative; \
		 p is defined as `1 - n` where n is defined in (0, 1]; \
		 p must be a value in [0, 1); \
		 qed.",
	);

	((BigUint::one() << 128usize) * numer / denom).to_u128().expect(
		"returns None if the underlying value cannot be represented with 128 bits; \
		 we start with 2^128 which is one more than can be represented with 128 bits; \
		 we multiple by p which is defined in [0, 1); \
		 the result must be lower than 2^128 by at least one and thus representable with 128 bits; \
		 qed.",
	)
}

/// Get the expected secondary author for the given slot and with given
/// authorities. This should always assign the slot to some authority unless the
/// authorities list is empty.
pub(super) fn secondary_slot_author(
	slot: Slot,
	authorities: &[(AuthorityId, BabeAuthorityWeight)],
	randomness: Randomness,
) -> Option<&AuthorityId> {
	if authorities.is_empty() {
		return None
	}

	let rand = U256::from((randomness, slot).using_encoded(blake2_256));

	let authorities_len = U256::from(authorities.len());
	let idx = rand % authorities_len;

	let expected_author = authorities.get(idx.as_u32() as usize).expect(
		"authorities not empty; index constrained to list length; \
				this is a valid index; qed",
	);

	Some(&expected_author.0)
}

/// Claim a secondary slot if it is our turn to propose, returning the
/// pre-digest to use when authoring the block, or `None` if it is not our turn
/// to propose.
fn claim_secondary_slot(
	slot: Slot,
	epoch: &Epoch,
	keys: &[(AuthorityId, usize)],
	keystore: &KeystorePtr,
	author_secondary_vrf: bool,
) -> Option<(PreDigest, AuthorityId)> {
	let Epoch { authorities, randomness, mut epoch_index, .. } = epoch;

	if authorities.is_empty() {
		return None
	}

	if epoch.end_slot() <= slot {
		// Slot doesn't strictly belong to the epoch, create a clone with fixed values.
		epoch_index = epoch.clone_for_slot(slot).epoch_index;
	}

	let expected_author = secondary_slot_author(slot, authorities, *randomness)?;

	for (authority_id, authority_index) in keys {
		if authority_id == expected_author {
			let pre_digest = if author_secondary_vrf {
				let data = make_vrf_sign_data(randomness, slot, epoch_index);
				let result =
					keystore.sr25519_vrf_sign(AuthorityId::ID, authority_id.as_ref(), &data);
				if let Ok(Some(vrf_signature)) = result {
					Some(PreDigest::SecondaryVRF(SecondaryVRFPreDigest {
						slot,
						authority_index: *authority_index as u32,
						vrf_signature,
					}))
				} else {
					None
				}
			} else if keystore.has_keys(&[(authority_id.to_raw_vec(), AuthorityId::ID)]) {
				Some(PreDigest::SecondaryPlain(SecondaryPlainPreDigest {
					slot,
					authority_index: *authority_index as u32,
				}))
			} else {
				None
			};

			if let Some(pre_digest) = pre_digest {
				return Some((pre_digest, authority_id.clone()))
			}
		}
	}

	None
}

/// Tries to claim the given slot number. This method starts by trying to claim
/// a primary VRF based slot. If we are not able to claim it, then if we have
/// secondary slots enabled for the given epoch, we will fallback to trying to
/// claim a secondary slot.
pub fn claim_slot(
	slot: Slot,
	epoch: &Epoch,
	keystore: &KeystorePtr,
) -> Option<(PreDigest, AuthorityId)> {
	let authorities = epoch
		.authorities
		.iter()
		.enumerate()
		.map(|(index, a)| (a.0.clone(), index))
		.collect::<Vec<_>>();
	claim_slot_using_keys(slot, epoch, keystore, &authorities)
}

/// Like `claim_slot`, but allows passing an explicit set of key pairs. Useful if we intend
/// to make repeated calls for different slots using the same key pairs.
pub fn claim_slot_using_keys(
	slot: Slot,
	epoch: &Epoch,
	keystore: &KeystorePtr,
	keys: &[(AuthorityId, usize)],
) -> Option<(PreDigest, AuthorityId)> {
	claim_primary_slot(slot, epoch, epoch.config.c, keystore, keys).or_else(|| {
		if epoch.config.allowed_slots.is_secondary_plain_slots_allowed() ||
			epoch.config.allowed_slots.is_secondary_vrf_slots_allowed()
		{
			claim_secondary_slot(
				slot,
				epoch,
				keys,
				keystore,
				epoch.config.allowed_slots.is_secondary_vrf_slots_allowed(),
			)
		} else {
			None
		}
	})
}

/// Claim a primary slot if it is our turn.  Returns `None` if it is not our turn.
/// This hashes the slot number, epoch, genesis hash, and chain randomness into
/// the VRF.  If the VRF produces a value less than `threshold`, it is our turn,
/// so it returns `Some(_)`. Otherwise, it returns `None`.
fn claim_primary_slot(
	slot: Slot,
	epoch: &Epoch,
	c: (u64, u64),
	keystore: &KeystorePtr,
	keys: &[(AuthorityId, usize)],
) -> Option<(PreDigest, AuthorityId)> {
	let Epoch { authorities, randomness, mut epoch_index, .. } = epoch;

	if epoch.end_slot() <= slot {
		// Slot doesn't strictly belong to the epoch, create a clone with fixed values.
		epoch_index = epoch.clone_for_slot(slot).epoch_index;
	}

	let data = make_vrf_sign_data(randomness, slot, epoch_index);

	for (authority_id, authority_index) in keys {
		let result = keystore.sr25519_vrf_sign(AuthorityId::ID, authority_id.as_ref(), &data);
		if let Ok(Some(vrf_signature)) = result {
			let threshold = calculate_primary_threshold(c, authorities, *authority_index);

			let can_claim = authority_id
				.as_inner_ref()
				.make_bytes::<AUTHORING_SCORE_LENGTH>(
					AUTHORING_SCORE_VRF_CONTEXT,
					&data.as_ref(),
					&vrf_signature.output,
				)
				.map(|bytes| u128::from_le_bytes(bytes) < threshold)
				.unwrap_or_default();

			if can_claim {
				let pre_digest = PreDigest::Primary(PrimaryPreDigest {
					slot,
					authority_index: *authority_index as u32,
					vrf_signature,
				});

				return Some((pre_digest, authority_id.clone()))
			}
		}
	}

	None
}

#[cfg(test)]
mod tests {
	use super::*;
	use sp_consensus_babe::{AllowedSlots, AuthorityId, BabeEpochConfiguration};
	use sp_core::{crypto::Pair as _, sr25519::Pair};
	use sp_keystore::testing::MemoryKeystore;

	#[test]
	fn claim_secondary_plain_slot_works() {
		let keystore: KeystorePtr = MemoryKeystore::new().into();
		let valid_public_key = keystore
			.sr25519_generate_new(AuthorityId::ID, Some(sp_core::crypto::DEV_PHRASE))
			.unwrap();

		let authorities = vec![
			(AuthorityId::from(Pair::generate().0.public()), 5),
			(AuthorityId::from(Pair::generate().0.public()), 7),
		];

		let mut epoch = Epoch {
			epoch_index: 10,
			start_slot: 0.into(),
			duration: 20,
			authorities: authorities.clone(),
			randomness: Default::default(),
			config: BabeEpochConfiguration {
				c: (3, 10),
				allowed_slots: AllowedSlots::PrimaryAndSecondaryPlainSlots,
			},
		};

		assert!(claim_slot(10.into(), &epoch, &keystore).is_none());

		epoch.authorities.push((valid_public_key.into(), 10));
		assert_eq!(claim_slot(10.into(), &epoch, &keystore).unwrap().1, valid_public_key.into());
	}
}