Fix DHT GetRecord query wedge on untrusted peer responses

Author: viquezclaudio

network-libp2p/src/error.rs

@@ -28,6 +28,15 @@ pub enum NetworkError {
     #[error("DHT query finished without a verified record")]
     DhtNoValidRecord,
 
+    #[error("DHT get query entered an inconsistent state")]
+    DhtGetInconsistentState,
+
+    #[error("DHT get record verification failed")]
+    DhtGetVerificationFailed,
+
+    #[error("DHT get query timed out")]
+    DhtGetTimeout,
+
     #[error("DHT PutRecord error: {0:?}")]
     DhtPutRecord(libp2p::kad::PutRecordError),
 

network-libp2p/src/network.rs

@@ -644,7 +644,10 @@ impl NetworkInterface for Network {
             })
             .await?;
 
-        let data = output_rx.await??;
+        let Ok(data) = timeout(REQUEST_TIMEOUT, output_rx).await else {
+            return Err(NetworkError::DhtGetTimeout);
+        };
+        let data = data??;
         // Now decode the signed record and returned the tagged signable record
         let signed_record: TaggedSigned<V, T> = Deserialize::deserialize_from_vec(&data)?;
         Ok(Some(signed_record.record))

network-libp2p/src/swarm.rs

@@ -539,27 +539,63 @@ fn handle_dht_event(event: kad::Event, event_info: EventInfo) {
     }
 }
 
+#[cfg(feature = "kad")]
+fn abort_dht_get(
+    id: QueryId,
+    step: ProgressStep,
+    error: NetworkError,
+    reason: &'static str,
+    event_info: &mut EventInfo,
+) {
+    if let Some(mut query) = event_info.swarm.behaviour_mut().dht.query_mut(&id) {
+        query.finish();
+    }
+
+    event_info.state.dht_get_results.remove(&id);
+
+    if let Some(output) = event_info.state.dht_gets.remove(&id) {
+        if output.send(Err(error)).is_err() {
+            error!(query_id = ?id, ?step, "could not send aborted get record query result to channel");
+        }
+    } else {
+        warn!(query_id = ?id, ?step, %reason, "GetRecord query abort for unknown query ID");
+    }
+}
+
 #[cfg(feature = "kad")]
 fn handle_dht_get(
     id: QueryId,
     result: Result<GetRecordOk, GetRecordError>,
     _stats: QueryStats,
     step: ProgressStep,
-    event_info: EventInfo,
+    mut event_info: EventInfo,
 ) {
+    if !event_info.state.dht_gets.contains_key(&id)
+        && !event_info.state.dht_get_results.contains_key(&id)
+    {
+        debug!(query_id = ?id, ?step, "Ignoring stale GetRecord query progress");
+        return;
+    }
+
     match result {
         Ok(GetRecordOk::FoundRecord(record)) => {
             // Verify incoming record
             let dht_record = match event_info.dht_verifier.verify(&record.record) {
                 Ok(record) => record,
                 Err(error) => {
                     warn!(?error, "DHT record verification failed");
+                    abort_dht_get(
+                        id,
+                        step,
+                        NetworkError::DhtGetVerificationFailed,
+                        "DHT get record verification failed",
+                        &mut event_info,
+                    );
                     return;
                 }
             };
 
             let count = store_dht_record(&mut event_info.state.dht_get_results, id, dht_record);
-
             // Check if we already have a quorum
             if count == event_info.state.dht_quorum {
                 event_info

network-libp2p/tests/network.rs

@@ -16,7 +16,7 @@ use nimiq_network_interface::{
 use nimiq_network_libp2p::{
     dht,
     discovery::{self, peer_contacts::PeerContact},
-    Config, Network,
+    Config, Network, NetworkError,
 };
 use nimiq_serde::{Deserialize, Serialize};
 use nimiq_test_log::test;
@@ -245,6 +245,17 @@ impl dht::Verifier for Verifier {
     }
 }
 
+struct RejectingVerifier;
+
+impl dht::Verifier for RejectingVerifier {
+    fn verify(
+        &self,
+        _record: &libp2p::kad::Record,
+    ) -> Result<dht::DhtRecord, dht::DhtVerifierError> {
+        Err(dht::DhtVerifierError::InvalidSignature)
+    }
+}
+
 async fn create_network_with_n_peers(
     n_peers: usize,
 ) -> (
@@ -515,6 +526,58 @@ async fn dht_put_and_get() {
     assert_eq!(fetched_record, Some(put_record));
 }
 
+#[test(tokio::test)]
+#[cfg(feature = "kad")]
+async fn dht_get_with_verifier_failure_returns_error() {
+    let addr1 = multiaddr![Memory(rand::random::<u64>())];
+    let addr2 = multiaddr![Memory(rand::random::<u64>())];
+    let keys = Arc::new(RwLock::new(BTreeMap::default()));
+
+    let net1 = Network::new(network_config(addr1.clone()), Verifier::new(&keys)).await;
+    net1.listen_on(vec![addr1.clone()]).await;
+
+    let net2 = Network::new(network_config(addr2.clone()), RejectingVerifier).await;
+    net2.listen_on(vec![addr2.clone()]).await;
+
+    let mut events1 = net1.subscribe_events();
+    let mut events2 = net2.subscribe_events();
+
+    net2.dial_address(addr1).await.unwrap();
+
+    let event1 = helper::get_next_peer_event(&mut events1).await;
+    helper::assert_peer_joined(&event1, &net2.get_local_peer_id());
+
+    let event2 = helper::get_next_peer_event(&mut events2).await;
+    helper::assert_peer_joined(&event2, &net1.get_local_peer_id());
+
+    sleep(Duration::from_secs(10)).await;
+
+    let mut rng = test_rng(false);
+    let keypair = KeyPair::generate(&mut rng);
+    let key: Address = (&keypair.public).into();
+
+    let put_record = ValidatorRecord {
+        peer_id: net1.get_local_peer_id(),
+        validator_address: key.clone(),
+        timestamp: 0x42u64,
+    };
+
+    assert!(keys.write().insert(key.clone(), keypair.public).is_none());
+    net1.dht_put(&key, &put_record, &keypair).await.unwrap();
+
+    let result = timeout(
+        Duration::from_secs(2),
+        net2.dht_get::<_, ValidatorRecord<PeerId>, KeyPair>(&key),
+    )
+    .await;
+
+    assert!(result.is_ok(), "dht_get future hung");
+    assert!(matches!(
+        result.unwrap(),
+        Err(NetworkError::DhtGetVerificationFailed)
+    ));
+}
+
 #[test(tokio::test)]
 async fn ban_peer() {
     let (net1, net2) = create_connected_networks().await;