sc_mixnet/
packet_dispatcher.rs

1// This file is part of Substrate.
2
3// Copyright (C) Parity Technologies (UK) Ltd.
4// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
5
6// This program is free software: you can redistribute it and/or modify
7// it under the terms of the GNU General Public License as published by
8// the Free Software Foundation, either version 3 of the License, or
9// (at your option) any later version.
10
11// This program is distributed in the hope that it will be useful,
12// but WITHOUT ANY WARRANTY; without even the implied warranty of
13// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14// GNU General Public License for more details.
15
16// You should have received a copy of the GNU General Public License
17// along with this program. If not, see <https://www.gnu.org/licenses/>.
18
19//! [`AddressedPacket`] dispatching.
20
21use super::peer_id::{from_core_peer_id, to_core_peer_id};
22use arrayvec::ArrayVec;
23use log::{debug, warn};
24use mixnet::core::{AddressedPacket, NetworkStatus, Packet, PeerId as CorePeerId};
25use parking_lot::Mutex;
26use sc_network::NotificationService;
27use sc_network_types::PeerId;
28use std::{collections::HashMap, future::Future, sync::Arc};
29
30const LOG_TARGET: &str = "mixnet";
31
32/// Packet queue for a peer.
33///
34/// Ideally we would use `Rc<RefCell<_>>`, but that would prevent the top-level future from being
35/// automatically marked `Send`. I believe it would be safe to manually mark it `Send`, but using
36/// `Arc<Mutex<_>>` here is not really a big deal.
37struct PeerQueue(Mutex<ArrayVec<Box<Packet>, 2>>);
38
39impl PeerQueue {
40	fn new() -> Self {
41		Self(Mutex::new(ArrayVec::new()))
42	}
43
44	/// Push `packet` onto the queue. Returns `true` if the queue was previously empty. Fails if
45	/// the queue is full.
46	fn push(&self, packet: Box<Packet>) -> Result<bool, ()> {
47		let mut queue = self.0.lock();
48		if queue.is_full() {
49			Err(())
50		} else {
51			let was_empty = queue.is_empty();
52			queue.push(packet);
53			Ok(was_empty)
54		}
55	}
56
57	/// Drop all packets from the queue.
58	fn clear(&self) {
59		let mut queue = self.0.lock();
60		queue.clear();
61	}
62
63	/// Pop the packet at the head of the queue and return it, or, if the queue is empty, return
64	/// `None`. Also returns `true` if there are more packets in the queue.
65	fn pop(&self) -> (Option<Box<Packet>>, bool) {
66		let mut queue = self.0.lock();
67		let packet = queue.pop();
68		(packet, !queue.is_empty())
69	}
70}
71
72/// A peer which has packets ready to send but is not currently being serviced.
73pub struct ReadyPeer {
74	id: PeerId,
75	/// The peer's packet queue. Not empty.
76	queue: Arc<PeerQueue>,
77}
78
79impl ReadyPeer {
80	/// If a future is returned, and if that future returns `Some`, this function should be
81	/// called again to send the next packet queued for the peer; `self` is placed in the `Some`
82	/// to make this straightforward. Otherwise, we have either sent or dropped all packets
83	/// queued for the peer, and it can be forgotten about for the time being.
84	pub fn send_packet(
85		self,
86		notification_service: &Box<dyn NotificationService>,
87	) -> Option<impl Future<Output = Option<Self>>> {
88		match notification_service.message_sink(&self.id) {
89			None => {
90				debug!(
91					target: LOG_TARGET,
92					"Failed to get message sink for peer ID {}", self.id,
93				);
94				self.queue.clear();
95				None
96			},
97			Some(sink) => Some(async move {
98				let (packet, more_packets) = self.queue.pop();
99				let packet = packet.expect("Should only be called if there is a packet to send");
100
101				match sink.send_async_notification((packet as Box<[_]>).into()).await {
102					Ok(_) => more_packets.then_some(self),
103					Err(err) => {
104						debug!(
105							target: LOG_TARGET,
106							"Failed to send packet to peer ID {}: {err}", self.id,
107						);
108						self.queue.clear();
109						None
110					},
111				}
112			}),
113		}
114	}
115}
116
117pub struct PacketDispatcher {
118	/// Peer ID of the local node. Only used to implement [`NetworkStatus`].
119	local_peer_id: CorePeerId,
120	/// Packet queue for each connected peer. These queues are very short and only exist to give
121	/// packets somewhere to sit while waiting for notification senders to be ready.
122	peer_queues: HashMap<CorePeerId, Arc<PeerQueue>>,
123}
124
125impl PacketDispatcher {
126	pub fn new(local_peer_id: &CorePeerId) -> Self {
127		Self { local_peer_id: *local_peer_id, peer_queues: HashMap::new() }
128	}
129
130	pub fn add_peer(&mut self, id: &PeerId) {
131		let Some(core_id) = to_core_peer_id(id) else {
132			debug!(target: LOG_TARGET,
133				"Cannot add peer; failed to convert libp2p peer ID {id} to mixnet peer ID");
134			return
135		};
136		if self.peer_queues.insert(core_id, Arc::new(PeerQueue::new())).is_some() {
137			warn!(target: LOG_TARGET, "Two stream opened notifications for peer ID {id}");
138		}
139	}
140
141	pub fn remove_peer(&mut self, id: &PeerId) {
142		let Some(core_id) = to_core_peer_id(id) else {
143			debug!(target: LOG_TARGET,
144				"Cannot remove peer; failed to convert libp2p peer ID {id} to mixnet peer ID");
145			return
146		};
147		if self.peer_queues.remove(&core_id).is_none() {
148			warn!(target: LOG_TARGET, "Stream closed notification for unknown peer ID {id}");
149		}
150	}
151
152	/// If the peer is not connected or the peer's packet queue is full, the packet is dropped.
153	/// Otherwise the packet is pushed onto the peer's queue, and if the queue was previously empty
154	/// a [`ReadyPeer`] is returned.
155	pub fn dispatch(&mut self, packet: AddressedPacket) -> Option<ReadyPeer> {
156		let Some(queue) = self.peer_queues.get_mut(&packet.peer_id) else {
157			debug!(target: LOG_TARGET, "Dropped packet to mixnet peer ID {:x?}; not connected",
158				packet.peer_id);
159			return None
160		};
161
162		match queue.push(packet.packet) {
163			Err(_) => {
164				debug!(
165					target: LOG_TARGET,
166					"Dropped packet to mixnet peer ID {:x?}; peer queue full", packet.peer_id
167				);
168				None
169			},
170			Ok(true) => {
171				// Queue was empty. Construct and return a ReadyPeer.
172				let Some(id) = from_core_peer_id(&packet.peer_id) else {
173					debug!(target: LOG_TARGET, "Cannot send packet; \
174						failed to convert mixnet peer ID {:x?} to libp2p peer ID",
175						packet.peer_id);
176					queue.clear();
177					return None
178				};
179				Some(ReadyPeer { id, queue: queue.clone() })
180			},
181			Ok(false) => None, // Queue was not empty
182		}
183	}
184}
185
186impl NetworkStatus for PacketDispatcher {
187	fn local_peer_id(&self) -> CorePeerId {
188		self.local_peer_id
189	}
190
191	fn is_connected(&self, peer_id: &CorePeerId) -> bool {
192		self.peer_queues.contains_key(peer_id)
193	}
194}
sc_mixnet/packet_dispatcher.rs

sc_mixnet/
packet_dispatcher.rs
