hotshot_libp2p_networking/network/
transport.rs

1use std::{
2    future::Future,
3    io::{Error as IoError, ErrorKind as IoErrorKind},
4    pin::Pin,
5    sync::Arc,
6    task::Poll,
7};
8
9use anyhow::{Context, Result as AnyhowResult, ensure};
10use bimap::BiMap;
11use futures::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, future::poll_fn};
12use hotshot_types::traits::signature_key::SignatureKey;
13use libp2p::{
14    Transport,
15    core::{
16        StreamMuxer,
17        muxing::StreamMuxerExt,
18        transport::{DialOpts, TransportEvent},
19    },
20    identity::PeerId,
21};
22use parking_lot::Mutex;
23use pin_project::pin_project;
24use serde::{Deserialize, Serialize};
25use tokio::time::timeout;
26use tracing::debug;
27
28use crate::network::log_summary::LogEvent;
29
30/// The maximum size of an authentication message. This is used to prevent
31/// DoS attacks by sending large messages.
32const MAX_AUTH_MESSAGE_SIZE: usize = 1024;
33
34/// The timeout for the authentication handshake. This is used to prevent
35/// attacks that keep connections open indefinitely by half-finishing the
36/// handshake.
37const AUTH_HANDSHAKE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
38
39/// A wrapper for a `Transport` that bidirectionally associates (and verifies)
40/// the corresponding consensus keys.
41#[pin_project]
42pub struct ConsensusKeyAuthentication<
43    T: Transport,
44    S: SignatureKey + 'static,
45    C: StreamMuxer + Unpin,
46> {
47    #[pin]
48    /// The underlying transport we are wrapping
49    pub inner: T,
50
51    /// A pre-signed message that we (depending on if it's specified or not) send to the remote peer for authentication
52    pub auth_message: Arc<Option<Vec<u8>>>,
53
54    /// The (verified) map of consensus keys to peer IDs
55    pub consensus_key_to_pid_map: Arc<Mutex<BiMap<S, PeerId>>>,
56
57    /// Phantom data for the connection type
58    pd: std::marker::PhantomData<(C, S)>,
59}
60
61/// A type alias for the future that upgrades a connection to perform the authentication handshake
62type UpgradeFuture<T> =
63    Pin<Box<dyn Future<Output = Result<<T as Transport>::Output, <T as Transport>::Error>> + Send>>;
64
65impl<T: Transport, S: SignatureKey + 'static, C: StreamMuxer + Unpin>
66    ConsensusKeyAuthentication<T, S, C>
67{
68    /// Create a new `ConsensusKeyAuthentication` transport that wraps the given transport
69    /// and authenticates connections against the stake table. If the auth message is `None`,
70    /// the authentication is disabled.
71    pub fn new(
72        inner: T,
73        auth_message: Option<Vec<u8>>,
74        consensus_key_to_pid_map: Arc<Mutex<BiMap<S, PeerId>>>,
75    ) -> Self {
76        Self {
77            inner,
78            auth_message: Arc::from(auth_message),
79            consensus_key_to_pid_map,
80            pd: std::marker::PhantomData,
81        }
82    }
83
84    /// Prove to the remote peer that we are in the stake table by sending
85    /// them our authentication message.
86    ///
87    /// # Errors
88    /// - If we fail to write the message to the stream
89    pub async fn authenticate_with_remote_peer<W: AsyncWrite + Unpin>(
90        stream: &mut W,
91        auth_message: &[u8],
92    ) -> AnyhowResult<()> {
93        // Write the length-delimited message
94        write_length_delimited(stream, auth_message).await?;
95
96        Ok(())
97    }
98
99    /// Verify that the remote peer is:
100    /// - In the stake table
101    /// - Sending us a valid authentication message
102    /// - Sending us a valid signature
103    /// - Matching the peer ID we expect
104    ///
105    /// # Errors
106    /// If the peer fails verification. This can happen if:
107    /// - We fail to read the message from the stream
108    /// - The message is too large
109    /// - The message is invalid
110    /// - The peer is not in the stake table
111    /// - The signature is invalid
112    pub async fn verify_peer_authentication<R: AsyncReadExt + Unpin>(
113        stream: &mut R,
114        required_peer_id: &PeerId,
115        consensus_key_to_pid_map: Arc<Mutex<BiMap<S, PeerId>>>,
116    ) -> AnyhowResult<()> {
117        // Read the length-delimited message from the remote peer
118        let message = read_length_delimited(stream, MAX_AUTH_MESSAGE_SIZE).await?;
119
120        // Deserialize the authentication message
121        let auth_message: AuthMessage<S> =
122            bincode::deserialize(&message).with_context(|| "Failed to deserialize auth message")?;
123
124        // Verify the signature on the public keys
125        let public_key = auth_message
126            .validate()
127            .with_context(|| "Failed to verify authentication message")?;
128
129        // Deserialize the `PeerId`
130        let peer_id = PeerId::from_bytes(&auth_message.peer_id_bytes)
131            .with_context(|| "Failed to deserialize peer ID")?;
132
133        // Verify that the peer ID is the same as the remote peer
134        if peer_id != *required_peer_id {
135            return Err(anyhow::anyhow!("Peer ID mismatch"));
136        }
137
138        // If we got here, the peer is authenticated. Add the consensus key to the map
139        consensus_key_to_pid_map.lock().insert(public_key, peer_id);
140
141        Ok(())
142    }
143
144    /// Wrap the supplied future in an upgrade that performs the authentication handshake.
145    ///
146    /// `outgoing` is a boolean that indicates if the connection is incoming or outgoing.
147    /// This is needed because the flow of the handshake is different for each.
148    fn gen_handshake<F: Future<Output = Result<T::Output, T::Error>> + Send + 'static>(
149        original_future: F,
150        outgoing: bool,
151        auth_message: Arc<Option<Vec<u8>>>,
152        consensus_key_to_pid_map: Arc<Mutex<BiMap<S, PeerId>>>,
153    ) -> UpgradeFuture<T>
154    where
155        T::Error: From<<C as StreamMuxer>::Error> + From<IoError>,
156        T::Output: AsOutput<C> + Send,
157        C::Substream: Unpin + Send,
158    {
159        // Create a new upgrade that performs the authentication handshake on top
160        Box::pin(async move {
161            // Wait for the original future to resolve
162            let mut stream = original_future.await?;
163
164            // Time out the authentication block
165            timeout(AUTH_HANDSHAKE_TIMEOUT, async {
166                // Open a substream for the handshake.
167                // The handshake order depends on whether the connection is incoming or outgoing.
168                let mut substream = if outgoing {
169                    poll_fn(|cx| stream.as_connection().poll_outbound_unpin(cx)).await?
170                } else {
171                    poll_fn(|cx| stream.as_connection().poll_inbound_unpin(cx)).await?
172                };
173
174                // Conditionally authenticate depending on whether we specified an auth message
175                if let Some(auth_message) = auth_message.as_ref() {
176                    if outgoing {
177                        // If the connection is outgoing, authenticate with the remote peer first
178                        Self::authenticate_with_remote_peer(&mut substream, auth_message)
179                            .await
180                            .map_err(|e| {
181                                LogEvent::AuthFailure.record();
182                                debug!("Failed to authenticate with remote peer: {e:?}");
183                                IoError::other(e)
184                            })?;
185
186                        // Verify the remote peer's authentication
187                        Self::verify_peer_authentication(
188                            &mut substream,
189                            stream.as_peer_id(),
190                            consensus_key_to_pid_map,
191                        )
192                        .await
193                        .map_err(|e| {
194                            LogEvent::VerifyFailure.record();
195                            debug!("Failed to verify remote peer: {e:?}");
196                            IoError::other(e)
197                        })?;
198                    } else {
199                        // If it is incoming, verify the remote peer's authentication first
200                        Self::verify_peer_authentication(
201                            &mut substream,
202                            stream.as_peer_id(),
203                            consensus_key_to_pid_map,
204                        )
205                        .await
206                        .map_err(|e| {
207                            LogEvent::VerifyFailure.record();
208                            debug!("Failed to verify remote peer: {e:?}");
209                            IoError::other(e)
210                        })?;
211
212                        // Authenticate with the remote peer
213                        Self::authenticate_with_remote_peer(&mut substream, auth_message)
214                            .await
215                            .map_err(|e| {
216                                LogEvent::AuthFailure.record();
217                                debug!("Failed to authenticate with remote peer: {e:?}");
218                                IoError::other(e)
219                            })?;
220                    }
221                }
222
223                Ok(stream)
224            })
225            .await
226            .map_err(|e| {
227                LogEvent::AuthHandshakeTimeout.record();
228                debug!("Timed out performing authentication handshake: {e:?}");
229                IoError::new(IoErrorKind::TimedOut, e)
230            })?
231        })
232    }
233}
234
235/// The deserialized form of an authentication message that is sent to the remote peer
236#[derive(Clone, Serialize, Deserialize)]
237struct AuthMessage<S: SignatureKey> {
238    /// The encoded (stake table) public key of the sender. This, along with the peer ID, is
239    /// signed. It is still encoded here to enable easy verification.
240    #[serde(with = "serde_bytes")]
241    public_key_bytes: Vec<u8>,
242
243    /// The encoded peer ID of the sender. This is appended to the public key before signing.
244    /// It is still encoded here to enable easy verification.
245    #[serde(with = "serde_bytes")]
246    peer_id_bytes: Vec<u8>,
247
248    /// The signature on the public key
249    signature: S::PureAssembledSignatureType,
250}
251
252impl<S: SignatureKey> AuthMessage<S> {
253    /// Validate the signature on the public key and return it if valid
254    pub fn validate(&self) -> AnyhowResult<S> {
255        // Deserialize the stake table public key
256        let public_key = S::from_bytes(&self.public_key_bytes)
257            .with_context(|| "Failed to deserialize public key")?;
258
259        // Reconstruct the signed message from the public key and peer ID
260        let mut signed_message = public_key.to_bytes();
261        signed_message.extend(self.peer_id_bytes.clone());
262
263        // Check if the signature is valid across both
264        if !public_key.validate(&self.signature, &signed_message) {
265            return Err(anyhow::anyhow!("Invalid signature"));
266        }
267
268        Ok(public_key)
269    }
270}
271
272/// Create an sign an authentication message to be sent to the remote peer
273///
274/// # Errors
275/// - If we fail to sign the public key
276/// - If we fail to serialize the authentication message
277pub fn construct_auth_message<S: SignatureKey + 'static>(
278    public_key: &S,
279    peer_id: &PeerId,
280    private_key: &S::PrivateKey,
281) -> AnyhowResult<Vec<u8>> {
282    // Serialize the stake table public key
283    let mut public_key_bytes = public_key.to_bytes();
284
285    // Serialize the peer ID and append it
286    let peer_id_bytes = peer_id.to_bytes();
287    public_key_bytes.extend_from_slice(&peer_id_bytes);
288
289    // Sign our public key
290    let signature =
291        S::sign(private_key, &public_key_bytes).with_context(|| "Failed to sign public key")?;
292
293    // Create the auth message
294    let auth_message = AuthMessage::<S> {
295        public_key_bytes,
296        peer_id_bytes,
297        signature,
298    };
299
300    // Serialize the auth message
301    bincode::serialize(&auth_message).with_context(|| "Failed to serialize auth message")
302}
303
304impl<T: Transport, S: SignatureKey + 'static, C: StreamMuxer + Unpin> Transport
305    for ConsensusKeyAuthentication<T, S, C>
306where
307    T::Dial: Future<Output = Result<T::Output, T::Error>> + Send + 'static,
308    T::ListenerUpgrade: Send + 'static,
309    T::Output: AsOutput<C> + Send,
310    T::Error: From<<C as StreamMuxer>::Error> + From<IoError>,
311    C::Substream: Unpin + Send,
312{
313    // `Dial` is for connecting out, `ListenerUpgrade` is for accepting incoming connections
314    type Dial = Pin<Box<dyn Future<Output = Result<T::Output, T::Error>> + Send>>;
315    type ListenerUpgrade = Pin<Box<dyn Future<Output = Result<T::Output, T::Error>> + Send>>;
316
317    // These are just passed through
318    type Output = T::Output;
319    type Error = T::Error;
320
321    /// Dial a remote peer. This function is changed to perform an authentication handshake
322    /// on top.
323    fn dial(
324        &mut self,
325        addr: libp2p::Multiaddr,
326        opts: DialOpts,
327    ) -> Result<Self::Dial, libp2p::TransportError<Self::Error>> {
328        // Perform the inner dial
329        let res = self.inner.dial(addr, opts);
330
331        // Clone the necessary fields
332        let auth_message = Arc::clone(&self.auth_message);
333
334        // If the dial was successful, perform the authentication handshake on top
335        match res {
336            Ok(dial) => Ok(Self::gen_handshake(
337                dial,
338                true,
339                auth_message,
340                Arc::clone(&self.consensus_key_to_pid_map),
341            )),
342            Err(err) => Err(err),
343        }
344    }
345
346    /// This function is where we perform the authentication handshake for _incoming_ connections.
347    /// The flow in this case is the reverse of the `dial` function: we first verify the remote peer's
348    /// authentication, and then authenticate with them.
349    fn poll(
350        mut self: std::pin::Pin<&mut Self>,
351        cx: &mut std::task::Context<'_>,
352    ) -> std::task::Poll<libp2p::core::transport::TransportEvent<Self::ListenerUpgrade, Self::Error>>
353    {
354        match Transport::poll(self.as_mut().project().inner, cx) {
355            Poll::Ready(event) => Poll::Ready(match event {
356                // If we have an incoming connection, we need to perform the authentication handshake
357                TransportEvent::Incoming {
358                    listener_id,
359                    upgrade,
360                    local_addr,
361                    send_back_addr,
362                } => {
363                    // Clone the necessary fields
364                    let auth_message = Arc::clone(&self.auth_message);
365
366                    // Generate the handshake upgrade future (inbound)
367                    let auth_upgrade = Self::gen_handshake(
368                        upgrade,
369                        false,
370                        auth_message,
371                        Arc::clone(&self.consensus_key_to_pid_map),
372                    );
373
374                    // Return the new event
375                    TransportEvent::Incoming {
376                        listener_id,
377                        upgrade: auth_upgrade,
378                        local_addr,
379                        send_back_addr,
380                    }
381                },
382
383                // We need to re-map the other events because we changed the type of the upgrade
384                TransportEvent::AddressExpired {
385                    listener_id,
386                    listen_addr,
387                } => TransportEvent::AddressExpired {
388                    listener_id,
389                    listen_addr,
390                },
391                TransportEvent::ListenerClosed {
392                    listener_id,
393                    reason,
394                } => TransportEvent::ListenerClosed {
395                    listener_id,
396                    reason,
397                },
398                TransportEvent::ListenerError { listener_id, error } => {
399                    TransportEvent::ListenerError { listener_id, error }
400                },
401                TransportEvent::NewAddress {
402                    listener_id,
403                    listen_addr,
404                } => TransportEvent::NewAddress {
405                    listener_id,
406                    listen_addr,
407                },
408            }),
409
410            Poll::Pending => Poll::Pending,
411        }
412    }
413
414    /// The below functions just pass through to the inner transport, but we had
415    /// to define them
416    fn remove_listener(&mut self, id: libp2p::core::transport::ListenerId) -> bool {
417        self.inner.remove_listener(id)
418    }
419    fn listen_on(
420        &mut self,
421        id: libp2p::core::transport::ListenerId,
422        addr: libp2p::Multiaddr,
423    ) -> Result<(), libp2p::TransportError<Self::Error>> {
424        self.inner.listen_on(id, addr)
425    }
426}
427
428/// A helper trait that allows us to access the underlying connection
429/// and `PeerId` from a transport output
430trait AsOutput<C: StreamMuxer + Unpin> {
431    /// Get a mutable reference to the underlying connection
432    fn as_connection(&mut self) -> &mut C;
433
434    /// Get a mutable reference to the underlying `PeerId`
435    fn as_peer_id(&mut self) -> &mut PeerId;
436}
437
438/// The implementation of the `AsConnection` trait for a tuple of a `PeerId`
439/// and a connection.
440impl<C: StreamMuxer + Unpin> AsOutput<C> for (PeerId, C) {
441    /// Get a mutable reference to the underlying connection
442    fn as_connection(&mut self) -> &mut C {
443        &mut self.1
444    }
445
446    /// Get a mutable reference to the underlying `PeerId`
447    fn as_peer_id(&mut self) -> &mut PeerId {
448        &mut self.0
449    }
450}
451
452/// A helper function to read a length-delimited message from a stream. Takes into
453/// account the maximum message size.
454///
455/// # Errors
456/// - If the message is too big
457/// - If we fail to read from the stream
458pub async fn read_length_delimited<S: AsyncRead + Unpin>(
459    stream: &mut S,
460    max_size: usize,
461) -> AnyhowResult<Vec<u8>> {
462    // Receive the first 8 bytes of the message, which is the length
463    let mut len_bytes = [0u8; 4];
464    stream
465        .read_exact(&mut len_bytes)
466        .await
467        .with_context(|| "Failed to read message length")?;
468
469    // Parse the length of the message as a `u32`
470    let len = usize::try_from(u32::from_be_bytes(len_bytes))?;
471
472    // Quit if the message is too large
473    ensure!(len <= max_size, "Message too large");
474
475    // Read the actual message
476    let mut message = vec![0u8; len];
477    stream
478        .read_exact(&mut message)
479        .await
480        .with_context(|| "Failed to read message")?;
481
482    Ok(message)
483}
484
485/// A helper function to write a length-delimited message to a stream.
486///
487/// # Errors
488/// - If we fail to write to the stream
489pub async fn write_length_delimited<S: AsyncWrite + Unpin>(
490    stream: &mut S,
491    message: &[u8],
492) -> AnyhowResult<()> {
493    // Write the length of the message
494    stream
495        .write_all(&u32::try_from(message.len())?.to_be_bytes())
496        .await
497        .with_context(|| "Failed to write message length")?;
498
499    // Write the actual message
500    stream
501        .write_all(message)
502        .await
503        .with_context(|| "Failed to write message")?;
504
505    Ok(())
506}
507
508#[cfg(test)]
509mod test {
510    use hotshot_types::{signature_key::BLSPubKey, traits::signature_key::SignatureKey};
511    use libp2p::{core::transport::dummy::DummyTransport, quic::Connection};
512    use rand::Rng;
513
514    use super::*;
515
516    /// A mock type to help with readability
517    type MockStakeTableAuth = ConsensusKeyAuthentication<DummyTransport, BLSPubKey, Connection>;
518
519    // Helper macro for generating a new identity and authentication message
520    macro_rules! new_identity {
521        () => {{
522            // Gen a new seed
523            let seed = rand::rngs::OsRng.r#gen::<[u8; 32]>();
524
525            // Create a new keypair
526            let keypair = BLSPubKey::generated_from_seed_indexed(seed, 1337);
527
528            // Create a peer ID
529            let peer_id = libp2p::identity::Keypair::generate_ed25519()
530                .public()
531                .to_peer_id();
532
533            // Construct an authentication message
534            let auth_message =
535                super::construct_auth_message(&keypair.0, &peer_id, &keypair.1).unwrap();
536
537            (keypair, peer_id, auth_message)
538        }};
539    }
540
541    // Helper macro to generator a cursor from a length-delimited message
542    macro_rules! cursor_from {
543        ($auth_message:expr) => {{
544            let mut stream = futures::io::Cursor::new(vec![]);
545            write_length_delimited(&mut stream, &$auth_message)
546                .await
547                .expect("Failed to write message");
548            stream.set_position(0);
549            stream
550        }};
551    }
552
553    /// Test valid construction and verification of an authentication message
554    #[test]
555    fn signature_verify() {
556        // Create a new identity
557        let (_, _, auth_message) = new_identity!();
558
559        // Verify the authentication message
560        let public_key = super::AuthMessage::<BLSPubKey>::validate(
561            &bincode::deserialize(&auth_message).unwrap(),
562        );
563        assert!(public_key.is_ok());
564    }
565
566    /// Test invalid construction and verification of an authentication message with
567    /// an invalid public key. This ensures we are signing over it correctly.
568    #[test]
569    fn signature_verify_invalid_public_key() {
570        // Create a new identity
571        let (_, _, auth_message) = new_identity!();
572
573        // Deserialize the authentication message
574        let mut auth_message: super::AuthMessage<BLSPubKey> =
575            bincode::deserialize(&auth_message).unwrap();
576
577        // Change the public key
578        auth_message.public_key_bytes[0] ^= 0x01;
579
580        // Serialize the message again
581        let auth_message = bincode::serialize(&auth_message).unwrap();
582
583        // Verify the authentication message
584        let public_key = super::AuthMessage::<BLSPubKey>::validate(
585            &bincode::deserialize(&auth_message).unwrap(),
586        );
587        assert!(public_key.is_err());
588    }
589
590    /// Test invalid construction and verification of an authentication message with
591    /// an invalid peer ID. This ensures we are signing over it correctly.
592    #[test]
593    fn signature_verify_invalid_peer_id() {
594        // Create a new identity
595        let (_, _, auth_message) = new_identity!();
596
597        // Deserialize the authentication message
598        let mut auth_message: super::AuthMessage<BLSPubKey> =
599            bincode::deserialize(&auth_message).unwrap();
600
601        // Change the peer ID
602        auth_message.peer_id_bytes[0] ^= 0x01;
603
604        // Serialize the message again
605        let auth_message = bincode::serialize(&auth_message).unwrap();
606
607        // Verify the authentication message
608        let public_key = super::AuthMessage::<BLSPubKey>::validate(
609            &bincode::deserialize(&auth_message).unwrap(),
610        );
611        assert!(public_key.is_err());
612    }
613
614    #[tokio::test(flavor = "multi_thread")]
615    async fn valid_authentication() {
616        // Create a new identity
617        let (keypair, peer_id, auth_message) = new_identity!();
618
619        // Create a stream and write the message to it
620        let mut stream = cursor_from!(auth_message);
621
622        // Create a map from consensus keys to peer IDs
623        let consensus_key_to_pid_map = Arc::new(parking_lot::Mutex::new(BiMap::new()));
624
625        // Verify the authentication message
626        let result = MockStakeTableAuth::verify_peer_authentication(
627            &mut stream,
628            &peer_id,
629            Arc::clone(&consensus_key_to_pid_map),
630        )
631        .await;
632
633        // Make sure the map has the correct entry
634        assert!(
635            consensus_key_to_pid_map
636                .lock()
637                .get_by_left(&keypair.0)
638                .unwrap()
639                == &peer_id,
640            "Map does not have the correct entry"
641        );
642
643        assert!(
644            result.is_ok(),
645            "Should have passed authentication but did not"
646        );
647    }
648
649    #[tokio::test(flavor = "multi_thread")]
650    async fn peer_id_mismatch() {
651        // Create a new identity and authentication message
652        let (_, _, auth_message) = new_identity!();
653
654        // Create a second (malicious) identity
655        let (_, malicious_peer_id, _) = new_identity!();
656
657        // Create a stream and write the message to it
658        let mut stream = cursor_from!(auth_message);
659
660        // Create a map from consensus keys to peer IDs
661        let consensus_key_to_pid_map = Arc::new(parking_lot::Mutex::new(BiMap::new()));
662
663        // Check against the malicious peer ID
664        let result = MockStakeTableAuth::verify_peer_authentication(
665            &mut stream,
666            &malicious_peer_id,
667            Arc::clone(&consensus_key_to_pid_map),
668        )
669        .await;
670
671        // Make sure it errored for the right reason
672        assert!(
673            result
674                .expect_err("Should have failed authentication but did not")
675                .to_string()
676                .contains("Peer ID mismatch"),
677            "Did not fail with the correct error"
678        );
679
680        // Make sure the map does not have the malicious peer ID
681        assert!(
682            consensus_key_to_pid_map.lock().is_empty(),
683            "Malicious peer ID should not be in the map"
684        );
685    }
686
687    #[tokio::test(flavor = "multi_thread")]
688    async fn read_and_write_length_delimited() {
689        // Create a message
690        let message = b"Hello, world!";
691
692        // Write the message to a buffer
693        let mut buffer = Vec::new();
694        write_length_delimited(&mut buffer, message).await.unwrap();
695
696        // Read the message from the buffer
697        let read_message = read_length_delimited(&mut buffer.as_slice(), 1024)
698            .await
699            .unwrap();
700
701        // Check if the messages are the same
702        assert_eq!(message, read_message.as_slice());
703    }
704}
hotshot_libp2p_networking/network/transport.rs

hotshot_libp2p_networking/network/
transport.rs
