Authenticate Libp2p connections

crates/hotshot/src/traits/networking/libp2p_network.rs

@@ -58,7 +58,9 @@ use libp2p_identity::{
 use libp2p_networking::{
     network::{
         behaviours::request_response::{Request, Response},
-        spawn_network_node, MeshParams,
+        spawn_network_node,
+        transport::construct_auth_message,
+        MeshParams,
         NetworkEvent::{self, DirectRequest, DirectResponse, GossipMsg},
         NetworkNodeConfig, NetworkNodeConfigBuilder, NetworkNodeHandle, NetworkNodeReceiver,
         NetworkNodeType, DEFAULT_REPLICATION_FACTOR,
@@ -144,7 +146,7 @@ struct Libp2pNetworkInner<K: SignatureKey + 'static> {
     /// this node's public key
     pk: K,
     /// handle to control the network
-    handle: Arc<NetworkNodeHandle>,
+    handle: Arc<NetworkNodeHandle<K>>,
     /// Message Receiver
     receiver: UnboundedReceiver<Vec<u8>>,
     /// Receiver for Requests for Data, includes the request and the response chan
@@ -393,7 +395,23 @@ impl<K: SignatureKey + 'static> Libp2pNetwork<K> {
         .parse()?;
 
         // Build our libp2p configuration from our global, network configuration
-        let mut config_builder: NetworkNodeConfigBuilder = NetworkNodeConfigBuilder::default();
+        let mut config_builder = NetworkNodeConfigBuilder::default();
+
+        // Extrapolate the stake table from the known nodes
+        let stake_table: HashSet<K> = config
+            .config
+            .known_nodes_with_stake
+            .iter()
+            .map(|node| K::public_key(&node.stake_table_entry))
+            .collect();
+
+        let auth_message = construct_auth_message(pub_key, priv_key)
+            .with_context(|| "Failed to construct auth message")?;
+
+        // Set the auth message and stake table
+        config_builder
+            .stake_table(Some(stake_table))
+            .auth_message(Some(auth_message));
 
         // The replication factor is the minimum of [the default and 2/3 the number of nodes]
         let Some(default_replication_factor) = DEFAULT_REPLICATION_FACTOR else {
@@ -486,7 +504,7 @@ impl<K: SignatureKey + 'static> Libp2pNetwork<K> {
     #[allow(clippy::too_many_arguments)]
     pub async fn new(
         metrics: Libp2pMetricsValue,
-        config: NetworkNodeConfig,
+        config: NetworkNodeConfig<K>,
         pk: K,
         bootstrap_addrs: BootstrapAddrs,
         id: usize,

crates/libp2p-networking/Cargo.toml

@@ -36,6 +36,7 @@ tide = { version = "0.16", optional = true, default-features = false, features =
 tracing = { workspace = true }
 void = "1"
 lazy_static = { workspace = true }
+pin-project = "1"
 
 [target.'cfg(all(async_executor_impl = "tokio"))'.dependencies]
 libp2p = { workspace = true, features = ["tokio"] }

crates/libp2p-networking/src/network/mod.rs

@@ -6,10 +6,13 @@ mod def;
 pub mod error;
 /// functionality of a libp2p network node
 mod node;
+/// Alternative Libp2p transport implementations
+pub mod transport;
 
 use std::{collections::HashSet, fmt::Debug, str::FromStr};
 
 use futures::channel::oneshot::{self, Sender};
+use hotshot_types::traits::signature_key::SignatureKey;
 #[cfg(async_executor_impl = "async-std")]
 use libp2p::dns::async_std::Transport as DnsTransport;
 #[cfg(async_executor_impl = "tokio")]
@@ -31,6 +34,7 @@ use quic::async_std::Transport as QuicTransport;
 use quic::tokio::Transport as QuicTransport;
 use serde::{Deserialize, Serialize};
 use tracing::instrument;
+use transport::StakeTableAuthentication;
 
 use self::behaviours::request_response::{Request, Response};
 pub use self::{
@@ -205,31 +209,43 @@ pub fn gen_multiaddr(port: u16) -> Multiaddr {
 /// This type is used to represent a transport in the libp2p network framework. The `PeerId` is a unique identifier for each peer in the network, and the `StreamMuxerBox` is a type of multiplexer that can handle multiple substreams over a single connection.
 type BoxedTransport = Boxed<(PeerId, StreamMuxerBox)>;
 
-/// Generate authenticated transport
+/// Generates an authenticated transport checked against the stake table.
+/// If the stake table or authentication message is not provided, the transport will
+/// not participate in stake table authentication.
+///
 /// # Errors
-/// could not sign the quic key with `identity`
+/// If we could not create a DNS transport
 #[instrument(skip(identity))]
-pub async fn gen_transport(identity: Keypair) -> Result<BoxedTransport, NetworkError> {
-    let quic_transport = {
+pub async fn gen_transport<K: SignatureKey + 'static>(
+    identity: Keypair,
+    stake_table: Option<HashSet<K>>,
+    auth_message: Option<Vec<u8>>,
+) -> Result<BoxedTransport, NetworkError> {
+    // Create the initial `Quic` transport
+    let transport = {
         let mut config = quic::Config::new(&identity);
         config.handshake_timeout = std::time::Duration::from_secs(20);
         QuicTransport::new(config)
     };
 
-    let dns_quic = {
+    // Require authentication against the stake table
+    let transport = StakeTableAuthentication::new(transport, stake_table, auth_message);
+
+    // Support DNS resolution
+    let transport = {
         #[cfg(async_executor_impl = "async-std")]
         {
-            DnsTransport::system(quic_transport).await
+            DnsTransport::system(transport).await
         }
 
         #[cfg(async_executor_impl = "tokio")]
         {
-            DnsTransport::system(quic_transport)
+            DnsTransport::system(transport)
         }
     }
     .map_err(|e| NetworkError::TransportLaunch { source: e })?;
 
-    Ok(dns_quic
+    Ok(transport
         .map(|(peer_id, connection), _| (peer_id, StreamMuxerBox::new(connection)))
         .boxed())
 }

crates/libp2p-networking/src/network/node.rs

@@ -17,7 +17,9 @@ use async_compatibility_layer::{
     channel::{unbounded, UnboundedReceiver, UnboundedRecvError, UnboundedSender},
 };
 use futures::{channel::mpsc, select, FutureExt, SinkExt, StreamExt};
-use hotshot_types::constants::KAD_DEFAULT_REPUB_INTERVAL_SEC;
+use hotshot_types::{
+    constants::KAD_DEFAULT_REPUB_INTERVAL_SEC, traits::signature_key::SignatureKey,
+};
 use libp2p::{
     autonat,
     core::transport::ListenerId,
@@ -76,7 +78,7 @@ pub const ESTABLISHED_LIMIT_UNWR: u32 = 10;
 
 /// Network definition
 #[derive(custom_debug::Debug)]
-pub struct NetworkNode {
+pub struct NetworkNode<K: SignatureKey + 'static> {
     /// pub/private key from with peer_id is derived
     identity: Keypair,
     /// peer id of network node
@@ -85,7 +87,7 @@ pub struct NetworkNode {
     #[debug(skip)]
     swarm: Swarm<NetworkDef>,
     /// the configuration parameters of the netework
-    config: NetworkNodeConfig,
+    config: NetworkNodeConfig<K>,
     /// the listener id we are listening on, if it exists
     listener_id: Option<ListenerId>,
     /// Handler for requests and response behavior events.
@@ -100,7 +102,7 @@ pub struct NetworkNode {
     bootstrap_tx: Option<mpsc::Sender<bootstrap::InputEvent>>,
 }
 
-impl NetworkNode {
+impl<K: SignatureKey + 'static> NetworkNode<K> {
     /// Returns number of peers this node is connected to
     pub fn num_connected(&self) -> usize {
         self.swarm.connected_peers().count()
@@ -158,17 +160,25 @@ impl NetworkNode {
     ///   * Generates a connection to the "broadcast" topic
     ///   * Creates a swarm to manage peers and events
     #[instrument]
-    pub async fn new(config: NetworkNodeConfig) -> Result<Self, NetworkError> {
-        // Generate a random PeerId
+    pub async fn new(config: NetworkNodeConfig<K>) -> Result<Self, NetworkError> {
+        // Generate a random `KeyPair` if one is not specified
         let identity = if let Some(ref kp) = config.identity {
             kp.clone()
         } else {
             Keypair::generate_ed25519()
         };
+
+        // Get the `PeerId` from the `KeyPair`
         let peer_id = PeerId::from(identity.public());
-        debug!(?peer_id);
-        let transport: BoxedTransport = gen_transport(identity.clone()).await?;
-        debug!("Launched network transport");
+
+        // Generate the transport from the identity, stake table, and auth message
+        let transport: BoxedTransport = gen_transport::<K>(
+            identity.clone(),
+            config.stake_table.clone(),
+            config.auth_message.clone(),
+        )
+        .await?;
+
         // Generate the swarm
         let mut swarm: Swarm<NetworkDef> = {
             // Use the hash of the message's contents as the ID

crates/libp2p-networking/src/network/node/config.rs

@@ -1,5 +1,6 @@
 use std::{collections::HashSet, num::NonZeroUsize, time::Duration};
 
+use hotshot_types::traits::signature_key::SignatureKey;
 use libp2p::{identity::Keypair, Multiaddr};
 use libp2p_identity::PeerId;
 
@@ -10,7 +11,7 @@ pub const DEFAULT_REPLICATION_FACTOR: Option<NonZeroUsize> = NonZeroUsize::new(1
 
 /// describe the configuration of the network
 #[derive(Clone, Default, derive_builder::Builder, custom_debug::Debug)]
-pub struct NetworkNodeConfig {
+pub struct NetworkNodeConfig<K: SignatureKey + 'static> {
     #[builder(default)]
     /// The type of node (bootstrap etc)
     pub node_type: NetworkNodeType,
@@ -40,6 +41,16 @@ pub struct NetworkNodeConfig {
     /// whether to start in libp2p::kad::Mode::Server mode
     #[builder(default = "false")]
     pub server_mode: bool,
+
+    /// The stake table. Used for authenticating other nodes. If not supplied
+    /// we will not check other nodes against the stake table
+    #[builder(default)]
+    pub stake_table: Option<HashSet<K>>,
+
+    /// The signed authentication message sent to the remote peer
+    /// If not supplied we will not send an authentication message during the handshake
+    #[builder(default)]
+    pub auth_message: Option<Vec<u8>>,
 }
 
 /// NOTE: `mesh_outbound_min <= mesh_n_low <= mesh_n <= mesh_n_high`

crates/libp2p-networking/src/network/node/handle.rs

@@ -5,7 +5,9 @@ use async_compatibility_layer::{
     channel::{Receiver, SendError, UnboundedReceiver, UnboundedRecvError, UnboundedSender},
 };
 use futures::channel::oneshot;
-use hotshot_types::traits::network::NetworkError as HotshotNetworkError;
+use hotshot_types::traits::{
+    network::NetworkError as HotshotNetworkError, signature_key::SignatureKey,
+};
 use libp2p::{request_response::ResponseChannel, Multiaddr};
 use libp2p_identity::PeerId;
 use snafu::{ResultExt, Snafu};
@@ -22,9 +24,9 @@ use crate::network::{
 /// - A reference to the state
 /// - Controls for the swarm
 #[derive(Debug, Clone)]
-pub struct NetworkNodeHandle {
+pub struct NetworkNodeHandle<K: SignatureKey + 'static> {
     /// network configuration
-    network_config: NetworkNodeConfig,
+    network_config: NetworkNodeConfig<K>,
 
     /// send an action to the networkbehaviour
     send_network: UnboundedSender<ClientRequest>,
@@ -70,10 +72,10 @@ impl NetworkNodeReceiver {
 /// Spawn a network node task task and return the handle and the receiver for it
 /// # Errors
 /// Errors if spawning the task fails
-pub async fn spawn_network_node(
-    config: NetworkNodeConfig,
+pub async fn spawn_network_node<K: SignatureKey + 'static>(
+    config: NetworkNodeConfig<K>,
     id: usize,
-) -> Result<(NetworkNodeReceiver, NetworkNodeHandle), NetworkNodeHandleError> {
+) -> Result<(NetworkNodeReceiver, NetworkNodeHandle<K>), NetworkNodeHandleError> {
     let mut network = NetworkNode::new(config.clone())
         .await
         .context(NetworkSnafu)?;
@@ -107,7 +109,7 @@ pub async fn spawn_network_node(
     Ok((receiver, handle))
 }
 
-impl NetworkNodeHandle {
+impl<K: SignatureKey + 'static> NetworkNodeHandle<K> {
     /// Cleanly shuts down a swarm node
     /// This is done by sending a message to
     /// the swarm itself to spin down
@@ -487,7 +489,7 @@ impl NetworkNodeHandle {
 
     /// Return a reference to the network config
     #[must_use]
-    pub fn config(&self) -> &NetworkNodeConfig {
+    pub fn config(&self) -> &NetworkNodeConfig<K> {
         &self.network_config
     }
 }