cuprate_p2p_core/
client.rs

1use std::{
2    fmt::{Debug, Display, Formatter},
3    sync::{Arc, Mutex},
4    task::{ready, Context, Poll},
5};
6
7use futures::channel::oneshot;
8use tokio::sync::{mpsc, OwnedSemaphorePermit, Semaphore};
9use tokio_util::sync::{PollSemaphore, PollSender};
10use tower::{Service, ServiceExt};
11use tracing::Instrument;
12
13use cuprate_helper::asynch::InfallibleOneshotReceiver;
14use cuprate_pruning::PruningSeed;
15use cuprate_wire::{BasicNodeData, CoreSyncData};
16
17use crate::{
18    handles::{ConnectionGuard, ConnectionHandle},
19    BroadcastMessage, ConnectionDirection, NetworkZone, PeerError, PeerRequest, PeerResponse,
20};
21
22mod connection;
23mod connector;
24pub mod handshaker;
25mod request_handler;
26mod sync_callback;
27mod timeout_monitor;
28
29pub use connector::{ConnectRequest, Connector};
30pub use handshaker::{DoHandshakeRequest, HandshakeError, HandshakerBuilder};
31pub use sync_callback::PeerSyncCallback;
32
33/// An internal identifier for a given peer, will be their address if known
34/// or a random u128 if not.
35#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
36pub enum InternalPeerID<A> {
37    /// A known address.
38    KnownAddr(A),
39    /// An unknown address (probably an inbound anonymity network connection).
40    Unknown([u8; 16]),
41}
42
43impl<A: Display> Display for InternalPeerID<A> {
44    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
45        match self {
46            Self::KnownAddr(addr) => addr.fmt(f),
47            Self::Unknown(id) => f.write_str(&format!("Unknown, ID: {}", hex::encode(id))),
48        }
49    }
50}
51
52/// Information on a connected peer.
53#[derive(Debug, Clone)]
54pub struct PeerInformation<A> {
55    /// The internal peer ID of this peer.
56    pub id: InternalPeerID<A>,
57    /// The [`ConnectionHandle`] for this peer, allows banning this peer and checking if it is still
58    /// alive.
59    pub handle: ConnectionHandle,
60    /// The direction of this connection (inbound|outbound).
61    pub direction: ConnectionDirection,
62    /// The peer's [`PruningSeed`].
63    pub pruning_seed: PruningSeed,
64    /// The peer's [`BasicNodeData`].
65    pub basic_node_data: BasicNodeData,
66    /// The [`CoreSyncData`] of this peer.
67    ///
68    /// Data across fields are not necessarily related, so [`CoreSyncData::top_id`] is not always the
69    /// block hash for the block at height one below [`CoreSyncData::current_height`].
70    ///
71    /// This value is behind a [`Mutex`] and is updated whenever the peer sends new information related
72    /// to their sync state. It is publicly accessible to anyone who has a peers [`Client`] handle. You
73    /// probably should not mutate this value unless you are creating a custom [`ProtocolRequestHandler`](crate::ProtocolRequestHandler).
74    pub core_sync_data: Arc<Mutex<CoreSyncData>>,
75}
76
77/// This represents a connection to a peer.
78///
79/// It allows sending requests to the peer, but does only does minimal checks that the data returned
80/// is the data asked for, i.e. for a certain request the only thing checked will be that the response
81/// is the correct response for that request, not that the response contains the correct data.
82pub struct Client<Z: NetworkZone> {
83    /// Information on the connected peer.
84    pub info: PeerInformation<Z::Addr>,
85
86    /// The channel to the [`Connection`](connection::Connection) task.
87    connection_tx: PollSender<connection::ConnectionTaskRequest>,
88    /// The semaphore that limits the requests sent to the peer.
89    semaphore: PollSemaphore,
90    /// A permit for the semaphore, will be [`Some`] after `poll_ready` returns ready.
91    permit: Option<OwnedSemaphorePermit>,
92}
93
94impl<N: NetworkZone> Clone for Client<N> {
95    fn clone(&self) -> Self {
96        Self {
97            info: self.info.clone(),
98            connection_tx: self.connection_tx.clone(),
99            semaphore: self.semaphore.clone(),
100            permit: None,
101        }
102    }
103}
104
105impl<Z: NetworkZone> Client<Z> {
106    /// Creates a new [`Client`].
107    pub(crate) fn new(
108        info: PeerInformation<Z::Addr>,
109        connection_tx: mpsc::Sender<connection::ConnectionTaskRequest>,
110        semaphore: Arc<Semaphore>,
111    ) -> Self {
112        Self {
113            info,
114            connection_tx: PollSender::new(connection_tx),
115            semaphore: PollSemaphore::new(semaphore),
116            permit: None,
117        }
118    }
119
120    /// Waits for the client to be ready for a [`PeerRequest`], then returns a service to handle one.
121    pub fn ready_peer_request(&mut self) -> tower::util::Ready<'_, Self, PeerRequest> {
122        ServiceExt::ready(self)
123    }
124
125    /// Waits for the client to be ready for a [`BroadcastMessage`], then returns a service to handle one.
126    pub fn ready_broadcast(&mut self) -> tower::util::Ready<'_, Self, BroadcastMessage> {
127        ServiceExt::ready(self)
128    }
129}
130
131impl<Z: NetworkZone> Service<PeerRequest> for Client<Z> {
132    type Response = PeerResponse;
133    type Error = tower::BoxError;
134    type Future = InfallibleOneshotReceiver<Result<Self::Response, Self::Error>>;
135
136    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
137        if self.permit.is_none() {
138            let permit = ready!(self.semaphore.poll_acquire(cx))
139                .expect("Client semaphore should not be closed!");
140
141            self.permit = Some(permit);
142        }
143
144        if ready!(self.connection_tx.poll_reserve(cx)).is_err() {
145            return Poll::Ready(Err(PeerError::ClientChannelClosed.into()));
146        }
147
148        Poll::Ready(Ok(()))
149    }
150
151    fn call(&mut self, request: PeerRequest) -> Self::Future {
152        let permit = self
153            .permit
154            .take()
155            .expect("poll_ready did not return ready before call to call");
156
157        let (tx, rx) = oneshot::channel();
158        let req = connection::ConnectionTaskRequest {
159            response_channel: tx,
160            request,
161            permit: Some(permit),
162        };
163
164        if let Err(req) = self.connection_tx.send_item(req) {
165            // The connection task could have closed between a call to `poll_ready` and the call to
166            // `call`, which means if we don't handle the error here the receiver would panic.
167            let resp = Err(PeerError::ClientChannelClosed.into());
168            drop(req.into_inner().unwrap().response_channel.send(resp));
169        }
170
171        rx.into()
172    }
173}
174
175impl<N: NetworkZone> Service<BroadcastMessage> for Client<N> {
176    type Response = PeerResponse;
177    type Error = tower::BoxError;
178    type Future = InfallibleOneshotReceiver<Result<Self::Response, Self::Error>>;
179
180    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
181        self.permit.take();
182
183        if ready!(self.connection_tx.poll_reserve(cx)).is_err() {
184            return Poll::Ready(Err(PeerError::ClientChannelClosed.into()));
185        }
186
187        Poll::Ready(Ok(()))
188    }
189
190    fn call(&mut self, request: BroadcastMessage) -> Self::Future {
191        let (tx, rx) = oneshot::channel();
192        let req = connection::ConnectionTaskRequest {
193            response_channel: tx,
194            request: request.into(),
195            // We don't need a permit as we only accept `BroadcastMessage`, which does not require a response.
196            permit: None,
197        };
198
199        if let Err(req) = self.connection_tx.send_item(req) {
200            // The connection task could have closed between a call to `poll_ready` and the call to
201            // `call`, which means if we don't handle the error here the receiver would panic.
202            let resp = Err(PeerError::ClientChannelClosed.into());
203            drop(req.into_inner().unwrap().response_channel.send(resp));
204        }
205
206        rx.into()
207    }
208}
209
210/// Creates a mock [`Client`] for testing purposes.
211///
212/// `request_handler` will be used to handle requests sent to the [`Client`]
213pub fn mock_client<Z: NetworkZone, S>(
214    info: PeerInformation<Z::Addr>,
215    connection_guard: ConnectionGuard,
216    mut request_handler: S,
217) -> Client<Z>
218where
219    S: Service<PeerRequest, Response = PeerResponse, Error = tower::BoxError> + Send + 'static,
220    S::Future: Send + 'static,
221{
222    let (tx, mut rx) = mpsc::channel(1);
223
224    let task_span = tracing::error_span!("mock_connection", addr = %info.id);
225
226    tokio::spawn(
227        async move {
228            let _guard = connection_guard;
229            loop {
230                let Some(req): Option<connection::ConnectionTaskRequest> = rx.recv().await else {
231                    tracing::debug!("Channel closed, closing mock connection");
232                    return;
233                };
234
235                tracing::debug!("Received new request: {:?}", req.request.id());
236                let res = request_handler
237                    .ready()
238                    .await
239                    .unwrap()
240                    .call(req.request)
241                    .await
242                    .unwrap();
243
244                tracing::debug!("Sending back response");
245
246                drop(req.response_channel.send(Ok(res)));
247            }
248        }
249        .instrument(task_span),
250    );
251
252    let semaphore = Arc::new(Semaphore::new(1));
253
254    Client::new(info, tx, semaphore)
255}
cuprate_p2p_core/client.rs

cuprate_p2p_core/
client.rs
