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
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
// This file is part of Substrate.

// Copyright (C) Parity Technologies (UK) Ltd.
// SPDX-License-Identifier: Apache-2.0

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// 	http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! This pallet is responsible for determining the current mixnet session and phase, and the
//! mixnode set for each session.

#![warn(missing_docs)]
#![cfg_attr(not(feature = "std"), no_std)]

extern crate alloc;

use alloc::vec::Vec;
use codec::{Decode, Encode, MaxEncodedLen};
use core::cmp::Ordering;
use frame_support::{
	traits::{EstimateNextSessionRotation, Get, OneSessionHandler},
	BoundedVec,
};
use frame_system::{
	offchain::{SendTransactionTypes, SubmitTransaction},
	pallet_prelude::BlockNumberFor,
};
pub use pallet::*;
use scale_info::TypeInfo;
use serde::{Deserialize, Serialize};
use sp_application_crypto::RuntimeAppPublic;
use sp_arithmetic::traits::{CheckedSub, Saturating, UniqueSaturatedInto, Zero};
use sp_io::MultiRemovalResults;
use sp_mixnet::types::{
	AuthorityId, AuthoritySignature, KxPublic, Mixnode, MixnodesErr, PeerId, SessionIndex,
	SessionPhase, SessionStatus, KX_PUBLIC_SIZE,
};
use sp_runtime::RuntimeDebug;

const LOG_TARGET: &str = "runtime::mixnet";

/// Index of an authority in the authority list for a session.
pub type AuthorityIndex = u32;

////////////////////////////////////////////////////////////////////////////////
// Bounded mixnode type
////////////////////////////////////////////////////////////////////////////////

/// Like [`Mixnode`], but encoded size is bounded.
#[derive(
	Clone, Decode, Encode, MaxEncodedLen, PartialEq, TypeInfo, RuntimeDebug, Serialize, Deserialize,
)]
pub struct BoundedMixnode<ExternalAddresses> {
	/// Key-exchange public key for the mixnode.
	pub kx_public: KxPublic,
	/// libp2p peer ID of the mixnode.
	pub peer_id: PeerId,
	/// External addresses for the mixnode, in multiaddr format, UTF-8 encoded.
	pub external_addresses: ExternalAddresses,
}

impl<MaxExternalAddressSize, MaxExternalAddresses> Into<Mixnode>
	for BoundedMixnode<BoundedVec<BoundedVec<u8, MaxExternalAddressSize>, MaxExternalAddresses>>
{
	fn into(self) -> Mixnode {
		Mixnode {
			kx_public: self.kx_public,
			peer_id: self.peer_id,
			external_addresses: self
				.external_addresses
				.into_iter()
				.map(BoundedVec::into_inner)
				.collect(),
		}
	}
}

impl<MaxExternalAddressSize: Get<u32>, MaxExternalAddresses: Get<u32>> From<Mixnode>
	for BoundedMixnode<BoundedVec<BoundedVec<u8, MaxExternalAddressSize>, MaxExternalAddresses>>
{
	fn from(mixnode: Mixnode) -> Self {
		Self {
			kx_public: mixnode.kx_public,
			peer_id: mixnode.peer_id,
			external_addresses: mixnode
				.external_addresses
				.into_iter()
				.flat_map(|addr| match addr.try_into() {
					Ok(addr) => Some(addr),
					Err(addr) => {
						log::debug!(
							target: LOG_TARGET,
							"Mixnode external address {addr:x?} too long; discarding",
						);
						None
					},
				})
				.take(MaxExternalAddresses::get() as usize)
				.collect::<Vec<_>>()
				.try_into()
				.expect("Excess external addresses discarded with take()"),
		}
	}
}

/// [`BoundedMixnode`] type for the given configuration.
pub type BoundedMixnodeFor<T> = BoundedMixnode<
	BoundedVec<
		BoundedVec<u8, <T as Config>::MaxExternalAddressSize>,
		<T as Config>::MaxExternalAddressesPerMixnode,
	>,
>;

////////////////////////////////////////////////////////////////////////////////
// Registration type
////////////////////////////////////////////////////////////////////////////////

/// A mixnode registration. A registration transaction is formed from one of these plus an
/// [`AuthoritySignature`].
#[derive(Clone, Decode, Encode, PartialEq, TypeInfo, RuntimeDebug)]
pub struct Registration<BlockNumber, BoundedMixnode> {
	/// Block number at the time of creation. When a registration transaction fails to make it on
	/// to the chain for whatever reason, we send out another one. We want this one to have a
	/// different hash in case the earlier transaction got banned somewhere; including the block
	/// number is a simple way of achieving this.
	pub block_number: BlockNumber,
	/// The session during which this registration should be processed. Note that on success the
	/// mixnode is registered for the _following_ session.
	pub session_index: SessionIndex,
	/// The index in the next session's authority list of the authority registering the mixnode.
	pub authority_index: AuthorityIndex,
	/// Mixnode information to register for the following session.
	pub mixnode: BoundedMixnode,
}

/// [`Registration`] type for the given configuration.
pub type RegistrationFor<T> = Registration<BlockNumberFor<T>, BoundedMixnodeFor<T>>;

////////////////////////////////////////////////////////////////////////////////
// Misc helper funcs
////////////////////////////////////////////////////////////////////////////////

fn check_removed_all(res: MultiRemovalResults) {
	debug_assert!(res.maybe_cursor.is_none());
}

fn twox<BlockNumber: UniqueSaturatedInto<u64>>(
	block_number: BlockNumber,
	kx_public: &KxPublic,
) -> u64 {
	let block_number: u64 = block_number.unique_saturated_into();
	let mut data = [0; 8 + KX_PUBLIC_SIZE];
	data[..8].copy_from_slice(&block_number.to_le_bytes());
	data[8..].copy_from_slice(kx_public);
	u64::from_le_bytes(sp_io::hashing::twox_64(&data))
}

////////////////////////////////////////////////////////////////////////////////
// The pallet
////////////////////////////////////////////////////////////////////////////////

#[frame_support::pallet(dev_mode)]
pub mod pallet {
	use super::*;
	use frame_support::pallet_prelude::*;
	use frame_system::pallet_prelude::*;

	#[pallet::pallet]
	pub struct Pallet<T>(_);

	#[pallet::config]
	pub trait Config: frame_system::Config + SendTransactionTypes<Call<Self>> {
		/// The maximum number of authorities per session.
		#[pallet::constant]
		type MaxAuthorities: Get<AuthorityIndex>;

		/// The maximum size of one of a mixnode's external addresses.
		#[pallet::constant]
		type MaxExternalAddressSize: Get<u32>;

		/// The maximum number of external addresses for a mixnode.
		#[pallet::constant]
		type MaxExternalAddressesPerMixnode: Get<u32>;

		/// Session progress/length estimation. Used to determine when to send registration
		/// transactions and the longevity of these transactions.
		type NextSessionRotation: EstimateNextSessionRotation<BlockNumberFor<Self>>;

		/// Length of the first phase of each session (`CoverToCurrent`), in blocks.
		#[pallet::constant]
		type NumCoverToCurrentBlocks: Get<BlockNumberFor<Self>>;

		/// Length of the second phase of each session (`RequestsToCurrent`), in blocks.
		#[pallet::constant]
		type NumRequestsToCurrentBlocks: Get<BlockNumberFor<Self>>;

		/// Length of the third phase of each session (`CoverToPrev`), in blocks.
		#[pallet::constant]
		type NumCoverToPrevBlocks: Get<BlockNumberFor<Self>>;

		/// The number of "slack" blocks at the start of each session, during which
		/// [`maybe_register`](Pallet::maybe_register) will not attempt to post registration
		/// transactions.
		#[pallet::constant]
		type NumRegisterStartSlackBlocks: Get<BlockNumberFor<Self>>;

		/// The number of "slack" blocks at the end of each session.
		/// [`maybe_register`](Pallet::maybe_register) will try to register before this slack
		/// period, but may post registration transactions during the slack period as a last
		/// resort.
		#[pallet::constant]
		type NumRegisterEndSlackBlocks: Get<BlockNumberFor<Self>>;

		/// Priority of unsigned transactions used to register mixnodes.
		#[pallet::constant]
		type RegistrationPriority: Get<TransactionPriority>;

		/// Minimum number of mixnodes. If there are fewer than this many mixnodes registered for a
		/// session, the mixnet will not be active during the session.
		#[pallet::constant]
		type MinMixnodes: Get<u32>;
	}

	/// Index of the current session. This may be offset relative to the session index tracked by
	/// eg `pallet_session`; mixnet session indices are independent.
	#[pallet::storage]
	pub(crate) type CurrentSessionIndex<T> = StorageValue<_, SessionIndex, ValueQuery>;

	/// Block in which the current session started.
	#[pallet::storage]
	pub(crate) type CurrentSessionStartBlock<T> = StorageValue<_, BlockNumberFor<T>, ValueQuery>;

	/// Authority list for the next session.
	#[pallet::storage]
	pub(crate) type NextAuthorityIds<T> = StorageMap<_, Identity, AuthorityIndex, AuthorityId>;

	/// Mixnode sets by session index. Only the mixnode sets for the previous, current, and next
	/// sessions are kept; older sets are discarded.
	///
	/// The mixnodes in each set are keyed by authority index so we can easily check if an
	/// authority has registered a mixnode. The authority indices should only be used during
	/// registration; the authority indices for the very first session are made up.
	#[pallet::storage]
	pub(crate) type Mixnodes<T> =
		StorageDoubleMap<_, Identity, SessionIndex, Identity, AuthorityIndex, BoundedMixnodeFor<T>>;

	#[pallet::genesis_config]
	#[derive(frame_support::DefaultNoBound)]
	pub struct GenesisConfig<T: Config> {
		/// The mixnode set for the very first session.
		pub mixnodes: BoundedVec<BoundedMixnodeFor<T>, T::MaxAuthorities>,
	}

	#[pallet::genesis_build]
	impl<T: Config> BuildGenesisConfig for GenesisConfig<T> {
		fn build(&self) {
			assert!(
				Mixnodes::<T>::iter_prefix_values(0).next().is_none(),
				"Initial mixnodes already set"
			);
			for (i, mixnode) in self.mixnodes.iter().enumerate() {
				// We just make up authority indices here. This doesn't matter as authority indices
				// are only used during registration to check an authority doesn't register twice.
				Mixnodes::<T>::insert(0, i as AuthorityIndex, mixnode);
			}
		}
	}

	#[pallet::call]
	impl<T: Config> Pallet<T> {
		/// Register a mixnode for the following session.
		#[pallet::call_index(0)]
		#[pallet::weight(1)] // TODO
		pub fn register(
			origin: OriginFor<T>,
			registration: RegistrationFor<T>,
			_signature: AuthoritySignature,
		) -> DispatchResult {
			ensure_none(origin)?;

			// Checked by ValidateUnsigned
			debug_assert_eq!(registration.session_index, CurrentSessionIndex::<T>::get());
			debug_assert!(registration.authority_index < T::MaxAuthorities::get());

			Mixnodes::<T>::insert(
				// Registering for the _following_ session
				registration.session_index + 1,
				registration.authority_index,
				registration.mixnode,
			);

			Ok(())
		}
	}

	#[pallet::validate_unsigned]
	impl<T: Config> ValidateUnsigned for Pallet<T> {
		type Call = Call<T>;

		fn validate_unsigned(_source: TransactionSource, call: &Self::Call) -> TransactionValidity {
			let Self::Call::register { registration, signature } = call else {
				return InvalidTransaction::Call.into()
			};

			// Check session index matches
			match registration.session_index.cmp(&CurrentSessionIndex::<T>::get()) {
				Ordering::Greater => return InvalidTransaction::Future.into(),
				Ordering::Less => return InvalidTransaction::Stale.into(),
				Ordering::Equal => (),
			}

			// Check authority index is valid
			if registration.authority_index >= T::MaxAuthorities::get() {
				return InvalidTransaction::BadProof.into()
			}
			let Some(authority_id) = NextAuthorityIds::<T>::get(registration.authority_index)
			else {
				return InvalidTransaction::BadProof.into()
			};

			// Check the authority hasn't registered a mixnode yet
			if Self::already_registered(registration.session_index, registration.authority_index) {
				return InvalidTransaction::Stale.into()
			}

			// Check signature. Note that we don't use regular signed transactions for registration
			// as we don't want validators to have to pay to register. Spam is prevented by only
			// allowing one registration per session per validator (see above).
			let signature_ok = registration.using_encoded(|encoded_registration| {
				authority_id.verify(&encoded_registration, signature)
			});
			if !signature_ok {
				return InvalidTransaction::BadProof.into()
			}

			ValidTransaction::with_tag_prefix("MixnetRegistration")
				.priority(T::RegistrationPriority::get())
				// Include both authority index _and_ ID in tag in case of forks with different
				// authority lists
				.and_provides((
					registration.session_index,
					registration.authority_index,
					authority_id,
				))
				.longevity(
					(T::NextSessionRotation::average_session_length() / 2_u32.into())
						.try_into()
						.unwrap_or(64_u64),
				)
				.build()
		}
	}
}

impl<T: Config> Pallet<T> {
	/// Returns the phase of the current session.
	fn session_phase() -> SessionPhase {
		let block_in_phase = frame_system::Pallet::<T>::block_number()
			.saturating_sub(CurrentSessionStartBlock::<T>::get());
		let Some(block_in_phase) = block_in_phase.checked_sub(&T::NumCoverToCurrentBlocks::get())
		else {
			return SessionPhase::CoverToCurrent
		};
		let Some(block_in_phase) =
			block_in_phase.checked_sub(&T::NumRequestsToCurrentBlocks::get())
		else {
			return SessionPhase::RequestsToCurrent
		};
		if block_in_phase < T::NumCoverToPrevBlocks::get() {
			SessionPhase::CoverToPrev
		} else {
			SessionPhase::DisconnectFromPrev
		}
	}

	/// Returns the index and phase of the current session.
	pub fn session_status() -> SessionStatus {
		SessionStatus {
			current_index: CurrentSessionIndex::<T>::get(),
			phase: Self::session_phase(),
		}
	}

	/// Returns the mixnode set for the given session (which should be either the previous or the
	/// current session).
	fn mixnodes(session_index: SessionIndex) -> Result<Vec<Mixnode>, MixnodesErr> {
		let mixnodes: Vec<_> =
			Mixnodes::<T>::iter_prefix_values(session_index).map(Into::into).collect();
		if mixnodes.len() < T::MinMixnodes::get() as usize {
			Err(MixnodesErr::InsufficientRegistrations {
				num: mixnodes.len() as u32,
				min: T::MinMixnodes::get(),
			})
		} else {
			Ok(mixnodes)
		}
	}

	/// Returns the mixnode set for the previous session.
	pub fn prev_mixnodes() -> Result<Vec<Mixnode>, MixnodesErr> {
		let Some(prev_session_index) = CurrentSessionIndex::<T>::get().checked_sub(1) else {
			return Err(MixnodesErr::InsufficientRegistrations {
				num: 0,
				min: T::MinMixnodes::get(),
			})
		};
		Self::mixnodes(prev_session_index)
	}

	/// Returns the mixnode set for the current session.
	pub fn current_mixnodes() -> Result<Vec<Mixnode>, MixnodesErr> {
		Self::mixnodes(CurrentSessionIndex::<T>::get())
	}

	/// Is now a good time to register, considering only session progress?
	fn should_register_by_session_progress(
		block_number: BlockNumberFor<T>,
		mixnode: &Mixnode,
	) -> bool {
		// At the start of each session there are some "slack" blocks during which we avoid
		// registering
		let block_in_session = block_number.saturating_sub(CurrentSessionStartBlock::<T>::get());
		if block_in_session < T::NumRegisterStartSlackBlocks::get() {
			return false
		}

		let (Some(end_block), _weight) =
			T::NextSessionRotation::estimate_next_session_rotation(block_number)
		else {
			// Things aren't going to work terribly well in this case as all the authorities will
			// just pile in after the slack period...
			return true
		};

		let remaining_blocks = end_block
			.saturating_sub(block_number)
			.saturating_sub(T::NumRegisterEndSlackBlocks::get());
		if remaining_blocks.is_zero() {
			// Into the slack time at the end of the session. Not necessarily too late;
			// registrations are accepted right up until the session ends.
			return true
		}

		// Want uniform distribution over the remaining blocks, so pick this block with probability
		// 1/remaining_blocks. maybe_register may be called multiple times per block; ensure the
		// same decision gets made each time by using a hash of the block number and the mixnode's
		// public key as the "random" source. This is slightly biased as remaining_blocks most
		// likely won't divide into 2^64, but it doesn't really matter...
		let random = twox(block_number, &mixnode.kx_public);
		(random % remaining_blocks.try_into().unwrap_or(u64::MAX)) == 0
	}

	fn next_local_authority() -> Option<(AuthorityIndex, AuthorityId)> {
		// In the case where multiple local IDs are in the next authority set, we just return the
		// first one. There's (currently at least) no point in registering multiple times.
		let mut local_ids = AuthorityId::all();
		local_ids.sort();
		NextAuthorityIds::<T>::iter().find(|(_index, id)| local_ids.binary_search(id).is_ok())
	}

	/// `session_index` should be the index of the current session. `authority_index` is the
	/// authority index in the _next_ session.
	fn already_registered(session_index: SessionIndex, authority_index: AuthorityIndex) -> bool {
		Mixnodes::<T>::contains_key(session_index + 1, authority_index)
	}

	/// Try to register a mixnode for the next session.
	///
	/// If a registration extrinsic is submitted, `true` is returned. The caller should avoid
	/// calling `maybe_register` again for a few blocks, to give the submitted extrinsic a chance
	/// to get included.
	///
	/// With the above exception, `maybe_register` is designed to be called every block. Most of
	/// the time it will not do anything, for example:
	///
	/// - If it is not an appropriate time to submit a registration extrinsic.
	/// - If the local node has already registered a mixnode for the next session.
	/// - If the local node is not permitted to register a mixnode for the next session.
	///
	/// `session_index` should match `session_status().current_index`; if it does not, `false` is
	/// returned immediately.
	pub fn maybe_register(session_index: SessionIndex, mixnode: Mixnode) -> bool {
		let current_session_index = CurrentSessionIndex::<T>::get();
		if session_index != current_session_index {
			log::trace!(
				target: LOG_TARGET,
				"Session {session_index} registration attempted, \
				but current session is {current_session_index}",
			);
			return false
		}

		let block_number = frame_system::Pallet::<T>::block_number();
		if !Self::should_register_by_session_progress(block_number, &mixnode) {
			log::trace!(
				target: LOG_TARGET,
				"Waiting for the session to progress further before registering",
			);
			return false
		}

		let Some((authority_index, authority_id)) = Self::next_local_authority() else {
			log::trace!(
				target: LOG_TARGET,
				"Not an authority in the next session; cannot register a mixnode",
			);
			return false
		};

		if Self::already_registered(session_index, authority_index) {
			log::trace!(
				target: LOG_TARGET,
				"Already registered a mixnode for the next session",
			);
			return false
		}

		let registration =
			Registration { block_number, session_index, authority_index, mixnode: mixnode.into() };
		let Some(signature) = authority_id.sign(&registration.encode()) else {
			log::debug!(target: LOG_TARGET, "Failed to sign registration");
			return false
		};
		let call = Call::register { registration, signature };
		match SubmitTransaction::<T, Call<T>>::submit_unsigned_transaction(call.into()) {
			Ok(()) => true,
			Err(()) => {
				log::debug!(
					target: LOG_TARGET,
					"Failed to submit registration transaction",
				);
				false
			},
		}
	}
}

impl<T: Config> sp_runtime::BoundToRuntimeAppPublic for Pallet<T> {
	type Public = AuthorityId;
}

impl<T: Config> OneSessionHandler<T::AccountId> for Pallet<T> {
	type Key = AuthorityId;

	fn on_genesis_session<'a, I: 'a>(validators: I)
	where
		I: Iterator<Item = (&'a T::AccountId, Self::Key)>,
	{
		assert!(
			NextAuthorityIds::<T>::iter().next().is_none(),
			"Initial authority IDs already set"
		);
		for (i, (_, authority_id)) in validators.enumerate() {
			NextAuthorityIds::<T>::insert(i as AuthorityIndex, authority_id);
		}
	}

	fn on_new_session<'a, I: 'a>(changed: bool, _validators: I, queued_validators: I)
	where
		I: Iterator<Item = (&'a T::AccountId, Self::Key)>,
	{
		let session_index = CurrentSessionIndex::<T>::mutate(|index| {
			*index += 1;
			*index
		});
		CurrentSessionStartBlock::<T>::put(frame_system::Pallet::<T>::block_number());

		// Discard the previous previous mixnode set, which we don't need any more
		if let Some(prev_prev_session_index) = session_index.checked_sub(2) {
			check_removed_all(Mixnodes::<T>::clear_prefix(
				prev_prev_session_index,
				T::MaxAuthorities::get(),
				None,
			));
		}

		if changed {
			// Save authority set for the next session. Note that we don't care about the authority
			// set for the current session; we just care about the key-exchange public keys that
			// were registered and are stored in Mixnodes.
			check_removed_all(NextAuthorityIds::<T>::clear(T::MaxAuthorities::get(), None));
			for (i, (_, authority_id)) in queued_validators.enumerate() {
				NextAuthorityIds::<T>::insert(i as AuthorityIndex, authority_id);
			}
		}
	}

	fn on_disabled(_i: u32) {
		// For now, to keep things simple, just ignore
		// TODO
	}
}