// 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/>.
//! [`AddressedPacket`] dispatching.
use super::peer_id::{from_core_peer_id, to_core_peer_id};
use arrayvec::ArrayVec;
use log::{debug, warn};
use mixnet::core::{AddressedPacket, NetworkStatus, Packet, PeerId as CorePeerId};
use parking_lot::Mutex;
use sc_network::NotificationService;
use sc_network_types::PeerId;
use std::{collections::HashMap, future::Future, sync::Arc};
const LOG_TARGET: &str = "mixnet";
/// Packet queue for a peer.
///
/// Ideally we would use `Rc<RefCell<_>>`, but that would prevent the top-level future from being
/// automatically marked `Send`. I believe it would be safe to manually mark it `Send`, but using
/// `Arc<Mutex<_>>` here is not really a big deal.
struct PeerQueue(Mutex<ArrayVec<Box<Packet>, 2>>);
impl PeerQueue {
fn new() -> Self {
Self(Mutex::new(ArrayVec::new()))
}
/// Push `packet` onto the queue. Returns `true` if the queue was previously empty. Fails if
/// the queue is full.
fn push(&self, packet: Box<Packet>) -> Result<bool, ()> {
let mut queue = self.0.lock();
if queue.is_full() {
Err(())
} else {
let was_empty = queue.is_empty();
queue.push(packet);
Ok(was_empty)
}
}
/// Drop all packets from the queue.
fn clear(&self) {
let mut queue = self.0.lock();
queue.clear();
}
/// Pop the packet at the head of the queue and return it, or, if the queue is empty, return
/// `None`. Also returns `true` if there are more packets in the queue.
fn pop(&self) -> (Option<Box<Packet>>, bool) {
let mut queue = self.0.lock();
let packet = queue.pop();
(packet, !queue.is_empty())
}
}
/// A peer which has packets ready to send but is not currently being serviced.
pub struct ReadyPeer {
id: PeerId,
/// The peer's packet queue. Not empty.
queue: Arc<PeerQueue>,
}
impl ReadyPeer {
/// If a future is returned, and if that future returns `Some`, this function should be
/// called again to send the next packet queued for the peer; `self` is placed in the `Some`
/// to make this straightforward. Otherwise, we have either sent or dropped all packets
/// queued for the peer, and it can be forgotten about for the time being.
pub fn send_packet(
self,
notification_service: &Box<dyn NotificationService>,
) -> Option<impl Future<Output = Option<Self>>> {
match notification_service.message_sink(&self.id) {
None => {
debug!(
target: LOG_TARGET,
"Failed to get message sink for peer ID {}", self.id,
);
self.queue.clear();
None
},
Some(sink) => Some(async move {
let (packet, more_packets) = self.queue.pop();
let packet = packet.expect("Should only be called if there is a packet to send");
match sink.send_async_notification((packet as Box<[_]>).into()).await {
Ok(_) => more_packets.then_some(self),
Err(err) => {
debug!(
target: LOG_TARGET,
"Failed to send packet to peer ID {}: {err}", self.id,
);
self.queue.clear();
None
},
}
}),
}
}
}
pub struct PacketDispatcher {
/// Peer ID of the local node. Only used to implement [`NetworkStatus`].
local_peer_id: CorePeerId,
/// Packet queue for each connected peer. These queues are very short and only exist to give
/// packets somewhere to sit while waiting for notification senders to be ready.
peer_queues: HashMap<CorePeerId, Arc<PeerQueue>>,
}
impl PacketDispatcher {
pub fn new(local_peer_id: &CorePeerId) -> Self {
Self { local_peer_id: *local_peer_id, peer_queues: HashMap::new() }
}
pub fn add_peer(&mut self, id: &PeerId) {
let Some(core_id) = to_core_peer_id(id) else {
debug!(target: LOG_TARGET,
"Cannot add peer; failed to convert libp2p peer ID {id} to mixnet peer ID");
return
};
if self.peer_queues.insert(core_id, Arc::new(PeerQueue::new())).is_some() {
warn!(target: LOG_TARGET, "Two stream opened notifications for peer ID {id}");
}
}
pub fn remove_peer(&mut self, id: &PeerId) {
let Some(core_id) = to_core_peer_id(id) else {
debug!(target: LOG_TARGET,
"Cannot remove peer; failed to convert libp2p peer ID {id} to mixnet peer ID");
return
};
if self.peer_queues.remove(&core_id).is_none() {
warn!(target: LOG_TARGET, "Stream closed notification for unknown peer ID {id}");
}
}
/// If the peer is not connected or the peer's packet queue is full, the packet is dropped.
/// Otherwise the packet is pushed onto the peer's queue, and if the queue was previously empty
/// a [`ReadyPeer`] is returned.
pub fn dispatch(&mut self, packet: AddressedPacket) -> Option<ReadyPeer> {
let Some(queue) = self.peer_queues.get_mut(&packet.peer_id) else {
debug!(target: LOG_TARGET, "Dropped packet to mixnet peer ID {:x?}; not connected",
packet.peer_id);
return None
};
match queue.push(packet.packet) {
Err(_) => {
debug!(
target: LOG_TARGET,
"Dropped packet to mixnet peer ID {:x?}; peer queue full", packet.peer_id
);
None
},
Ok(true) => {
// Queue was empty. Construct and return a ReadyPeer.
let Some(id) = from_core_peer_id(&packet.peer_id) else {
debug!(target: LOG_TARGET, "Cannot send packet; \
failed to convert mixnet peer ID {:x?} to libp2p peer ID",
packet.peer_id);
queue.clear();
return None
};
Some(ReadyPeer { id, queue: queue.clone() })
},
Ok(false) => None, // Queue was not empty
}
}
}
impl NetworkStatus for PacketDispatcher {
fn local_peer_id(&self) -> CorePeerId {
self.local_peer_id
}
fn is_connected(&self, peer_id: &CorePeerId) -> bool {
self.peer_queues.contains_key(peer_id)
}
}