1use derive_more::Display;
20use polkadot_primitives::Hash;
21use sc_network::{
22	config::SetConfig, peer_store::PeerStoreProvider, service::NotificationMetrics,
23	types::ProtocolName, NetworkBackend, NotificationService,
24};
25use sp_runtime::traits::Block;
26use std::{
27	collections::{hash_map::Entry, HashMap},
28	ops::{Index, IndexMut},
29	sync::Arc,
30};
31use strum::{EnumIter, IntoEnumIterator};
32
33const LEGACY_COLLATION_PROTOCOL_V1: &str = "/polkadot/collation/1";
35
36const LEGACY_COLLATION_PROTOCOL_VERSION_V1: u32 = 1;
38
39pub const MAX_NOTIFICATION_SIZE: u64 = 100 * 1024;
41
42#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, EnumIter)]
44pub enum PeerSet {
45	Validation,
48	Collation,
50}
51
52#[derive(Copy, Clone, Debug, Eq, PartialEq)]
56pub enum IsAuthority {
57	Yes,
59	No,
61}
62
63impl PeerSet {
64	pub fn get_info<B: Block, N: NetworkBackend<B, <B as Block>::Hash>>(
69		self,
70		is_authority: IsAuthority,
71		peerset_protocol_names: &PeerSetProtocolNames,
72		metrics: NotificationMetrics,
73		peer_store_handle: Arc<dyn PeerStoreProvider>,
74	) -> (N::NotificationProtocolConfig, (PeerSet, Box<dyn NotificationService>)) {
75		let protocol = peerset_protocol_names.get_main_name(self);
78		let fallback_names = PeerSetProtocolNames::get_fallback_names(
79			self,
80			&peerset_protocol_names.genesis_hash,
81			peerset_protocol_names.fork_id.as_deref(),
82		);
83		let max_notification_size = self.get_max_notification_size(is_authority);
84
85		match self {
86			PeerSet::Validation => {
87				let (config, notification_service) = N::notification_config(
88					protocol,
89					fallback_names,
90					max_notification_size,
91					None,
92					SetConfig {
93						in_peers: super::MIN_GOSSIP_PEERS as u32 / 2 - 1,
98						out_peers: super::MIN_GOSSIP_PEERS as u32 / 2 - 1,
99						reserved_nodes: Vec::new(),
100						non_reserved_mode: sc_network::config::NonReservedPeerMode::Accept,
101					},
102					metrics,
103					peer_store_handle,
104				);
105
106				(config, (PeerSet::Validation, notification_service))
107			},
108			PeerSet::Collation => {
109				let (config, notification_service) = N::notification_config(
110					protocol,
111					fallback_names,
112					max_notification_size,
113					None,
114					SetConfig {
115						in_peers: if is_authority == IsAuthority::Yes { 100 } else { 0 },
118						out_peers: 0,
119						reserved_nodes: Vec::new(),
120						non_reserved_mode: if is_authority == IsAuthority::Yes {
121							sc_network::config::NonReservedPeerMode::Accept
122						} else {
123							sc_network::config::NonReservedPeerMode::Deny
124						},
125					},
126					metrics,
127					peer_store_handle,
128				);
129
130				(config, (PeerSet::Collation, notification_service))
131			},
132		}
133	}
134
135	pub fn get_main_version(self) -> ProtocolVersion {
140		match self {
141			PeerSet::Validation => ValidationVersion::V3.into(),
142			PeerSet::Collation => CollationVersion::V2.into(),
143		}
144	}
145
146	pub fn get_max_notification_size(self, _: IsAuthority) -> u64 {
148		MAX_NOTIFICATION_SIZE
149	}
150
151	pub fn get_label(self) -> &'static str {
153		match self {
154			PeerSet::Validation => "validation",
155			PeerSet::Collation => "collation",
156		}
157	}
158
159	pub fn get_protocol_label(self, version: ProtocolVersion) -> Option<&'static str> {
161		match self {
164			PeerSet::Validation =>
165				if version == ValidationVersion::V3.into() {
166					Some("validation/3")
167				} else {
168					None
169				},
170			PeerSet::Collation =>
171				if version == CollationVersion::V1.into() {
172					Some("collation/1")
173				} else if version == CollationVersion::V2.into() {
174					Some("collation/2")
175				} else {
176					None
177				},
178		}
179	}
180}
181
182#[derive(Debug, Default)]
184pub struct PerPeerSet<T> {
185	validation: T,
186	collation: T,
187}
188
189impl<T> Index<PeerSet> for PerPeerSet<T> {
190	type Output = T;
191	fn index(&self, index: PeerSet) -> &T {
192		match index {
193			PeerSet::Validation => &self.validation,
194			PeerSet::Collation => &self.collation,
195		}
196	}
197}
198
199impl<T> IndexMut<PeerSet> for PerPeerSet<T> {
200	fn index_mut(&mut self, index: PeerSet) -> &mut T {
201		match index {
202			PeerSet::Validation => &mut self.validation,
203			PeerSet::Collation => &mut self.collation,
204		}
205	}
206}
207
208pub fn peer_sets_info<B: Block, N: NetworkBackend<B, <B as Block>::Hash>>(
213	is_authority: IsAuthority,
214	peerset_protocol_names: &PeerSetProtocolNames,
215	metrics: NotificationMetrics,
216	peer_store_handle: Arc<dyn PeerStoreProvider>,
217) -> Vec<(N::NotificationProtocolConfig, (PeerSet, Box<dyn NotificationService>))> {
218	PeerSet::iter()
219		.map(|s| {
220			s.get_info::<B, N>(
221				is_authority,
222				&peerset_protocol_names,
223				metrics.clone(),
224				Arc::clone(&peer_store_handle),
225			)
226		})
227		.collect()
228}
229
230#[derive(Debug, Clone, Copy, Display, PartialEq, Eq, Hash)]
232pub struct ProtocolVersion(u32);
233
234impl From<ProtocolVersion> for u32 {
235	fn from(version: ProtocolVersion) -> u32 {
236		version.0
237	}
238}
239
240#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, EnumIter)]
242pub enum ValidationVersion {
243	V3 = 3,
245}
246
247#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, EnumIter)]
249pub enum CollationVersion {
250	V1 = 1,
252	V2 = 2,
254}
255
256#[derive(Debug, Clone, Copy, PartialEq, Eq)]
258pub struct UnknownVersion;
259
260impl TryFrom<ProtocolVersion> for ValidationVersion {
261	type Error = UnknownVersion;
262
263	fn try_from(p: ProtocolVersion) -> Result<Self, UnknownVersion> {
264		for v in Self::iter() {
265			if v as u32 == p.0 {
266				return Ok(v)
267			}
268		}
269
270		Err(UnknownVersion)
271	}
272}
273
274impl TryFrom<ProtocolVersion> for CollationVersion {
275	type Error = UnknownVersion;
276
277	fn try_from(p: ProtocolVersion) -> Result<Self, UnknownVersion> {
278		for v in Self::iter() {
279			if v as u32 == p.0 {
280				return Ok(v)
281			}
282		}
283
284		Err(UnknownVersion)
285	}
286}
287
288impl From<ValidationVersion> for ProtocolVersion {
289	fn from(version: ValidationVersion) -> ProtocolVersion {
290		ProtocolVersion(version as u32)
291	}
292}
293
294impl From<CollationVersion> for ProtocolVersion {
295	fn from(version: CollationVersion) -> ProtocolVersion {
296		ProtocolVersion(version as u32)
297	}
298}
299
300#[derive(Debug, Clone)]
302pub struct PeerSetProtocolNames {
303	protocols: HashMap<ProtocolName, (PeerSet, ProtocolVersion)>,
304	names: HashMap<(PeerSet, ProtocolVersion), ProtocolName>,
305	genesis_hash: Hash,
306	fork_id: Option<String>,
307}
308
309impl PeerSetProtocolNames {
310	pub fn new(genesis_hash: Hash, fork_id: Option<&str>) -> Self {
312		let mut protocols = HashMap::new();
313		let mut names = HashMap::new();
314		for protocol in PeerSet::iter() {
315			match protocol {
316				PeerSet::Validation =>
317					for version in ValidationVersion::iter() {
318						Self::register_main_protocol(
319							&mut protocols,
320							&mut names,
321							protocol,
322							version.into(),
323							&genesis_hash,
324							fork_id,
325						);
326					},
327				PeerSet::Collation => {
328					for version in CollationVersion::iter() {
329						Self::register_main_protocol(
330							&mut protocols,
331							&mut names,
332							protocol,
333							version.into(),
334							&genesis_hash,
335							fork_id,
336						);
337					}
338					Self::register_legacy_collation_protocol(&mut protocols, protocol);
339				},
340			}
341		}
342		Self { protocols, names, genesis_hash, fork_id: fork_id.map(|fork_id| fork_id.into()) }
343	}
344
345	fn register_main_protocol(
347		protocols: &mut HashMap<ProtocolName, (PeerSet, ProtocolVersion)>,
348		names: &mut HashMap<(PeerSet, ProtocolVersion), ProtocolName>,
349		protocol: PeerSet,
350		version: ProtocolVersion,
351		genesis_hash: &Hash,
352		fork_id: Option<&str>,
353	) {
354		let protocol_name = Self::generate_name(genesis_hash, fork_id, protocol, version);
355		names.insert((protocol, version), protocol_name.clone());
356		Self::insert_protocol_or_panic(protocols, protocol_name, protocol, version);
357	}
358
359	fn register_legacy_collation_protocol(
361		protocols: &mut HashMap<ProtocolName, (PeerSet, ProtocolVersion)>,
362		protocol: PeerSet,
363	) {
364		Self::insert_protocol_or_panic(
365			protocols,
366			LEGACY_COLLATION_PROTOCOL_V1.into(),
367			protocol,
368			ProtocolVersion(LEGACY_COLLATION_PROTOCOL_VERSION_V1),
369		)
370	}
371
372	fn insert_protocol_or_panic(
374		protocols: &mut HashMap<ProtocolName, (PeerSet, ProtocolVersion)>,
375		name: ProtocolName,
376		protocol: PeerSet,
377		version: ProtocolVersion,
378	) {
379		match protocols.entry(name) {
380			Entry::Vacant(entry) => {
381				entry.insert((protocol, version));
382			},
383			Entry::Occupied(entry) => {
384				panic!(
385					"Protocol {:?} (version {}) has the same on-the-wire name as protocol {:?} (version {}): `{}`.",
386					protocol,
387					version,
388					entry.get().0,
389					entry.get().1,
390					entry.key(),
391				);
392			},
393		}
394	}
395
396	pub fn try_get_protocol(&self, name: &ProtocolName) -> Option<(PeerSet, ProtocolVersion)> {
398		self.protocols.get(name).map(ToOwned::to_owned)
399	}
400
401	pub fn get_main_name(&self, protocol: PeerSet) -> ProtocolName {
404		self.get_name(protocol, protocol.get_main_version())
405	}
406
407	pub fn get_name(&self, protocol: PeerSet, version: ProtocolVersion) -> ProtocolName {
409		self.names
410			.get(&(protocol, version))
411			.expect("Protocols & versions are specified via enums defined above, and they are all registered in `new()`; qed")
412			.clone()
413	}
414
415	fn generate_name(
417		genesis_hash: &Hash,
418		fork_id: Option<&str>,
419		protocol: PeerSet,
420		version: ProtocolVersion,
421	) -> ProtocolName {
422		let prefix = if let Some(fork_id) = fork_id {
423			format!("/{}/{}", hex::encode(genesis_hash), fork_id)
424		} else {
425			format!("/{}", hex::encode(genesis_hash))
426		};
427
428		let short_name = match protocol {
429			PeerSet::Validation => "validation",
430			PeerSet::Collation => "collation",
431		};
432
433		format!("{}/{}/{}", prefix, short_name, version).into()
434	}
435
436	fn get_fallback_names(
439		protocol: PeerSet,
440		_genesis_hash: &Hash,
441		_fork_id: Option<&str>,
442	) -> Vec<ProtocolName> {
443		let mut fallbacks = vec![];
444		match protocol {
445			PeerSet::Validation => {
446				},
449			PeerSet::Collation => {
450				fallbacks.push(LEGACY_COLLATION_PROTOCOL_V1.into());
451			},
452		};
453		fallbacks
454	}
455}
456
457#[cfg(test)]
458mod tests {
459	use super::{
460		CollationVersion, Hash, PeerSet, PeerSetProtocolNames, ProtocolVersion, ValidationVersion,
461	};
462	use strum::IntoEnumIterator;
463
464	struct TestVersion(u32);
465
466	impl From<TestVersion> for ProtocolVersion {
467		fn from(version: TestVersion) -> ProtocolVersion {
468			ProtocolVersion(version.0)
469		}
470	}
471
472	#[test]
473	fn protocol_names_are_correctly_generated() {
474		let genesis_hash = Hash::from([
475			122, 200, 116, 29, 232, 183, 20, 109, 138, 86, 23, 253, 70, 41, 20, 85, 127, 230, 60,
476			38, 90, 127, 28, 16, 231, 218, 227, 40, 88, 238, 187, 128,
477		]);
478		let name = PeerSetProtocolNames::generate_name(
479			&genesis_hash,
480			None,
481			PeerSet::Validation,
482			TestVersion(3).into(),
483		);
484		let expected =
485			"/7ac8741de8b7146d8a5617fd462914557fe63c265a7f1c10e7dae32858eebb80/validation/3";
486		assert_eq!(name, expected.into());
487
488		let name = PeerSetProtocolNames::generate_name(
489			&genesis_hash,
490			None,
491			PeerSet::Collation,
492			TestVersion(5).into(),
493		);
494		let expected =
495			"/7ac8741de8b7146d8a5617fd462914557fe63c265a7f1c10e7dae32858eebb80/collation/5";
496		assert_eq!(name, expected.into());
497
498		let fork_id = Some("test-fork");
499		let name = PeerSetProtocolNames::generate_name(
500			&genesis_hash,
501			fork_id,
502			PeerSet::Validation,
503			TestVersion(7).into(),
504		);
505		let expected =
506			"/7ac8741de8b7146d8a5617fd462914557fe63c265a7f1c10e7dae32858eebb80/test-fork/validation/7";
507		assert_eq!(name, expected.into());
508
509		let name = PeerSetProtocolNames::generate_name(
510			&genesis_hash,
511			fork_id,
512			PeerSet::Collation,
513			TestVersion(11).into(),
514		);
515		let expected =
516			"/7ac8741de8b7146d8a5617fd462914557fe63c265a7f1c10e7dae32858eebb80/test-fork/collation/11";
517		assert_eq!(name, expected.into());
518	}
519
520	#[test]
521	fn all_protocol_names_are_known() {
522		let genesis_hash = Hash::from([
523			122, 200, 116, 29, 232, 183, 20, 109, 138, 86, 23, 253, 70, 41, 20, 85, 127, 230, 60,
524			38, 90, 127, 28, 16, 231, 218, 227, 40, 88, 238, 187, 128,
525		]);
526		let protocol_names = PeerSetProtocolNames::new(genesis_hash, None);
527
528		let validation_main =
529			"/7ac8741de8b7146d8a5617fd462914557fe63c265a7f1c10e7dae32858eebb80/validation/3";
530		assert_eq!(
531			protocol_names.try_get_protocol(&validation_main.into()),
532			Some((PeerSet::Validation, TestVersion(3).into())),
533		);
534
535		let validation_legacy = "/polkadot/validation/1";
536		assert!(protocol_names.try_get_protocol(&validation_legacy.into()).is_none());
537
538		let collation_main =
539			"/7ac8741de8b7146d8a5617fd462914557fe63c265a7f1c10e7dae32858eebb80/collation/1";
540		assert_eq!(
541			protocol_names.try_get_protocol(&collation_main.into()),
542			Some((PeerSet::Collation, TestVersion(1).into())),
543		);
544
545		let collation_legacy = "/polkadot/collation/1";
546		assert_eq!(
547			protocol_names.try_get_protocol(&collation_legacy.into()),
548			Some((PeerSet::Collation, TestVersion(1).into())),
549		);
550	}
551
552	#[test]
553	fn all_protocol_versions_are_registered() {
554		let genesis_hash = Hash::from([
555			122, 200, 116, 29, 232, 183, 20, 109, 138, 86, 23, 253, 70, 41, 20, 85, 127, 230, 60,
556			38, 90, 127, 28, 16, 231, 218, 227, 40, 88, 238, 187, 128,
557		]);
558		let protocol_names = PeerSetProtocolNames::new(genesis_hash, None);
559
560		for protocol in PeerSet::iter() {
561			match protocol {
562				PeerSet::Validation =>
563					for version in ValidationVersion::iter() {
564						assert_eq!(
565							protocol_names.get_name(protocol, version.into()),
566							PeerSetProtocolNames::generate_name(
567								&genesis_hash,
568								None,
569								protocol,
570								version.into(),
571							),
572						);
573					},
574				PeerSet::Collation =>
575					for version in CollationVersion::iter() {
576						assert_eq!(
577							protocol_names.get_name(protocol, version.into()),
578							PeerSetProtocolNames::generate_name(
579								&genesis_hash,
580								None,
581								protocol,
582								version.into(),
583							),
584						);
585					},
586			}
587		}
588	}
589
590	#[test]
591	fn all_protocol_versions_have_labels() {
592		for protocol in PeerSet::iter() {
593			match protocol {
594				PeerSet::Validation =>
595					for version in ValidationVersion::iter() {
596						protocol
597							.get_protocol_label(version.into())
598							.expect("All validation protocol versions must have a label.");
599					},
600				PeerSet::Collation =>
601					for version in CollationVersion::iter() {
602						protocol
603							.get_protocol_label(version.into())
604							.expect("All collation protocol versions must have a label.");
605					},
606			}
607		}
608	}
609}