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