Skip to main content

zakura_network/peer/
handshake.rs

1//! Initial [`Handshake`]s with Zebra peers over a `PeerTransport`.
2
3use std::{
4    cmp::min,
5    collections::HashSet,
6    fmt,
7    future::Future,
8    net::{IpAddr, Ipv4Addr, SocketAddr},
9    panic,
10    pin::Pin,
11    sync::Arc,
12    task::{Context, Poll},
13    time::Duration,
14};
15
16use chrono::{TimeZone, Utc};
17use futures::{channel::oneshot, future, pin_mut, FutureExt, SinkExt, StreamExt};
18use indexmap::IndexSet;
19use rand::{rngs::OsRng, RngCore};
20use tokio::{
21    io::{AsyncRead, AsyncWrite},
22    sync::broadcast,
23    task::JoinError,
24    time::{error, timeout, Instant},
25};
26use tokio_stream::wrappers::IntervalStream;
27use tokio_util::codec::Framed;
28use tower::Service;
29use tracing::Instrument;
30use tracing::{span, Level, Span};
31
32use zakura_chain::{
33    block,
34    chain_tip::{ChainTip, NoChainTip},
35    parameters::Network,
36    serialization::{DateTime32, SerializationError},
37};
38
39use crate::{
40    constants,
41    meta_addr::MetaAddrChange,
42    peer::{
43        CancelHeartbeatTask, Client, ClientRequest, Connection, ErrorSlot, HandshakeError,
44        MinimumPeerVersion, PeerError,
45    },
46    peer_registry::PeerRegistry,
47    peer_set::{ConnectionTracker, InventoryChange},
48    protocol::{
49        external::{canonical_ip, types::*, AddrInVersion, Codec, InventoryHash, Message},
50        internal::{Request, Response},
51    },
52    types::MetaAddr,
53    zakura::{
54        P2pV2Upgrade, P2pV2UpgradeAccept, P2pV2UpgradeInit, P2pV2UpgradeReject,
55        ZakuraHandshakeConfig, ZakuraLegacyNonces, ZakuraPeerId, ZakuraProtocolError,
56        P2P_V2_UPGRADE_COMMAND, PRELUDE_MAGIC,
57    },
58    zakura::{ZakuraHandshakeConnector, ZakuraRejectReason, ZakuraUpgradeOutcome},
59    BoxError, Config, PeerSocketAddr, VersionMessage,
60};
61
62#[cfg(test)]
63mod tests;
64
65/// A [`Service`] that handshakes with a remote peer and constructs a
66/// client/server pair.
67///
68/// CORRECTNESS
69///
70/// To avoid hangs, each handshake (or its connector) should be:
71/// - launched in a separate task, and
72/// - wrapped in a timeout.
73pub struct Handshake<S, C = NoChainTip>
74where
75    S: Service<Request, Response = Response, Error = BoxError> + Clone + Send + 'static,
76    S::Future: Send,
77    C: ChainTip + Clone + Send + 'static,
78{
79    config: Config,
80    user_agent: String,
81    our_services: PeerServices,
82    relay: bool,
83
84    inbound_service: S,
85    address_book_updater: tokio::sync::mpsc::Sender<MetaAddrChange>,
86    inv_collector: broadcast::Sender<InventoryChange>,
87    minimum_peer_version: MinimumPeerVersion<C>,
88    nonces: Arc<futures::lock::Mutex<IndexSet<Nonce>>>,
89    zakura_handshake_connector: Option<ZakuraHandshakeConnector>,
90    peer_registry: Option<PeerRegistry>,
91
92    /// Inbound peer IPs that are exempt from the inbound-overload connection
93    /// drop (operator-configured block-gossip / zcashd-compat sidecars),
94    /// canonicalized for IPv4-mapped matching. Empty when unconfigured.
95    protected_peer_ips: Arc<HashSet<IpAddr>>,
96
97    parent_span: Span,
98}
99
100impl<S, C> fmt::Debug for Handshake<S, C>
101where
102    S: Service<Request, Response = Response, Error = BoxError> + Clone + Send + 'static,
103    S::Future: Send,
104    C: ChainTip + Clone + Send + 'static,
105{
106    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
107        // skip the channels, they don't tell us anything useful
108        f.debug_struct(std::any::type_name::<Handshake<S, C>>())
109            .field("config", &self.config)
110            .field("user_agent", &self.user_agent)
111            .field("our_services", &self.our_services)
112            .field("relay", &self.relay)
113            .field("minimum_peer_version", &self.minimum_peer_version)
114            .field(
115                "zakura_handshake_connector",
116                &self.zakura_handshake_connector,
117            )
118            .field("parent_span", &self.parent_span)
119            .finish()
120    }
121}
122
123impl<S, C> Clone for Handshake<S, C>
124where
125    S: Service<Request, Response = Response, Error = BoxError> + Clone + Send + 'static,
126    S::Future: Send,
127    C: ChainTip + Clone + Send + 'static,
128{
129    fn clone(&self) -> Self {
130        Self {
131            config: self.config.clone(),
132            user_agent: self.user_agent.clone(),
133            our_services: self.our_services,
134            relay: self.relay,
135            inbound_service: self.inbound_service.clone(),
136            address_book_updater: self.address_book_updater.clone(),
137            inv_collector: self.inv_collector.clone(),
138            minimum_peer_version: self.minimum_peer_version.clone(),
139            nonces: self.nonces.clone(),
140            zakura_handshake_connector: self.zakura_handshake_connector.clone(),
141            peer_registry: self.peer_registry.clone(),
142            protected_peer_ips: self.protected_peer_ips.clone(),
143            parent_span: self.parent_span.clone(),
144        }
145    }
146}
147
148impl<S, C> Handshake<S, C>
149where
150    S: Service<Request, Response = Response, Error = BoxError> + Clone + Send + 'static,
151    S::Future: Send,
152    C: ChainTip + Clone + Send + 'static,
153{
154    /// Returns a connection address label using this handshake's configured privacy policy.
155    pub(crate) fn addr_label(&self, connected_addr: &ConnectedAddr) -> String {
156        connected_addr.addr_label(self.config.expose_peer_addresses)
157    }
158}
159
160/// The metadata for a peer connection.
161#[derive(Clone, Debug, PartialEq, Eq)]
162pub struct ConnectionInfo {
163    /// The connected peer address, if known.
164    /// This address might not be valid for outbound connections.
165    ///
166    /// Peers can be connected via a transient inbound or proxy address,
167    /// which will appear as the connected address to the OS and Zebra.
168    pub connected_addr: ConnectedAddr,
169
170    /// The network protocol [`VersionMessage`] sent by the remote peer.
171    pub remote: VersionMessage,
172
173    /// The network protocol [`VersionMessage`] sent by this node.
174    pub local: VersionMessage,
175
176    /// The network protocol version negotiated with the remote peer.
177    ///
178    /// Derived from `remote.version` and the
179    /// [current `zakura_network` protocol version](constants::CURRENT_NETWORK_PROTOCOL_VERSION).
180    pub negotiated_version: Version,
181
182    /// Whether this is an inbound connection from an operator-configured
183    /// protected peer (a block-gossip / zcashd-compat sidecar).
184    ///
185    /// Protected peers follow this node and learn the chain tip only from it,
186    /// so they are exempt from the inbound-overload connection drop (their
187    /// requests are still shed for backpressure, but the connection is not
188    /// severed). See [`Connection::handle_inbound_overload`].
189    pub is_protected_peer: bool,
190}
191
192/// The peer address that we are handshaking with.
193///
194/// Typically, we can rely on outbound addresses, but inbound addresses don't
195/// give us enough information to reconnect to that peer.
196#[derive(Copy, Clone, PartialEq, Eq)]
197pub enum ConnectedAddr {
198    /// The address we used to make a direct outbound connection.
199    ///
200    /// In an honest network, a Zcash peer is listening on this exact address
201    /// and port.
202    OutboundDirect {
203        /// The connected outbound remote address and port.
204        addr: PeerSocketAddr,
205    },
206
207    /// The address we received from the OS, when a remote peer directly
208    /// connected to our Zcash listener port.
209    ///
210    /// In an honest network, a Zcash peer might be listening on this address,
211    /// if its outbound address is the same as its listener address. But the port
212    /// is an ephemeral outbound TCP port, not a listener port.
213    InboundDirect {
214        /// The connected inbound remote address and ephemeral port.
215        ///
216        /// The IP address might be the address of a Zcash peer, but the port is an ephemeral port.
217        addr: PeerSocketAddr,
218    },
219
220    /// The proxy address we used to make an outbound connection.
221    ///
222    /// The proxy address can be used by many connections, but our own ephemeral
223    /// outbound address and port can be used as an identifier for the duration
224    /// of this connection.
225    OutboundProxy {
226        /// The remote address and port of the proxy.
227        proxy_addr: SocketAddr,
228
229        /// The local address and transient port we used to connect to the proxy.
230        transient_local_addr: SocketAddr,
231    },
232
233    /// The address we received from the OS, when a remote peer connected via an
234    /// inbound proxy.
235    ///
236    /// The proxy's ephemeral outbound address can be used as an identifier for
237    /// the duration of this connection.
238    InboundProxy {
239        /// The local address and transient port we used to connect to the proxy.
240        transient_addr: SocketAddr,
241    },
242
243    /// An isolated connection, where we deliberately don't have any connection metadata.
244    Isolated,
245    //
246    // TODO: handle Tor onion addresses
247}
248
249/// Get an unspecified IPv4 address for `network`
250pub fn get_unspecified_ipv4_addr(network: Network) -> SocketAddr {
251    (Ipv4Addr::UNSPECIFIED, network.default_port()).into()
252}
253
254use ConnectedAddr::*;
255
256impl ConnectedAddr {
257    /// Returns a new outbound directly connected addr.
258    pub fn new_outbound_direct(addr: PeerSocketAddr) -> ConnectedAddr {
259        OutboundDirect { addr }
260    }
261
262    /// Returns a new inbound directly connected addr.
263    pub fn new_inbound_direct(addr: PeerSocketAddr) -> ConnectedAddr {
264        InboundDirect { addr }
265    }
266
267    /// Returns a new outbound connected addr via `proxy`.
268    ///
269    /// `local_addr` is the ephemeral local address of the connection.
270    #[allow(unused)]
271    pub fn new_outbound_proxy(proxy: SocketAddr, local_addr: SocketAddr) -> ConnectedAddr {
272        OutboundProxy {
273            proxy_addr: proxy,
274            transient_local_addr: local_addr,
275        }
276    }
277
278    /// Returns a new inbound connected addr from `proxy`.
279    //
280    // TODO: distinguish between direct listeners and proxy listeners in the
281    //       rest of zakura-network
282    #[allow(unused)]
283    pub fn new_inbound_proxy(proxy: SocketAddr) -> ConnectedAddr {
284        InboundProxy {
285            transient_addr: proxy,
286        }
287    }
288
289    /// Returns a new isolated connected addr, with no metadata.
290    pub fn new_isolated() -> ConnectedAddr {
291        Isolated
292    }
293
294    /// Returns a `PeerSocketAddr` that can be used to track this connection in the
295    /// `AddressBook`.
296    ///
297    /// `None` for inbound connections, proxy connections, and isolated
298    /// connections.
299    ///
300    /// # Correctness
301    ///
302    /// This address can be used for reconnection attempts, or as a permanent
303    /// identifier.
304    ///
305    /// # Security
306    ///
307    /// This address must not depend on the canonical address from the `Version`
308    /// message. Otherwise, malicious peers could interfere with other peers
309    /// `AddressBook` state.
310    ///
311    /// TODO: remove the `get_` from these methods (Rust style avoids `get` prefixes)
312    pub fn get_address_book_addr(&self) -> Option<PeerSocketAddr> {
313        match self {
314            OutboundDirect { addr } | InboundDirect { addr } => Some(*addr),
315            // TODO: consider using the canonical address of the peer to track
316            //       outbound proxy connections
317            OutboundProxy { .. } | InboundProxy { .. } | Isolated => None,
318        }
319    }
320
321    /// Returns a `PeerSocketAddr` that can be used to temporarily identify a
322    /// connection.
323    ///
324    /// Isolated connections must not change Zebra's peer set or address book
325    /// state, so they do not have an identifier.
326    ///
327    /// # Correctness
328    ///
329    /// The returned address is only valid while the original connection is
330    /// open. It must not be used in the `AddressBook`, for outbound connection
331    /// attempts, or as a permanent identifier.
332    ///
333    /// # Security
334    ///
335    /// This address must not depend on the canonical address from the `Version`
336    /// message. Otherwise, malicious peers could interfere with other peers'
337    /// `PeerSet` state.
338    pub fn get_transient_addr(&self) -> Option<PeerSocketAddr> {
339        match self {
340            OutboundDirect { addr } => Some(*addr),
341            InboundDirect { addr } => Some(*addr),
342            OutboundProxy {
343                transient_local_addr,
344                ..
345            } => Some(PeerSocketAddr::from(*transient_local_addr)),
346            InboundProxy { transient_addr } => Some(PeerSocketAddr::from(*transient_addr)),
347            Isolated => None,
348        }
349    }
350
351    /// Returns the remote socket address when the transport identifies the peer.
352    ///
353    /// Proxy connections do not expose the peer's remote socket. In particular,
354    /// an outbound proxy connection stores this node's local ephemeral socket as
355    /// its transient identifier. Callers must not report that identifier as the
356    /// remote peer address.
357    pub(crate) fn diagnostic_remote_addr(&self) -> Option<PeerSocketAddr> {
358        match self {
359            OutboundDirect { addr } | InboundDirect { addr } => Some(*addr),
360            OutboundProxy { .. } | InboundProxy { .. } | Isolated => None,
361        }
362    }
363
364    /// Returns the redacted label for this connection's address.
365    pub fn get_transient_addr_label(&self) -> String {
366        self.get_transient_addr()
367            .map_or_else(|| "isolated".to_string(), |addr| addr.to_string())
368    }
369
370    /// Returns this connection's address using the configured privacy policy.
371    pub(crate) fn addr_label(&self, expose_peer_addresses: bool) -> String {
372        self.get_transient_addr().map_or_else(
373            || "isolated".to_string(),
374            |addr| addr.addr_label(expose_peer_addresses),
375        )
376    }
377
378    /// Returns a short label for the kind of connection.
379    pub fn get_short_kind_label(&self) -> &'static str {
380        match self {
381            OutboundDirect { .. } => "Out",
382            InboundDirect { .. } => "In",
383            OutboundProxy { .. } => "ProxOut",
384            InboundProxy { .. } => "ProxIn",
385            Isolated => "Isol",
386        }
387    }
388
389    /// Returns a list of alternate remote peer addresses, which can be used for
390    /// reconnection attempts.
391    ///
392    /// Uses the connected address, and the remote canonical address.
393    ///
394    /// Skips duplicates. If this is an outbound connection, also skips the
395    /// remote address that we're currently connected to.
396    pub fn get_alternate_addrs(
397        &self,
398        mut canonical_remote: PeerSocketAddr,
399    ) -> impl Iterator<Item = PeerSocketAddr> {
400        let addrs = match self {
401            OutboundDirect { addr } => {
402                // Fixup unspecified addresses and ports using known good data
403                if canonical_remote.ip().is_unspecified() {
404                    canonical_remote.set_ip(addr.ip());
405                }
406                if canonical_remote.port() == 0 {
407                    canonical_remote.set_port(addr.port());
408                }
409
410                // Try the canonical remote address, if it is different from the
411                // outbound address (which we already have in our address book)
412                if &canonical_remote != addr {
413                    vec![canonical_remote]
414                } else {
415                    // we didn't learn a new address from the handshake:
416                    // it's the same as the outbound address, which is already in our address book
417                    Vec::new()
418                }
419            }
420
421            InboundDirect { addr } => {
422                // Use the IP from the TCP connection, and the port the peer told us
423                let maybe_addr = SocketAddr::new(addr.ip(), canonical_remote.port()).into();
424
425                // Try both addresses, but remove one duplicate if they match
426                if canonical_remote != maybe_addr {
427                    vec![canonical_remote, maybe_addr]
428                } else {
429                    vec![canonical_remote]
430                }
431            }
432
433            // Proxy addresses can't be used for reconnection attempts, but we
434            // can try the canonical remote address
435            OutboundProxy { .. } | InboundProxy { .. } => vec![canonical_remote],
436
437            // Hide all metadata for isolated connections
438            Isolated => Vec::new(),
439        };
440
441        addrs.into_iter()
442    }
443
444    /// Returns true if the [`ConnectedAddr`] was created for an inbound connection.
445    pub fn is_inbound(&self) -> bool {
446        matches!(self, InboundDirect { .. } | InboundProxy { .. })
447    }
448
449    /// Returns `true` if this is an inbound connection whose peer IP is in the
450    /// operator-configured `protected_peer_ips` set (block-gossip / zcashd-compat
451    /// sidecars).
452    ///
453    /// The peer's transient inbound IP is canonicalized (IPv4-mapped IPv6 →
454    /// IPv4) before lookup, so a dual-stack `::ffff:` sidecar still matches an
455    /// IPv4 configuration entry. Outbound connections are never protected: the
456    /// set describes inbound sidecars that dial this node.
457    pub fn is_protected_peer(&self, protected_peer_ips: &HashSet<IpAddr>) -> bool {
458        if protected_peer_ips.is_empty() || !self.is_inbound() {
459            return false;
460        }
461
462        self.get_transient_addr()
463            .map(|addr| protected_peer_ips.contains(&canonical_ip(addr.ip())))
464            .unwrap_or(false)
465    }
466}
467
468impl fmt::Debug for ConnectedAddr {
469    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
470        let kind = self.get_short_kind_label();
471        let addr = self.get_transient_addr_label();
472
473        if matches!(self, Isolated) {
474            f.write_str(kind)
475        } else {
476            f.debug_tuple(kind).field(&addr).finish()
477        }
478    }
479}
480
481/// A builder for `Handshake`.
482pub struct Builder<S, C = NoChainTip>
483where
484    S: Service<Request, Response = Response, Error = BoxError> + Clone + Send + 'static,
485    S::Future: Send,
486    C: ChainTip + Clone + Send + 'static,
487{
488    config: Option<Config>,
489    our_services: Option<PeerServices>,
490    user_agent: Option<String>,
491    relay: Option<bool>,
492
493    inbound_service: Option<S>,
494    address_book_updater: Option<tokio::sync::mpsc::Sender<MetaAddrChange>>,
495    inv_collector: Option<broadcast::Sender<InventoryChange>>,
496    zakura_handshake_connector: Option<ZakuraHandshakeConnector>,
497    peer_registry: Option<PeerRegistry>,
498    protected_peer_ips: Option<Arc<HashSet<IpAddr>>>,
499    latest_chain_tip: C,
500}
501
502impl<S, C> Builder<S, C>
503where
504    S: Service<Request, Response = Response, Error = BoxError> + Clone + Send + 'static,
505    S::Future: Send,
506    C: ChainTip + Clone + Send + 'static,
507{
508    /// Provide a config.  Mandatory.
509    pub fn with_config(mut self, config: Config) -> Self {
510        self.config = Some(config);
511        self
512    }
513
514    /// Provide a service to handle inbound requests. Mandatory.
515    pub fn with_inbound_service(mut self, inbound_service: S) -> Self {
516        self.inbound_service = Some(inbound_service);
517        self
518    }
519
520    /// Provide a channel for registering inventory advertisements. Optional.
521    ///
522    /// This channel takes transient remote addresses, which the `PeerSet` uses
523    /// to look up peers that have specific inventory.
524    pub fn with_inventory_collector(
525        mut self,
526        inv_collector: broadcast::Sender<InventoryChange>,
527    ) -> Self {
528        self.inv_collector = Some(inv_collector);
529        self
530    }
531
532    /// Provide a hook for timestamp collection. Optional.
533    ///
534    /// This channel takes `MetaAddr`s, permanent addresses which can be used to
535    /// make outbound connections to peers.
536    pub fn with_address_book_updater(
537        mut self,
538        address_book_updater: tokio::sync::mpsc::Sender<MetaAddrChange>,
539    ) -> Self {
540        self.address_book_updater = Some(address_book_updater);
541        self
542    }
543
544    /// Provide the services this node advertises to other peers.  Optional.
545    ///
546    /// If this is unset, the node will advertise itself as a client.
547    pub fn with_advertised_services(mut self, services: PeerServices) -> Self {
548        self.our_services = Some(services);
549        self
550    }
551
552    /// Provide this node's user agent.  Optional.
553    ///
554    /// This must be a valid BIP14 string.  If it is unset, the user-agent will be empty.
555    pub fn with_user_agent(mut self, user_agent: String) -> Self {
556        self.user_agent = Some(user_agent);
557        self
558    }
559
560    /// Provide a realtime endpoint to obtain the current best chain tip block height. Optional.
561    ///
562    /// If this is unset, the minimum accepted protocol version for peer connections is kept
563    /// constant over network upgrade activations.
564    ///
565    /// Use [`NoChainTip`] to explicitly provide no chain tip.
566    pub fn with_latest_chain_tip<NewC>(self, latest_chain_tip: NewC) -> Builder<S, NewC>
567    where
568        NewC: ChainTip + Clone + Send + 'static,
569    {
570        Builder {
571            latest_chain_tip,
572
573            // TODO: Until Rust RFC 2528 reaches stable, we can't do `..self`
574            config: self.config,
575            inbound_service: self.inbound_service,
576            address_book_updater: self.address_book_updater,
577            our_services: self.our_services,
578            user_agent: self.user_agent,
579            relay: self.relay,
580            inv_collector: self.inv_collector,
581            zakura_handshake_connector: self.zakura_handshake_connector,
582            peer_registry: self.peer_registry,
583            protected_peer_ips: self.protected_peer_ips,
584        }
585    }
586
587    /// Provide a handle for routing mutually capable peers to Zakura P2P v2.
588    pub fn with_zakura_handshake_connector(
589        mut self,
590        zakura_handshake_connector: ZakuraHandshakeConnector,
591    ) -> Self {
592        self.zakura_handshake_connector = Some(zakura_handshake_connector);
593        self
594    }
595
596    /// Provide the active peer registry used by local diagnostics.
597    pub(crate) fn with_peer_registry(mut self, peer_registry: PeerRegistry) -> Self {
598        self.peer_registry = Some(peer_registry);
599        self
600    }
601
602    /// Provide the set of inbound peer IPs that are exempt from the
603    /// inbound-overload connection drop. Optional.
604    ///
605    /// These are operator-configured block-gossip / zcashd-compat sidecar IPs
606    /// (canonicalized for IPv4-mapped matching). If this is unset, no peer is
607    /// exempted and every connection behaves exactly as before.
608    pub fn with_protected_peer_ips(mut self, protected_peer_ips: Arc<HashSet<IpAddr>>) -> Self {
609        self.protected_peer_ips = Some(protected_peer_ips);
610        self
611    }
612
613    /// Whether to request that peers relay transactions to our node.  Optional.
614    ///
615    /// If this is unset, the node will not request transactions.
616    pub fn want_transactions(mut self, relay: bool) -> Self {
617        self.relay = Some(relay);
618        self
619    }
620
621    /// Consume this builder and produce a [`Handshake`].
622    ///
623    /// Returns an error only if any mandatory field was unset.
624    pub fn finish(self) -> Result<Handshake<S, C>, &'static str> {
625        let config = self.config.ok_or("did not specify config")?;
626        let inbound_service = self
627            .inbound_service
628            .ok_or("did not specify inbound service")?;
629        let inv_collector = self.inv_collector.unwrap_or_else(|| {
630            let (tx, _) = broadcast::channel(100);
631            tx
632        });
633        let address_book_updater = self.address_book_updater.unwrap_or_else(|| {
634            // No `AddressBookUpdater` for timestamp collection was passed, so create a stub
635            // channel. Dropping the receiver means sends will fail, but we don't care.
636            let (tx, _rx) = tokio::sync::mpsc::channel(1);
637            tx
638        });
639        let nonces = Arc::new(futures::lock::Mutex::new(IndexSet::new()));
640        let user_agent = self.user_agent.unwrap_or_default();
641        let our_services = configured_advertised_services(
642            &config,
643            self.our_services.unwrap_or_else(PeerServices::empty),
644        );
645        let relay = self.relay.unwrap_or(false);
646        let network = config.network.clone();
647        let minimum_peer_version = MinimumPeerVersion::new(self.latest_chain_tip, &network);
648
649        Ok(Handshake {
650            config,
651            user_agent,
652            our_services,
653            relay,
654            inbound_service,
655            address_book_updater,
656            inv_collector,
657            minimum_peer_version,
658            nonces,
659            zakura_handshake_connector: self.zakura_handshake_connector,
660            peer_registry: self.peer_registry,
661            protected_peer_ips: self.protected_peer_ips.unwrap_or_default(),
662            parent_span: Span::current(),
663        })
664    }
665}
666
667impl<S> Handshake<S, NoChainTip>
668where
669    S: Service<Request, Response = Response, Error = BoxError> + Clone + Send + 'static,
670    S::Future: Send,
671{
672    /// Create a builder that configures a [`Handshake`] service.
673    pub fn builder() -> Builder<S, NoChainTip> {
674        // We don't derive `Default` because the derive inserts a `where S:
675        // Default` bound even though `Option<S>` implements `Default` even if
676        // `S` does not.
677        Builder {
678            config: None,
679            our_services: None,
680            user_agent: None,
681            relay: None,
682            inbound_service: None,
683            address_book_updater: None,
684            inv_collector: None,
685            zakura_handshake_connector: None,
686            peer_registry: None,
687            protected_peer_ips: None,
688            latest_chain_tip: NoChainTip,
689        }
690    }
691}
692
693/// Return the services Zakura should advertise for this handshake.
694fn configured_advertised_services(config: &Config, mut services: PeerServices) -> PeerServices {
695    services.remove(PeerServices::NODE_P2P_V2);
696
697    if config.v2_p2p() {
698        services |= PeerServices::NODE_P2P_V2;
699    }
700
701    services
702}
703
704fn inbound_error_address_change(
705    addr: PeerSocketAddr,
706    remote_services: PeerServices,
707    error: &SerializationError,
708) -> MetaAddrChange {
709    // Strict parsing rejects this fault before semantic verification can assign
710    // its existing ban score, so preserve the same peer penalty here.
711    if matches!(error, SerializationError::NonCanonicalShieldedProofSize) {
712        MetaAddr::new_misbehavior(addr, constants::MAX_PEER_MISBEHAVIOR_SCORE)
713    } else {
714        MetaAddr::new_errored(addr, remote_services)
715    }
716}
717
718/// Returns true when the legacy handshake should try to route this peer to Zakura P2P v2.
719fn should_attempt_zakura_upgrade(config: &Config, connection_info: &ConnectionInfo) -> bool {
720    config.v2_p2p()
721        && connection_info
722            .remote
723            .services
724            .contains(PeerServices::NODE_P2P_V2)
725}
726
727/// Negotiate the Zcash network protocol version with the remote peer at `connected_addr`, using
728/// the connection `peer_conn`.
729///
730/// We split `Handshake` into its components before calling this function, to avoid infectious
731/// `Sync` bounds on the returned future.
732///
733/// Returns the [`VersionMessage`] sent by the remote peer, and the [`Version`] negotiated with the
734/// remote peer, inside a [`ConnectionInfo`] struct.
735#[allow(clippy::too_many_arguments)]
736pub async fn negotiate_version<PeerTransport>(
737    peer_conn: &mut Framed<PeerTransport, Codec>,
738    connected_addr: &ConnectedAddr,
739    config: Config,
740    nonces: Arc<futures::lock::Mutex<IndexSet<Nonce>>>,
741    user_agent: String,
742    our_services: PeerServices,
743    relay: bool,
744    mut minimum_peer_version: MinimumPeerVersion<impl ChainTip>,
745    is_protected_peer: bool,
746) -> Result<Arc<ConnectionInfo>, HandshakeError>
747where
748    PeerTransport: AsyncRead + AsyncWrite + Unpin + Send + 'static,
749{
750    // Create a random nonce for this connection
751    let local_nonce = Nonce::default();
752
753    // Insert the nonce for this handshake into the shared nonce set.
754    // Each connection has its own connection state, and handshakes execute concurrently.
755    //
756    // # Correctness
757    //
758    // It is ok to wait for the lock here, because handshakes have a short
759    // timeout, and the async mutex will be released when the task times
760    // out.
761    {
762        let mut locked_nonces = nonces.lock().await;
763
764        // Duplicate nonces are very rare, because they require a 64-bit random number collision,
765        // and the nonce set is limited to a few hundred entries.
766        let is_unique_nonce = locked_nonces.insert(local_nonce);
767        if !is_unique_nonce {
768            return Err(HandshakeError::LocalDuplicateNonce);
769        }
770
771        // # Security
772        //
773        // Limit the amount of memory used for nonces.
774        // Nonces can be left in the set if the connection fails or times out between
775        // the nonce being inserted, and it being removed.
776        //
777        // Zebra has strict connection limits, so we limit the number of nonces to
778        // the configured connection limit.
779        // This is a tradeoff between:
780        // - avoiding memory denial of service attacks which make large numbers of connections,
781        //   for example, 100 failed inbound connections takes 1 second.
782        // - memory usage: 16 bytes per `Nonce`, 3.2 kB for 200 nonces
783        // - collision probability: two hundred 64-bit nonces have a very low collision probability
784        //   <https://en.wikipedia.org/wiki/Birthday_problem#Probability_of_a_shared_birthday_(collision)>
785        while locked_nonces.len() > config.peerset_total_connection_limit() {
786            locked_nonces.shift_remove_index(0);
787        }
788
789        std::mem::drop(locked_nonces);
790    }
791
792    // Don't leak our exact clock skew to our peers. On the other hand,
793    // we can't deviate too much, or zcashd will get confused.
794    // Inspection of the zcashd source code reveals that the timestamp
795    // is only ever used at the end of parsing the version message, in
796    //
797    // pfrom->nTimeOffset = timeWarning.AddTimeData(pfrom->addr, nTime, GetTime());
798    //
799    // AddTimeData is defined in src/timedata.cpp and is a no-op as long
800    // as the difference between the specified timestamp and the
801    // zcashd's local time is less than TIMEDATA_WARNING_THRESHOLD, set
802    // to 10 * 60 seconds (10 minutes).
803    //
804    // nTimeOffset is peer metadata that is never used, except for
805    // statistics.
806    //
807    // To try to stay within the range where zcashd will ignore our clock skew,
808    // truncate the timestamp to the nearest 5 minutes.
809    let now = Utc::now().timestamp();
810    let timestamp = Utc
811        .timestamp_opt(now - now.rem_euclid(5 * 60), 0)
812        .single()
813        .expect("in-range number of seconds and valid nanosecond");
814
815    let (their_addr, our_services, our_listen_addr) = match connected_addr {
816        // Version messages require an address, so we use
817        // an unspecified address for Isolated connections
818        Isolated => {
819            let unspec_ipv4 = get_unspecified_ipv4_addr(config.network);
820            (unspec_ipv4.into(), PeerServices::empty(), unspec_ipv4)
821        }
822        _ => {
823            let their_addr = connected_addr
824                .get_transient_addr()
825                .expect("non-Isolated connections have a remote addr");
826
827            // Include the configured external address in our version message, if any, otherwise, include our listen address.
828            let advertise_addr = match config.external_addr {
829                Some(external_addr) => {
830                    info!(
831                        peer = %their_addr.addr_label(config.expose_peer_addresses),
832                        ?config.listen_addr,
833                        "using external address for Version messages"
834                    );
835                    external_addr
836                }
837                None => config.listen_addr,
838            };
839
840            (their_addr, our_services, advertise_addr)
841        }
842    };
843
844    let start_height = minimum_peer_version
845        .chain_tip()
846        .best_tip_height()
847        .unwrap_or(block::Height(0));
848
849    let our_version = VersionMessage {
850        version: constants::CURRENT_NETWORK_PROTOCOL_VERSION,
851        services: our_services,
852        timestamp,
853        address_recv: AddrInVersion::new(their_addr, PeerServices::NODE_NETWORK),
854        // TODO: detect external address (#1893)
855        address_from: AddrInVersion::new(our_listen_addr, our_services),
856        nonce: local_nonce,
857        user_agent: user_agent.clone(),
858        start_height,
859        relay,
860    };
861
862    debug!(?our_version, "sending initial version message");
863    if our_services.contains(PeerServices::NODE_P2P_V2) {
864        metrics::counter!("zakura.p2p.handshake.service_bit.advertised").increment(1);
865    }
866    peer_conn.send(our_version.clone().into()).await?;
867
868    let mut remote_msg = peer_conn
869        .next()
870        .await
871        .ok_or(HandshakeError::ConnectionClosed)??;
872
873    // Wait for next message if the one we got is not Version
874    let remote: VersionMessage = loop {
875        match remote_msg {
876            Message::Version(version_message) => {
877                debug!(?version_message, "got version message from remote peer");
878                break version_message;
879            }
880            _ => {
881                remote_msg = peer_conn
882                    .next()
883                    .await
884                    .ok_or(HandshakeError::ConnectionClosed)??;
885                debug!(?remote_msg, "ignoring non-version message from remote peer");
886            }
887        }
888    };
889
890    let remote_address_services = remote.address_from.untrusted_services();
891    let addr_label = their_addr.addr_label(config.expose_peer_addresses);
892
893    if remote_address_services != remote.services {
894        info!(
895            ?remote.services,
896            ?remote_address_services,
897            ?remote.user_agent,
898            "peer with inconsistent version services and version address services",
899        );
900    }
901    if remote.services.contains(PeerServices::NODE_P2P_V2) {
902        metrics::counter!("zakura.p2p.handshake.service_bit.remote").increment(1);
903    }
904
905    // Check for nonce reuse, indicating self-connection
906    //
907    // # Correctness
908    //
909    // We must wait for the lock before we continue with the connection, to avoid
910    // self-connection. If the connection times out, the async lock will be
911    // released.
912    //
913    // # Security
914    //
915    // We don't remove the nonce here, because peers that observe our network traffic could
916    // maliciously remove nonces, and force us to make self-connections.
917    let nonce_reuse = nonces.lock().await.contains(&remote.nonce);
918    if nonce_reuse {
919        info!(
920            peer = %addr_label,
921            connection_kind = connected_addr.get_short_kind_label(),
922            "rejecting self-connection attempt"
923        );
924        Err(HandshakeError::RemoteNonceReuse)?;
925    }
926
927    // # Security
928    //
929    // Reject connections to peers on old versions, because they might not know about all
930    // network upgrades and could lead to chain forks or slower block propagation.
931    let min_version = minimum_peer_version.current();
932    if remote.version < min_version {
933        debug!(
934            remote_ip = %addr_label,
935            ?remote.version,
936            ?min_version,
937            ?remote.user_agent,
938            "disconnecting from peer with obsolete network protocol version",
939        );
940
941        // Handshake rejects by protocol version. Not labeled by peer address or
942        // user-agent: the Prometheus exporter never prunes series.
943        metrics::counter!(
944            "zcash.net.peers.obsolete",
945            "remote_version" => remote.version.to_string(),
946            "min_version" => min_version.to_string(),
947        )
948        .increment(1);
949
950        metrics::gauge!("zcash.net.peers.version.obsolete").set(remote.version.0 as f64);
951
952        // Disconnect if peer is using an obsolete version.
953        return Err(HandshakeError::ObsoleteVersion(remote.version));
954    }
955
956    let negotiated_version = min(constants::CURRENT_NETWORK_PROTOCOL_VERSION, remote.version);
957
958    // Limit containing struct size, and avoid multiple duplicates of 300+ bytes of data.
959    let connection_info = Arc::new(ConnectionInfo {
960        connected_addr: *connected_addr,
961        remote,
962        local: our_version,
963        negotiated_version,
964        is_protected_peer,
965    });
966
967    debug!(
968        remote_ip = %addr_label,
969        ?connection_info.remote.version,
970        ?negotiated_version,
971        ?min_version,
972        ?connection_info.remote.user_agent,
973        "negotiated network protocol version with peer",
974    );
975
976    // Handshake count by protocol version. Not labeled by peer address or
977    // user-agent: the Prometheus exporter never prunes series.
978    metrics::counter!(
979        "zcash.net.peers.connected",
980        "remote_version" => connection_info.remote.version.to_string(),
981        "negotiated_version" => negotiated_version.to_string(),
982        "min_version" => min_version.to_string(),
983    )
984    .increment(1);
985
986    metrics::gauge!("zcash.net.peers.version.connected")
987        .set(connection_info.remote.version.0 as f64);
988
989    peer_conn.send(Message::Verack).await?;
990
991    let mut remote_msg = peer_conn
992        .next()
993        .await
994        .ok_or(HandshakeError::ConnectionClosed)??;
995
996    // Wait for next message if the one we got is not Verack
997    loop {
998        match remote_msg {
999            Message::Verack => {
1000                debug!(?remote_msg, "got verack message from remote peer");
1001                break;
1002            }
1003            _ => {
1004                remote_msg = peer_conn
1005                    .next()
1006                    .await
1007                    .ok_or(HandshakeError::ConnectionClosed)??;
1008                debug!(?remote_msg, "ignoring non-verack message from remote peer");
1009            }
1010        }
1011    }
1012
1013    Ok(connection_info)
1014}
1015
1016/// Route a mutually P2P-v2-capable peer into the Zakura handshake path.
1017///
1018/// After the legacy `version`/`verack` exchange, two mutually capable peers swap
1019/// a bounded [`P2pV2Upgrade`] prelude over the legacy TCP stream to learn each
1020/// other's Zakura (iroh) node address. The TCP initiator then dials the
1021/// responder over QUIC; the responder's iroh router accepts that dial and
1022/// registers the peer. On success the caller drops the legacy stream and the
1023/// peers continue over Zakura. Any neutral problem (no local endpoint, malformed
1024/// or rejected prelude) returns [`ZakuraUpgradeOutcome::Rejected`] with
1025/// [`ZakuraRejectReason::TemporaryUnavailable`], so the caller keeps the legacy
1026/// connection instead.
1027async fn upgrade_to_zakura_handshake<PeerTransport>(
1028    peer_conn: &mut Framed<PeerTransport, Codec>,
1029    connection_info: &ConnectionInfo,
1030    connected_addr: ConnectedAddr,
1031    node_config: &Config,
1032    zakura_handshake_connector: Option<ZakuraHandshakeConnector>,
1033    peer_registry: Option<&PeerRegistry>,
1034    address_book_updater: &tokio::sync::mpsc::Sender<MetaAddrChange>,
1035) -> Result<ZakuraUpgradeOutcome, HandshakeError>
1036where
1037    PeerTransport: AsyncRead + AsyncWrite + Unpin + Send + 'static,
1038{
1039    let Some(connector) = zakura_handshake_connector else {
1040        metrics::counter!("zakura.p2p.handshake.upgrade.error").increment(1);
1041        return Err(HandshakeError::ZakuraUpgrade(
1042            crate::zakura::ZakuraUpgradeError::Unavailable,
1043        ));
1044    };
1045
1046    // Handshake routing tests inject a deterministic outcome instead of running
1047    // the real prelude exchange over a live Zakura endpoint.
1048    #[cfg(test)]
1049    if let Some(outcome) = connector.consume_test_outcome() {
1050        record_zakura_upgrade(
1051            &outcome,
1052            connection_info,
1053            connected_addr,
1054            &connector,
1055            peer_registry,
1056            address_book_updater,
1057        );
1058        return Ok(outcome);
1059    }
1060
1061    // Our own iroh dial hints. Without a live Zakura endpoint we cannot upgrade,
1062    // so we stay on the legacy connection.
1063    let Some((local_node_id, local_direct_addresses)) = connector.local_iroh_hints().await else {
1064        return Ok(neutral_upgrade_fallback());
1065    };
1066
1067    let addr_label = connected_addr.addr_label(node_config.expose_peer_addresses);
1068    let upgrade_config = ZakuraHandshakeConfig::for_network_with_dev_cohort(
1069        &node_config.network,
1070        node_config.zakura.dev_network.as_deref(),
1071    );
1072    let nonces = ZakuraLegacyNonces {
1073        local_zebra_nonce: connection_info.local.nonce,
1074        remote_zebra_nonce: connection_info.remote.nonce,
1075    };
1076
1077    // The side that opened the legacy TCP connection initiates the prelude and
1078    // dials over QUIC; the accepting side responds and is dialed.
1079    let outcome = if connected_addr.is_inbound() {
1080        run_responder_upgrade(
1081            peer_conn,
1082            &connector,
1083            &upgrade_config,
1084            nonces,
1085            local_node_id,
1086            local_direct_addresses,
1087            ResponderRegistrationWait::Production,
1088        )
1089        .await?
1090    } else {
1091        run_initiator_upgrade(
1092            peer_conn,
1093            &connector,
1094            &upgrade_config,
1095            nonces,
1096            local_node_id,
1097            local_direct_addresses,
1098        )
1099        .await?
1100    };
1101
1102    match &outcome {
1103        ZakuraUpgradeOutcome::Upgraded { peer_id, conn_id } => {
1104            // The success metric `zakura.p2p.handshake.upgraded` is incremented by
1105            // the supervisor when the dialed/accepted QUIC connection registers.
1106            info!(
1107                peer = %addr_label,
1108                connection_kind = connected_addr.get_short_kind_label(),
1109                ?peer_id,
1110                conn_id,
1111                remote_services = ?connection_info.remote.services,
1112                "upgraded mutually P2P-v2-capable peer to Zakura",
1113            );
1114        }
1115        ZakuraUpgradeOutcome::Duplicate { peer_id } => {
1116            info!(
1117                peer = %addr_label,
1118                connection_kind = connected_addr.get_short_kind_label(),
1119                ?peer_id,
1120                "closing duplicate Zakura peer neutrally"
1121            );
1122            metrics::counter!("zakura.p2p.handshake.duplicate").increment(1);
1123        }
1124        ZakuraUpgradeOutcome::Rejected { reason } => {
1125            debug!(
1126                peer = %addr_label,
1127                connection_kind = connected_addr.get_short_kind_label(),
1128                ?reason,
1129                "Zakura upgrade not completed; continuing on the legacy connection",
1130            );
1131            metrics::counter!(
1132                "zakura.p2p.upgrade.prelude.rejected",
1133                "reason" => format!("{reason:?}"),
1134                "network" => upgrade_config.network_label(),
1135            )
1136            .increment(1);
1137        }
1138    }
1139
1140    record_zakura_upgrade(
1141        &outcome,
1142        connection_info,
1143        connected_addr,
1144        &connector,
1145        peer_registry,
1146        address_book_updater,
1147    );
1148
1149    Ok(outcome)
1150}
1151
1152/// Record active native metadata and suppress redundant outbound legacy dials.
1153fn record_zakura_upgrade(
1154    outcome: &ZakuraUpgradeOutcome,
1155    connection_info: &ConnectionInfo,
1156    connected_addr: ConnectedAddr,
1157    connector: &ZakuraHandshakeConnector,
1158    peer_registry: Option<&PeerRegistry>,
1159    address_book_updater: &tokio::sync::mpsc::Sender<MetaAddrChange>,
1160) {
1161    let ZakuraUpgradeOutcome::Upgraded { peer_id, conn_id } = outcome else {
1162        return;
1163    };
1164
1165    if let (Some(peer_registry), Some(peer)) = (
1166        peer_registry,
1167        crate::ConnectedPeer::from_connection_info(connection_info),
1168    ) {
1169        // The maintained outbound dial reuses this metadata across transport reconnects.
1170        // An inbound upgrade binds its metadata only to the current connection.
1171        let retain_for_redial = !connected_addr.is_inbound();
1172        if !peer_registry.attach_native_metadata(peer_id.clone(), *conn_id, peer, retain_for_redial)
1173        {
1174            debug!(
1175                ?peer_id,
1176                conn_id, "native connection changed before metadata attachment",
1177            );
1178        }
1179    }
1180
1181    if connected_addr.is_inbound() {
1182        return;
1183    }
1184
1185    let Some(book_addr) = connected_addr.get_address_book_addr() else {
1186        return;
1187    };
1188
1189    // Keep the upgraded peer's legacy address-book entry live for as long as
1190    // its Zakura connection is registered. This prevents the outbound crawler
1191    // from re-running the upgrade.
1192    connector.spawn_legacy_liveness_keeper(
1193        peer_id.clone(),
1194        book_addr,
1195        address_book_updater.clone(),
1196    );
1197}
1198
1199/// The neutral upgrade fallback outcome: keep the legacy connection.
1200fn neutral_upgrade_fallback() -> ZakuraUpgradeOutcome {
1201    ZakuraUpgradeOutcome::Rejected {
1202        reason: ZakuraRejectReason::TemporaryUnavailable,
1203    }
1204}
1205
1206/// The TCP initiator side of the legacy Zakura upgrade prelude exchange.
1207///
1208/// Sends our [`P2pV2UpgradeInit`], reads the responder's [`P2pV2UpgradeAccept`],
1209/// and dials the responder's advertised Zakura node address over QUIC.
1210async fn run_initiator_upgrade<PeerTransport>(
1211    peer_conn: &mut Framed<PeerTransport, Codec>,
1212    connector: &ZakuraHandshakeConnector,
1213    config: &ZakuraHandshakeConfig,
1214    nonces: ZakuraLegacyNonces,
1215    local_node_id: Vec<u8>,
1216    local_direct_addresses: Vec<Vec<u8>>,
1217) -> Result<ZakuraUpgradeOutcome, HandshakeError>
1218where
1219    PeerTransport: AsyncRead + AsyncWrite + Unpin + Send + 'static,
1220{
1221    let mut upgrade_nonce = [0u8; 32];
1222    OsRng.fill_bytes(&mut upgrade_nonce);
1223
1224    let init = P2pV2UpgradeInit {
1225        magic: PRELUDE_MAGIC,
1226        prelude_version: config.prelude_version,
1227        zakura_protocol_min: config.zakura_protocol_min,
1228        zakura_protocol_max: config.zakura_protocol_max,
1229        network_id: config.network_id,
1230        chain_id: config.chain_id,
1231        capabilities: config.supported_capabilities,
1232        local_zebra_nonce: nonces.local_zebra_nonce,
1233        remote_zebra_nonce: nonces.remote_zebra_nonce,
1234        upgrade_nonce,
1235        iroh_node_id: local_node_id,
1236        iroh_direct_addresses: local_direct_addresses,
1237        iroh_relay_hint: None,
1238        max_control_frame_bytes: config.max_control_frame_bytes,
1239        max_open_streams: config.max_open_streams,
1240    };
1241
1242    let Ok(init_bytes) = P2pV2Upgrade::Init(init.clone()).encode() else {
1243        return Ok(neutral_upgrade_fallback());
1244    };
1245    peer_conn.send(Message::P2pV2Upgrade(init_bytes)).await?;
1246
1247    let Some(P2pV2Upgrade::Accept(accept)) = read_upgrade_prelude(peer_conn).await? else {
1248        // A well-formed reject, an unexpected variant, or no prelude within the
1249        // skip window: keep the legacy connection. A malformed prelude is not
1250        // reached here; `read_upgrade_prelude` disconnects on the first
1251        // malformed upgrade message instead of falling back.
1252        return Ok(neutral_upgrade_fallback());
1253    };
1254
1255    if accept.validate(config, nonces, &init).is_err() {
1256        return Ok(neutral_upgrade_fallback());
1257    }
1258
1259    let Ok(peer_id) = ZakuraPeerId::new(accept.iroh_node_id.clone()) else {
1260        return Ok(neutral_upgrade_fallback());
1261    };
1262
1263    // Dial the responder's Zakura endpoint over QUIC and wait for the local
1264    // supervisor to register a usable outbound handle before dropping the
1265    // legacy connection.
1266    let handoff = connector
1267        .spawn_zakura_dial_to_hints_and_wait(
1268            &peer_id,
1269            &accept.iroh_node_id,
1270            &accept.iroh_direct_addresses,
1271        )
1272        .await;
1273    match handoff {
1274        crate::zakura::ZakuraNativeHandoff::Registered(conn_id) => {
1275            Ok(ZakuraUpgradeOutcome::Upgraded { peer_id, conn_id })
1276        }
1277        crate::zakura::ZakuraNativeHandoff::Duplicate => {
1278            Ok(ZakuraUpgradeOutcome::Duplicate { peer_id })
1279        }
1280        crate::zakura::ZakuraNativeHandoff::Failed => Ok(neutral_upgrade_fallback()),
1281    }
1282}
1283
1284/// Selects the registration wait used by the responder upgrade path.
1285enum ResponderRegistrationWait {
1286    Production,
1287    #[cfg(test)]
1288    TestTimeout(Duration),
1289}
1290
1291impl ResponderRegistrationWait {
1292    async fn wait(
1293        self,
1294        connector: &ZakuraHandshakeConnector,
1295        peer_id: &ZakuraPeerId,
1296    ) -> crate::zakura::ZakuraNativeHandoff {
1297        match self {
1298            Self::Production => connector.wait_for_zakura_registration(peer_id).await,
1299            #[cfg(test)]
1300            Self::TestTimeout(timeout) => {
1301                tokio::time::timeout(timeout, connector.wait_for_zakura_registration(peer_id))
1302                    .await
1303                    .unwrap_or(crate::zakura::ZakuraNativeHandoff::Failed)
1304            }
1305        }
1306    }
1307}
1308
1309/// The TCP responder side of the legacy Zakura upgrade prelude exchange.
1310///
1311/// Reads the initiator's [`P2pV2UpgradeInit`] and replies with our
1312/// [`P2pV2UpgradeAccept`], advertising our Zakura node address so the initiator
1313/// can dial us. Our iroh router accepts that inbound dial separately.
1314async fn run_responder_upgrade<PeerTransport>(
1315    peer_conn: &mut Framed<PeerTransport, Codec>,
1316    connector: &ZakuraHandshakeConnector,
1317    config: &ZakuraHandshakeConfig,
1318    nonces: ZakuraLegacyNonces,
1319    local_node_id: Vec<u8>,
1320    local_direct_addresses: Vec<Vec<u8>>,
1321    registration_wait: ResponderRegistrationWait,
1322) -> Result<ZakuraUpgradeOutcome, HandshakeError>
1323where
1324    PeerTransport: AsyncRead + AsyncWrite + Unpin + Send + 'static,
1325{
1326    // A well-formed unexpected variant or no prelude within the skip window: send
1327    // a neutral reject and keep legacy. A malformed prelude is not reached here;
1328    // `read_upgrade_prelude` disconnects on the first malformed upgrade message
1329    // instead of replying with a neutral reject and falling back.
1330    let Some(P2pV2Upgrade::Init(init)) = read_upgrade_prelude(peer_conn).await? else {
1331        send_upgrade_reject(peer_conn, config).await?;
1332        return Ok(neutral_upgrade_fallback());
1333    };
1334
1335    let selected_zakura_protocol = match init.validate(config, nonces) {
1336        Ok(selected) => selected,
1337        Err(_) => {
1338            send_upgrade_reject(peer_conn, config).await?;
1339            return Ok(neutral_upgrade_fallback());
1340        }
1341    };
1342
1343    let mut responder_upgrade_nonce = [0u8; 32];
1344    OsRng.fill_bytes(&mut responder_upgrade_nonce);
1345
1346    let accept = P2pV2UpgradeAccept {
1347        magic: PRELUDE_MAGIC,
1348        prelude_version: config.prelude_version,
1349        selected_zakura_protocol,
1350        network_id: config.network_id,
1351        chain_id: config.chain_id,
1352        capabilities: config.supported_capabilities,
1353        initiator_upgrade_nonce: init.upgrade_nonce,
1354        responder_upgrade_nonce,
1355        local_zebra_nonce: nonces.local_zebra_nonce,
1356        remote_zebra_nonce: nonces.remote_zebra_nonce,
1357        iroh_node_id: local_node_id,
1358        iroh_direct_addresses: local_direct_addresses,
1359        iroh_relay_hint: None,
1360        max_control_frame_bytes: config.max_control_frame_bytes,
1361        max_open_streams: config.max_open_streams,
1362    };
1363
1364    let Ok(accept_bytes) = P2pV2Upgrade::Accept(accept).encode() else {
1365        return Ok(neutral_upgrade_fallback());
1366    };
1367    peer_conn.send(Message::P2pV2Upgrade(accept_bytes)).await?;
1368
1369    let Ok(peer_id) = ZakuraPeerId::new(init.iroh_node_id.clone()) else {
1370        return Ok(neutral_upgrade_fallback());
1371    };
1372
1373    // The peer dials our advertised Zakura endpoint over QUIC after receiving
1374    // `Accept`, and our iroh router registers that inbound connection
1375    // separately. Wait for that native registration before reporting the
1376    // upgrade: the outer handshake drops the legacy TCP connection on
1377    // `Upgraded`, so without this wait an inbound peer that sends a valid `Init`
1378    // and then never completes the native dial would make us discard a working
1379    // legacy connection with no Zakura replacement. This mirrors the initiator's
1380    // `spawn_zakura_dial_to_hints_and_wait` hand-off wait.
1381    match registration_wait.wait(connector, &peer_id).await {
1382        crate::zakura::ZakuraNativeHandoff::Registered(conn_id) => {
1383            Ok(ZakuraUpgradeOutcome::Upgraded { peer_id, conn_id })
1384        }
1385        crate::zakura::ZakuraNativeHandoff::Duplicate => {
1386            Ok(ZakuraUpgradeOutcome::Duplicate { peer_id })
1387        }
1388        crate::zakura::ZakuraNativeHandoff::Failed => Ok(neutral_upgrade_fallback()),
1389    }
1390}
1391
1392/// Sends a neutral [`P2pV2UpgradeReject`] so the peer stops waiting for an accept
1393/// and falls back to the legacy connection.
1394async fn send_upgrade_reject<PeerTransport>(
1395    peer_conn: &mut Framed<PeerTransport, Codec>,
1396    config: &ZakuraHandshakeConfig,
1397) -> Result<(), HandshakeError>
1398where
1399    PeerTransport: AsyncRead + AsyncWrite + Unpin + Send + 'static,
1400{
1401    let reject = P2pV2UpgradeReject {
1402        magic: PRELUDE_MAGIC,
1403        prelude_version: config.prelude_version,
1404        reason: ZakuraRejectReason::TemporaryUnavailable,
1405    };
1406    if let Ok(reject_bytes) = P2pV2Upgrade::Reject(reject).encode() {
1407        peer_conn.send(Message::P2pV2Upgrade(reject_bytes)).await?;
1408    }
1409    Ok(())
1410}
1411
1412/// Decode a peer-controlled Zakura upgrade prelude.
1413///
1414/// A panic is terminal for this legacy connection: callers propagate the
1415/// returned serialization error and drop the framed transport without reusing
1416/// its parser state.
1417fn decode_upgrade_prelude(payload: &[u8]) -> Result<P2pV2Upgrade, HandshakeError> {
1418    decode_upgrade_prelude_with(|| P2pV2Upgrade::decode(payload))
1419}
1420
1421/// Run the structured upgrade prelude decoder inside its panic boundary.
1422fn decode_upgrade_prelude_with(
1423    decode: impl FnOnce() -> Result<P2pV2Upgrade, ZakuraProtocolError> + panic::UnwindSafe,
1424) -> Result<P2pV2Upgrade, HandshakeError> {
1425    match panic::catch_unwind(decode) {
1426        Ok(Ok(prelude)) => Ok(prelude),
1427        Ok(Err(error)) => {
1428            metrics::counter!("zakura.p2p.upgrade.prelude.malformed").increment(1);
1429            Err(HandshakeError::ZakuraUpgradePreludeMalformed(error))
1430        }
1431        Err(_panic_payload) => {
1432            metrics::counter!(
1433                "peer.message.parse.panics",
1434                "parser" => "upgrade_prelude",
1435            )
1436            .increment(1);
1437            tracing::error!(
1438                command = P2P_V2_UPGRADE_COMMAND,
1439                "peer-controlled message parser panicked; disconnecting legacy peer"
1440            );
1441            Err(SerializationError::Parse("Zakura P2P v2 upgrade prelude parser panicked").into())
1442        }
1443    }
1444}
1445
1446/// Reads the next legacy [`P2pV2Upgrade`] prelude from the peer.
1447///
1448/// Skips a small bounded number of unrelated messages (the overall handshake
1449/// timeout also applies), so a peer cannot stall the upgrade by streaming other
1450/// messages.
1451///
1452/// Returns `Ok(None)` only when no prelude arrives within the skip bound: a peer
1453/// that advertised `NODE_P2P_V2` but never frames a `p2pv2up` message is treated
1454/// as a neutral legacy fallback (compatibility).
1455///
1456/// In contrast, a peer that *does* frame a `p2pv2up` message whose payload fails
1457/// to decode has violated the upgrade protocol, so this returns
1458/// [`HandshakeError::ZakuraUpgradePreludeMalformed`] rather than erasing the
1459/// decode error to `None`. Erasing it would let a peer force a downgrade to
1460/// legacy by sending malformed upgrade bytes (SR-7 fail-closed); surfacing it
1461/// disconnects the peer on the first malformed upgrade message.
1462async fn read_upgrade_prelude<PeerTransport>(
1463    peer_conn: &mut Framed<PeerTransport, Codec>,
1464) -> Result<Option<P2pV2Upgrade>, HandshakeError>
1465where
1466    PeerTransport: AsyncRead + AsyncWrite + Unpin + Send + 'static,
1467{
1468    // Bound on unrelated messages tolerated before the prelude.
1469    const MAX_SKIPPED_MESSAGES: usize = 4;
1470
1471    for _ in 0..MAX_SKIPPED_MESSAGES {
1472        let message = peer_conn
1473            .next()
1474            .await
1475            .ok_or(HandshakeError::ConnectionClosed)??;
1476        if let Message::P2pV2Upgrade(payload) = message {
1477            return decode_upgrade_prelude(&payload).map(Some);
1478        }
1479    }
1480
1481    Ok(None)
1482}
1483
1484/// A handshake request.
1485/// Contains the information needed to handshake with the peer.
1486pub struct HandshakeRequest<PeerTransport>
1487where
1488    PeerTransport: AsyncRead + AsyncWrite + Unpin + Send + 'static,
1489{
1490    /// The tokio [`TcpStream`](tokio::net::TcpStream) or Tor
1491    /// `arti_client::DataStream` to the peer.
1492    // Use [`arti_client::DataStream`] when #5492 is done.
1493    pub data_stream: PeerTransport,
1494
1495    /// The address of the peer, and other related information.
1496    pub connected_addr: ConnectedAddr,
1497
1498    /// A connection tracker that reduces the open connection count when dropped.
1499    ///
1500    /// Used to limit the number of open connections in Zebra.
1501    pub connection_tracker: ConnectionTracker,
1502}
1503
1504impl<S, PeerTransport, C> Service<HandshakeRequest<PeerTransport>> for Handshake<S, C>
1505where
1506    S: Service<Request, Response = Response, Error = BoxError> + Clone + Send + 'static,
1507    S::Future: Send,
1508    C: ChainTip + Clone + Send + 'static,
1509    PeerTransport: AsyncRead + AsyncWrite + Unpin + Send + 'static,
1510{
1511    type Response = Client;
1512    type Error = BoxError;
1513    type Future =
1514        Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;
1515
1516    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
1517        Poll::Ready(Ok(()))
1518    }
1519
1520    fn call(&mut self, req: HandshakeRequest<PeerTransport>) -> Self::Future {
1521        let HandshakeRequest {
1522            data_stream,
1523            connected_addr,
1524            connection_tracker,
1525        } = req;
1526
1527        let addr_label = connected_addr.addr_label(self.config.expose_peer_addresses);
1528        let negotiator_span = debug_span!(
1529            "negotiator",
1530            peer = %addr_label,
1531            connection_kind = connected_addr.get_short_kind_label(),
1532        );
1533        // set the peer connection span's parent to the global span, as it
1534        // should exist independently of its creation source (inbound
1535        // connection, crawler, initial peer, ...)
1536        let connection_span = span!(
1537            parent: &self.parent_span,
1538            Level::INFO,
1539            "",
1540            peer = %addr_label,
1541            connection_kind = connected_addr.get_short_kind_label(),
1542        );
1543
1544        // Clone these upfront, so they can be moved into the future.
1545        let nonces = self.nonces.clone();
1546        let inbound_service = self.inbound_service.clone();
1547        let address_book_updater = self.address_book_updater.clone();
1548        let inv_collector = self.inv_collector.clone();
1549        let config = self.config.clone();
1550        let user_agent = self.user_agent.clone();
1551        let our_services = self.our_services;
1552        let relay = self.relay;
1553        let minimum_peer_version = self.minimum_peer_version.clone();
1554        let zakura_handshake_connector = self.zakura_handshake_connector.clone();
1555        let peer_registry = self.peer_registry.clone();
1556
1557        // Whether this peer is exempt from the inbound-overload connection drop.
1558        // Computed here (not in the future) so only the resulting `bool` is moved
1559        // into the handshake task.
1560        let is_protected_peer = connected_addr.is_protected_peer(&self.protected_peer_ips);
1561
1562        // # Security
1563        //
1564        // `zakura_network::init()` implements a connection timeout on this future.
1565        // Any code outside this future does not have a timeout.
1566        let fut = async move {
1567            debug!(
1568                addr = %addr_label,
1569                connection_kind = connected_addr.get_short_kind_label(),
1570                "negotiating protocol version with remote peer"
1571            );
1572
1573            // Start timing the handshake for metrics
1574            let handshake_start = Instant::now();
1575
1576            let mut peer_conn = Framed::new(
1577                data_stream,
1578                Codec::builder().for_network(&config.network).finish(),
1579            );
1580            let mut connection_tracker = connection_tracker;
1581
1582            let connection_info = match negotiate_version(
1583                &mut peer_conn,
1584                &connected_addr,
1585                config.clone(),
1586                nonces,
1587                user_agent,
1588                our_services,
1589                relay,
1590                minimum_peer_version,
1591                is_protected_peer,
1592            )
1593            .await
1594            {
1595                Ok(info) => {
1596                    // Record successful handshake duration
1597                    let duration = handshake_start.elapsed().as_secs_f64();
1598                    metrics::histogram!(
1599                        "zcash.net.peer.handshake.duration_seconds",
1600                        "result" => "success"
1601                    )
1602                    .record(duration);
1603                    info
1604                }
1605                Err(err) => {
1606                    // Record failed handshake duration and failure reason
1607                    let duration = handshake_start.elapsed().as_secs_f64();
1608                    let reason = match &err {
1609                        HandshakeError::UnexpectedMessage(_) => "unexpected_message",
1610                        HandshakeError::RemoteNonceReuse => "nonce_reuse",
1611                        HandshakeError::LocalDuplicateNonce => "duplicate_nonce",
1612                        HandshakeError::ConnectionClosed => "connection_closed",
1613                        HandshakeError::Io(_) => "io_error",
1614                        HandshakeError::Serialization(_) => "serialization",
1615                        HandshakeError::ObsoleteVersion(_) => "obsolete_version",
1616                        HandshakeError::Timeout => "timeout",
1617                        HandshakeError::ZakuraUpgradeSelected
1618                        | HandshakeError::ZakuraUpgrade(_)
1619                        | HandshakeError::ZakuraUpgradePreludeMalformed(_) => {
1620                            unreachable!("negotiate_version returns before Zakura upgrade routing")
1621                        }
1622                    };
1623                    metrics::histogram!(
1624                        "zcash.net.peer.handshake.duration_seconds",
1625                        "result" => "failure"
1626                    )
1627                    .record(duration);
1628                    metrics::counter!(
1629                        "zcash.net.peer.handshake.failures.total",
1630                        "reason" => reason
1631                    )
1632                    .increment(1);
1633                    return Err(err);
1634                }
1635            };
1636
1637            let remote_services = connection_info.remote.services;
1638
1639            if should_attempt_zakura_upgrade(&config, &connection_info) {
1640                match upgrade_to_zakura_handshake(
1641                    &mut peer_conn,
1642                    &connection_info,
1643                    connected_addr,
1644                    &config,
1645                    zakura_handshake_connector,
1646                    peer_registry.as_ref(),
1647                    &address_book_updater,
1648                )
1649                .await
1650                {
1651                    Ok(
1652                        ZakuraUpgradeOutcome::Upgraded { .. }
1653                        | ZakuraUpgradeOutcome::Duplicate { .. },
1654                    ) => {
1655                        // Returning here drops the legacy stream and connection tracker, cleanly
1656                        // closing the Zebra path and releasing the connection limit exactly once.
1657                        return Err(HandshakeError::ZakuraUpgradeSelected);
1658                    }
1659                    Ok(ZakuraUpgradeOutcome::Rejected {
1660                        reason: ZakuraRejectReason::TemporaryUnavailable,
1661                    }) => {
1662                        debug!(
1663                            peer = %addr_label,
1664                            connection_kind = connected_addr.get_short_kind_label(),
1665                            "Zakura upgrade is temporarily unavailable; continuing legacy handshake"
1666                        );
1667                    }
1668                    Ok(ZakuraUpgradeOutcome::Rejected { .. }) => {
1669                        // Returning here drops the legacy stream and connection tracker, cleanly
1670                        // closing the Zebra path and releasing the connection limit exactly once.
1671                        return Err(HandshakeError::ZakuraUpgradeSelected);
1672                    }
1673                    Err(error) => return Err(error),
1674                }
1675            }
1676
1677            // The handshake succeeded: update the peer status from AttemptPending to Responded,
1678            // send initial connection info, and update the active connection counter.
1679            connection_tracker.mark_open();
1680            if let Some(book_addr) = connected_addr.get_address_book_addr() {
1681                // the collector doesn't depend on network activity,
1682                // so this await should not hang
1683                let _ = address_book_updater
1684                    .send(MetaAddr::new_connected(
1685                        book_addr,
1686                        &remote_services,
1687                        connected_addr.is_inbound(),
1688                    ))
1689                    .await;
1690            }
1691
1692            let (peer_registry_guard, peer_registry_updater) = peer_registry
1693                .as_ref()
1694                .and_then(|registry| {
1695                    crate::ConnectedPeer::from_connection_info(&connection_info)
1696                        .map(|peer| registry.register_legacy(peer))
1697                })
1698                .unzip();
1699
1700            // Reconfigure the codec to use the negotiated version.
1701            //
1702            // TODO: The tokio documentation says not to do this while any frames are still being processed.
1703            // Since we don't know that here, another way might be to release the tcp
1704            // stream from the unversioned Framed wrapper and construct a new one with a versioned codec.
1705            let bare_codec = peer_conn.codec_mut();
1706            bare_codec.reconfigure_version(connection_info.negotiated_version);
1707            bare_codec.reconfigure_full_body_len();
1708
1709            debug!("constructing client, spawning server");
1710
1711            // These channels communicate between the inbound and outbound halves of the connection,
1712            // and between the different connection tasks. We create separate tasks and channels
1713            // for each new connection.
1714            let (server_tx, server_rx) = futures::channel::mpsc::channel(0);
1715            let (shutdown_tx, shutdown_rx) = oneshot::channel();
1716            let error_slot = ErrorSlot::default();
1717
1718            let (peer_tx, peer_rx) = peer_conn.split();
1719
1720            // Instrument the peer's rx and tx streams.
1721
1722            let inner_conn_span = connection_span.clone();
1723            let peer_tx = peer_tx.with(move |msg: Message| {
1724                let span = debug_span!(parent: inner_conn_span.clone(), "outbound_metric");
1725                // Add a metric for outbound messages.
1726                metrics::counter!(
1727                    "zcash.net.out.messages",
1728                    "command" => msg.command(),
1729                )
1730                .increment(1);
1731                // We need to use future::ready rather than an async block here,
1732                // because we need the sink to be Unpin, and the With<Fut, ...>
1733                // returned by .with is Unpin only if Fut is Unpin, and the
1734                // futures generated by async blocks are not Unpin.
1735                future::ready(Ok(msg)).instrument(span)
1736            });
1737
1738            // CORRECTNESS
1739            //
1740            // Ping/Pong messages and every error must update the peer address state via
1741            // the inbound_ts_collector.
1742            //
1743            // The heartbeat task sends regular Ping/Pong messages,
1744            // and it ends the connection if the heartbeat times out.
1745            // So we can just track peer activity based on Ping and Pong.
1746            // (This significantly improves performance, by reducing time system calls.)
1747            let inbound_ts_collector = address_book_updater.clone();
1748            let inbound_inv_collector = inv_collector.clone();
1749            let ts_inner_conn_span = connection_span.clone();
1750            let inv_inner_conn_span = connection_span.clone();
1751            let peer_rx = peer_rx
1752                .then(move |msg| {
1753                    // Add a metric for inbound messages and errors.
1754                    // Fire a timestamp or failure event.
1755                    let inbound_ts_collector = inbound_ts_collector.clone();
1756                    let span =
1757                        debug_span!(parent: ts_inner_conn_span.clone(), "inbound_ts_collector");
1758
1759                    async move {
1760                        match &msg {
1761                            Ok(msg) => {
1762                                metrics::counter!(
1763                                    "zcash.net.in.messages",
1764                                    "command" => msg.command(),
1765                                )
1766                                .increment(1);
1767
1768                                // # Security
1769                                //
1770                                // Peer messages are not rate-limited, so we can't send anything
1771                                // to a shared channel or do anything expensive here.
1772                            }
1773                            Err(err) => {
1774                                metrics::counter!(
1775                                    "zakura.net.in.errors",
1776                                    "error" => err.to_string(),
1777                                )
1778                                .increment(1);
1779
1780                                // # Security
1781                                //
1782                                // Peer errors are rate-limited because:
1783                                // - opening connections is rate-limited
1784                                // - the number of connections is limited
1785                                // - after the first error, the peer is disconnected
1786                                if let Some(book_addr) = connected_addr.get_address_book_addr() {
1787                                    let change = inbound_error_address_change(
1788                                        book_addr,
1789                                        remote_services,
1790                                        err,
1791                                    );
1792                                    let _ = inbound_ts_collector.send(change).await;
1793                                }
1794                            }
1795                        }
1796                        msg
1797                    }
1798                    .instrument(span)
1799                })
1800                .then(move |msg| {
1801                    let inbound_inv_collector = inbound_inv_collector.clone();
1802                    let span = debug_span!(parent: inv_inner_conn_span.clone(), "inventory_filter");
1803                    register_inventory_status(msg, connected_addr, inbound_inv_collector)
1804                        .instrument(span)
1805                })
1806                .boxed();
1807
1808            // If we've learned potential peer addresses from the inbound connection remote address
1809            // or the handshake version message, add those addresses to the peer cache for this
1810            // peer.
1811            //
1812            // # Security
1813            //
1814            // We can't add these alternate addresses directly to the address book. If we did,
1815            // malicious peers could interfere with the address book state of other peers by
1816            // providing their addresses in `Version` messages. Or they could fill the address book
1817            // with fake addresses.
1818            //
1819            // These peer addresses are rate-limited because:
1820            // - opening connections is rate-limited
1821            // - these addresses are put in the peer address cache
1822            // - the peer address cache is only used when Zebra requests addresses from that peer
1823            let remote_canonical_addr = connection_info.remote.address_from.addr();
1824            let alternate_addrs = connected_addr
1825                .get_alternate_addrs(remote_canonical_addr)
1826                .map(|addr| {
1827                    // Assume the connecting node is a server node, and it's available now.
1828                    MetaAddr::new_gossiped_meta_addr(
1829                        addr,
1830                        PeerServices::NODE_NETWORK,
1831                        DateTime32::now(),
1832                    )
1833                });
1834
1835            let server = Connection::new(
1836                inbound_service,
1837                server_rx,
1838                error_slot.clone(),
1839                peer_tx,
1840                connection_tracker,
1841                connection_info.clone(),
1842                addr_label,
1843                alternate_addrs.collect(),
1844            );
1845
1846            let connection_task = tokio::spawn(
1847                async move {
1848                    let _peer_registry_guard = peer_registry_guard;
1849                    server.run(peer_rx).await;
1850                }
1851                .instrument(connection_span.clone())
1852                .boxed(),
1853            );
1854
1855            let heartbeat_task = tokio::spawn(
1856                send_periodic_heartbeats_with_shutdown_handle(
1857                    connected_addr,
1858                    shutdown_rx,
1859                    server_tx.clone(),
1860                    address_book_updater.clone(),
1861                    peer_registry_updater,
1862                )
1863                .instrument(tracing::debug_span!(parent: connection_span, "heartbeat"))
1864                .boxed(),
1865            );
1866
1867            let client = Client {
1868                connection_info,
1869                shutdown_tx: Some(shutdown_tx),
1870                server_tx,
1871                inv_collector,
1872                error_slot,
1873                connection_task,
1874                heartbeat_task,
1875            };
1876
1877            Ok(client)
1878        };
1879
1880        // Correctness: As a defence-in-depth against hangs, wrap the entire handshake in a timeout.
1881        let fut = timeout(constants::HANDSHAKE_TIMEOUT, fut);
1882
1883        // Spawn a new task to drive this handshake, forwarding panics to the calling task.
1884        tokio::spawn(fut.instrument(negotiator_span))
1885            .map(
1886                |join_result: Result<
1887                    Result<Result<Client, HandshakeError>, error::Elapsed>,
1888                    JoinError,
1889                >| {
1890                    match join_result {
1891                        Ok(Ok(Ok(connection_client))) => Ok(connection_client),
1892                        Ok(Ok(Err(handshake_error))) => Err(handshake_error.into()),
1893                        Ok(Err(timeout_error)) => Err(timeout_error.into()),
1894                        Err(join_error) => match join_error.try_into_panic() {
1895                            // Forward panics to the calling task
1896                            Ok(panic_reason) => panic::resume_unwind(panic_reason),
1897                            Err(join_error) => Err(join_error.into()),
1898                        },
1899                    }
1900                },
1901            )
1902            .boxed()
1903    }
1904}
1905
1906/// Register any advertised or missing inventory in `msg` for `connected_addr`.
1907pub(crate) async fn register_inventory_status(
1908    msg: Result<Message, SerializationError>,
1909    connected_addr: ConnectedAddr,
1910    inv_collector: broadcast::Sender<InventoryChange>,
1911) -> Result<Message, SerializationError> {
1912    match (&msg, connected_addr.get_transient_addr()) {
1913        (Ok(Message::Inv(advertised)), Some(transient_addr)) => {
1914            // We ignore inventory messages with more than one
1915            // block, because they are most likely replies to a
1916            // query, rather than a newly gossiped block.
1917            //
1918            // (We process inventory messages with any number of
1919            // transactions.)
1920            //
1921            // https://zebra.zfnd.org/dev/rfcs/0003-inventory-tracking.html#inventory-monitoring
1922            //
1923            // Note: zcashd has a bug where it merges queued inv messages of
1924            // the same or different types. Zebra compensates by sending `notfound`
1925            // responses to the inv collector. (#2156, #1768)
1926            //
1927            // (We can't split `inv`s, because that fills the inventory registry
1928            // with useless entries that the whole network has, making it large and slow.)
1929            match advertised.as_slice() {
1930                [advertised @ InventoryHash::Block(_)] => {
1931                    debug!(
1932                        ?advertised,
1933                        "registering gossiped advertised block inventory for peer"
1934                    );
1935
1936                    // The peer set and inv collector use the peer's remote
1937                    // address as an identifier
1938                    // If all receivers have been dropped, `send` returns an error.
1939                    // When that happens, Zebra is shutting down, so we want to ignore this error.
1940                    let _ = inv_collector
1941                        .send(InventoryChange::new_available(*advertised, transient_addr));
1942                }
1943                advertised => {
1944                    let advertised = advertised
1945                        .iter()
1946                        .filter(|advertised| advertised.unmined_tx_id().is_some());
1947
1948                    debug!(
1949                        ?advertised,
1950                        "registering advertised unmined transaction inventory for peer",
1951                    );
1952
1953                    if let Some(change) =
1954                        InventoryChange::new_available_multi(advertised, transient_addr)
1955                    {
1956                        // Ignore channel errors that should only happen during shutdown.
1957                        let _ = inv_collector.send(change);
1958                    }
1959                }
1960            }
1961        }
1962
1963        (Ok(Message::NotFound(missing)), Some(transient_addr)) => {
1964            // Ignore Errors and the unsupported FilteredBlock type
1965            let missing = missing.iter().filter(|missing| {
1966                missing.unmined_tx_id().is_some() || missing.block_hash().is_some()
1967            });
1968
1969            debug!(?missing, "registering missing inventory for peer");
1970
1971            if let Some(change) = InventoryChange::new_missing_multi(missing, transient_addr) {
1972                let _ = inv_collector.send(change);
1973            }
1974        }
1975        _ => {}
1976    }
1977
1978    msg
1979}
1980
1981/// Send periodical heartbeats to `server_tx`, and update the peer status through
1982/// `heartbeat_ts_collector`.
1983///
1984/// # Correctness
1985///
1986/// To prevent hangs:
1987/// - every await that depends on the network must have a timeout (or interval)
1988/// - every error/shutdown must update the address book state and return
1989///
1990/// The address book state can be updated via `ClientRequest.tx`, or the
1991/// heartbeat_ts_collector.
1992///
1993/// Returning from this function terminates the connection's heartbeat task.
1994async fn send_periodic_heartbeats_with_shutdown_handle(
1995    connected_addr: ConnectedAddr,
1996    shutdown_rx: oneshot::Receiver<CancelHeartbeatTask>,
1997    server_tx: futures::channel::mpsc::Sender<ClientRequest>,
1998    heartbeat_ts_collector: tokio::sync::mpsc::Sender<MetaAddrChange>,
1999    peer_registry_updater: Option<crate::peer_registry::PeerRegistryUpdater>,
2000) -> Result<(), BoxError> {
2001    use futures::future::Either;
2002
2003    let heartbeat_run_loop = send_periodic_heartbeats_run_loop(
2004        connected_addr,
2005        server_tx,
2006        heartbeat_ts_collector.clone(),
2007        peer_registry_updater,
2008    );
2009
2010    pin_mut!(shutdown_rx);
2011    pin_mut!(heartbeat_run_loop);
2012
2013    // CORRECTNESS
2014    //
2015    // Currently, select prefers the first future if multiple
2016    // futures are ready.
2017    //
2018    // Starvation is impossible here, because interval has a
2019    // slow rate, and shutdown is a oneshot. If both futures
2020    // are ready, we want the shutdown to take priority over
2021    // sending a useless heartbeat.
2022    match future::select(shutdown_rx, heartbeat_run_loop).await {
2023        Either::Left((Ok(CancelHeartbeatTask), _unused_run_loop)) => {
2024            tracing::trace!("shutting down because Client requested shut down");
2025            handle_heartbeat_shutdown(
2026                PeerError::ClientCancelledHeartbeatTask,
2027                &heartbeat_ts_collector,
2028                &connected_addr,
2029            )
2030            .await
2031        }
2032        Either::Left((Err(oneshot::Canceled), _unused_run_loop)) => {
2033            tracing::trace!("shutting down because Client was dropped");
2034            handle_heartbeat_shutdown(
2035                PeerError::ClientDropped,
2036                &heartbeat_ts_collector,
2037                &connected_addr,
2038            )
2039            .await
2040        }
2041        Either::Right((result, _unused_shutdown)) => {
2042            tracing::trace!("shutting down due to heartbeat failure");
2043            // heartbeat_timeout() already send an error on the timestamp collector channel
2044
2045            result
2046        }
2047    }
2048}
2049
2050/// Send periodical heartbeats to `server_tx`, and update the peer status through
2051/// `heartbeat_ts_collector`.
2052///
2053/// See `send_periodic_heartbeats_with_shutdown_handle` for details.
2054async fn send_periodic_heartbeats_run_loop(
2055    connected_addr: ConnectedAddr,
2056    mut server_tx: futures::channel::mpsc::Sender<ClientRequest>,
2057    heartbeat_ts_collector: tokio::sync::mpsc::Sender<MetaAddrChange>,
2058    peer_registry_updater: Option<crate::peer_registry::PeerRegistryUpdater>,
2059) -> Result<(), BoxError> {
2060    // Don't send the first heartbeat immediately - we've just completed the handshake!
2061    let mut interval = tokio::time::interval_at(
2062        Instant::now() + constants::HEARTBEAT_INTERVAL,
2063        constants::HEARTBEAT_INTERVAL,
2064    );
2065    // If the heartbeat is delayed, also delay all future heartbeats.
2066    // (Shorter heartbeat intervals just add load, without any benefit.)
2067    interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
2068
2069    let mut interval_stream = IntervalStream::new(interval);
2070
2071    while let Some(_instant) = interval_stream.next().await {
2072        // We've reached another heartbeat interval without
2073        // shutting down, so do a heartbeat request.
2074        let ping_sent_at = Instant::now();
2075        if let Some(peer_registry_updater) = &peer_registry_updater {
2076            peer_registry_updater.record_ping_sent(ping_sent_at.into_std());
2077        }
2078        if let Some(book_addr) = connected_addr.get_address_book_addr() {
2079            let _ = heartbeat_ts_collector
2080                .send(MetaAddr::new_ping_sent(book_addr, ping_sent_at.into()))
2081                .await;
2082        }
2083
2084        let heartbeat = send_one_heartbeat(&mut server_tx);
2085        let rtt = heartbeat_timeout(heartbeat, &heartbeat_ts_collector, &connected_addr).await?;
2086
2087        if let (Some(peer_registry_updater), Some(rtt)) = (&peer_registry_updater, rtt) {
2088            peer_registry_updater.record_response(rtt);
2089        }
2090
2091        // # Security
2092        //
2093        // Peer heartbeats are rate-limited because:
2094        // - opening connections is rate-limited
2095        // - the number of connections is limited
2096        // - Zebra initiates each heartbeat using a timer
2097        if let Some(book_addr) = connected_addr.get_address_book_addr() {
2098            if let Some(rtt) = rtt {
2099                // the collector doesn't depend on network activity,
2100                // so this await should not hang
2101                let _ = heartbeat_ts_collector
2102                    .send(MetaAddr::new_responded(book_addr, Some(rtt)))
2103                    .await;
2104            }
2105        }
2106    }
2107
2108    unreachable!("unexpected IntervalStream termination")
2109}
2110
2111/// Send one heartbeat using `server_tx`.
2112async fn send_one_heartbeat(
2113    server_tx: &mut futures::channel::mpsc::Sender<ClientRequest>,
2114) -> Result<Response, BoxError> {
2115    // We just reached a heartbeat interval, so start sending
2116    // a heartbeat.
2117    let (tx, rx) = oneshot::channel();
2118
2119    // Try to send the heartbeat request
2120    let request = Request::Ping(Nonce::default());
2121    tracing::trace!(?request, "queueing heartbeat request");
2122    match server_tx.try_send(ClientRequest {
2123        request,
2124        tx,
2125        // we're not requesting inventory, so we don't need to update the registry
2126        inv_collector: None,
2127        transient_addr: None,
2128        span: tracing::Span::current(),
2129    }) {
2130        Ok(()) => {}
2131        Err(e) => {
2132            if e.is_disconnected() {
2133                Err(PeerError::ConnectionClosed)?;
2134            } else if e.is_full() {
2135                // Send the message when the Client becomes ready.
2136                // If sending takes too long, the heartbeat timeout will elapse
2137                // and close the connection, reducing our load to busy peers.
2138                server_tx.send(e.into_inner()).await?;
2139            } else {
2140                // we need to map unexpected error types to PeerErrors
2141                warn!(?e, "unexpected try_send error");
2142                Err(e)?;
2143            };
2144        }
2145    }
2146
2147    // Flush the heartbeat request from the queue
2148    server_tx.flush().await?;
2149    tracing::trace!("sent heartbeat request");
2150
2151    // Heartbeats are checked internally to the
2152    // connection logic, but we need to wait on the
2153    // response to avoid canceling the request.
2154    let response = rx.await??;
2155    tracing::trace!(?response, "got heartbeat response");
2156
2157    Ok(response)
2158}
2159
2160/// Wrap `fut` in a timeout, handing any inner or outer errors using
2161/// `handle_heartbeat_error`.
2162async fn heartbeat_timeout(
2163    fut: impl Future<Output = Result<Response, BoxError>>,
2164    address_book_updater: &tokio::sync::mpsc::Sender<MetaAddrChange>,
2165    connected_addr: &ConnectedAddr,
2166) -> Result<Option<Duration>, BoxError> {
2167    let response = match timeout(constants::HEARTBEAT_INTERVAL, fut).await {
2168        Ok(inner_result) => {
2169            handle_heartbeat_error(inner_result, address_book_updater, connected_addr).await?
2170        }
2171        Err(elapsed) => {
2172            handle_heartbeat_error(Err(elapsed), address_book_updater, connected_addr).await?
2173        }
2174    };
2175
2176    let rtt = match response {
2177        Response::Pong(rtt) => Some(rtt),
2178        _ => None,
2179    };
2180
2181    Ok(rtt)
2182}
2183
2184/// If `result.is_err()`, mark `connected_addr` as failed using `address_book_updater`.
2185async fn handle_heartbeat_error<T, E>(
2186    result: Result<T, E>,
2187    address_book_updater: &tokio::sync::mpsc::Sender<MetaAddrChange>,
2188    connected_addr: &ConnectedAddr,
2189) -> Result<T, E>
2190where
2191    E: std::fmt::Debug,
2192{
2193    match result {
2194        Ok(t) => Ok(t),
2195        Err(err) => {
2196            tracing::debug!(?err, "heartbeat error, shutting down");
2197
2198            // # Security
2199            //
2200            // Peer errors and shutdowns are rate-limited because:
2201            // - opening connections is rate-limited
2202            // - the number of connections is limited
2203            // - after the first error or shutdown, the peer is disconnected
2204            if let Some(book_addr) = connected_addr.get_address_book_addr() {
2205                let _ = address_book_updater
2206                    .send(MetaAddr::new_errored(book_addr, None))
2207                    .await;
2208            }
2209            Err(err)
2210        }
2211    }
2212}
2213
2214/// Mark `connected_addr` as shut down using `address_book_updater`.
2215async fn handle_heartbeat_shutdown(
2216    peer_error: PeerError,
2217    address_book_updater: &tokio::sync::mpsc::Sender<MetaAddrChange>,
2218    connected_addr: &ConnectedAddr,
2219) -> Result<(), BoxError> {
2220    tracing::debug!(?peer_error, "client shutdown, shutting down heartbeat");
2221
2222    if let Some(book_addr) = connected_addr.get_address_book_addr() {
2223        let _ = address_book_updater
2224            .send(MetaAddr::new_shutdown(book_addr))
2225            .await;
2226    }
2227
2228    Err(peer_error.into())
2229}