diff --git a/zebra-network/src/peer/client.rs b/zebra-network/src/peer/client.rs index f890ac88..239103ed 100644 --- a/zebra-network/src/peer/client.rs +++ b/zebra-network/src/peer/client.rs @@ -70,7 +70,7 @@ impl Service for PeerClient { } } Ok(()) => { - // The reciever end of the oneshot is itself a future. + // The receiver end of the oneshot is itself a future. rx.map(|oneshot_recv_result| { oneshot_recv_result .expect("ClientRequest oneshot sender must not be dropped before send") diff --git a/zebra-network/src/peer/error.rs b/zebra-network/src/peer/error.rs index f9f828ec..e4aa54ad 100644 --- a/zebra-network/src/peer/error.rs +++ b/zebra-network/src/peer/error.rs @@ -33,6 +33,9 @@ pub enum PeerError { /// already complete. #[error("Remote peer sent handshake messages after handshake")] DuplicateHandshake, + /// A badly-behaved remote peer sent the wrong nonce in response to a heartbeat `Ping`. + #[error("Remote peer sent the wrong heartbeat nonce")] + HeartbeatNonceMismatch, /// This node's internal services were overloaded, so the connection was dropped /// to shed load. #[error("Internal services over capacity")] diff --git a/zebra-network/src/peer/server.rs b/zebra-network/src/peer/server.rs index 5ec078ff..bcdc39ab 100644 --- a/zebra-network/src/peer/server.rs +++ b/zebra-network/src/peer/server.rs @@ -121,6 +121,7 @@ where Either::Right(((), _peer_fut)) => { trace!("client request timed out"); // Re-matching lets us take ownership of tx + // XXX check here for Ping heartbeat timeout? self.state = match self.state { ServerState::AwaitingResponse(_, tx) => { let e = PeerError::ClientRequestTimeout; @@ -216,6 +217,13 @@ where let _ = tx.send(Ok(Response::Ok)); AwaitingRequest }), + (AwaitingRequest, Ping(nonce)) => self + .peer_tx + .send(Message::Ping(nonce)) + .await + .map_err(|e| e.into()) + .map(|()| AwaitingResponse(Ping(nonce), tx)), + // XXX timeout handling here? } { Ok(new_state) => { self.state = new_state; @@ -250,6 +258,12 @@ where .expect("response oneshot should be unused"); AwaitingRequest } + (AwaitingResponse(Ping(req_nonce), tx), Message::Pong(res_nonce)) => { + if req_nonce != res_nonce { + self.fail_with(PeerError::HeartbeatNonceMismatch); + } + AwaitingRequest + } // By default, messages are not responses. (state, msg) => { ignored_msg = Some(msg); diff --git a/zebra-network/src/protocol/internal.rs b/zebra-network/src/protocol/internal.rs index 64bd9e49..622f01b0 100644 --- a/zebra-network/src/protocol/internal.rs +++ b/zebra-network/src/protocol/internal.rs @@ -6,6 +6,8 @@ use crate::meta_addr::MetaAddr; +use super::types::Nonce; + /// A network request, represented in internal format. #[derive(Clone, Debug, Eq, PartialEq)] pub enum Request { @@ -13,6 +15,8 @@ pub enum Request { GetPeers, /// Advertises peers to the remote server. PushPeers(Vec), + /// Heartbeats triggered on peer connection start. + Ping(Nonce), } /// A response to a network request, represented in internal format.