p2p: Deal with shortid collisions for reconciliation sets

If a transaction to be added to a peer's recon set has a shot id collisions (a previously
added wtxid maps to the same short id), both transaction should be fanout, given
our peer may have added the opposite transaction to our recon set, and these two
transaction won't be reconciled.

src/node/txreconciliation.cpp

@@ -128,13 +128,12 @@ private:
 public:
     explicit Impl(uint32_t recon_version) : m_recon_version(recon_version) {}
 
-    uint64_t PreRegisterPeer(NodeId peer_id) EXCLUSIVE_LOCKS_REQUIRED(!m_txreconciliation_mutex)
+    uint64_t PreRegisterPeer(NodeId peer_id, uint64_t local_salt) EXCLUSIVE_LOCKS_REQUIRED(!m_txreconciliation_mutex)
     {
         AssertLockNotHeld(m_txreconciliation_mutex);
         LOCK(m_txreconciliation_mutex);
 
         LogPrintLevel(BCLog::TXRECONCILIATION, BCLog::Level::Debug, "Pre-register peer=%d\n", peer_id);
-        const uint64_t local_salt{FastRandomContext().rand64()};
 
         // We do this exactly once per peer (which are unique by NodeId, see GetNewNodeId) so it's
         // safe to assume we don't have this record yet.
@@ -142,6 +141,21 @@ public:
         return local_salt;
     }
 
+    bool HasCollision(TxReconciliationState *peer_state, const Wtxid& wtxid, Wtxid& collision, uint32_t &short_id) EXCLUSIVE_LOCKS_REQUIRED(m_txreconciliation_mutex)
+    {
+        AssertLockHeld(m_txreconciliation_mutex);
+
+        short_id = peer_state->ComputeShortID(wtxid);
+        const auto iter = peer_state->m_short_id_mapping.find(short_id);
+
+        if (iter != peer_state->m_short_id_mapping.end()) {
+            collision = iter->second;
+            return true;
+        }
+
+        return false;
+    }
+
     ReconciliationRegisterResult RegisterPeer(NodeId peer_id, bool is_peer_inbound, uint32_t peer_recon_version,
                                               uint64_t remote_salt) EXCLUSIVE_LOCKS_REQUIRED(!m_txreconciliation_mutex)
     {
@@ -196,8 +210,20 @@ public:
         auto peer_state = GetRegisteredPeerState(peer_id);
         if (!peer_state) return AddToSetResult::Failed();
 
-        // TODO: We should compute the short_id here here first and see if there's any collision
-        // if so, return AddToSetResult::Collision(wtxid)
+        // Bypass if the wtxid is already in the set
+        if (peer_state->m_local_set.contains(wtxid)) {
+            LogPrintLevel(BCLog::TXRECONCILIATION, BCLog::Level::Debug, "%s already in reconciliation set for peer=%d. Bypassing.\n",
+                          wtxid.ToString(), peer_id);
+            return AddToSetResult::Succeeded();
+        }
+
+        // Make sure there is no short id collision between the wtxid we are trying to add
+        // and any existing one in the reconciliation set
+        Wtxid collision;
+        uint32_t short_id;
+        if (HasCollision(peer_state, wtxid, collision, short_id)) {
+            return AddToSetResult::Collision(collision);
+        }
 
         // Transactions which don't make it to the set due to the limit are announced via fanout.
         if (peer_state->m_local_set.size() >= MAX_RECONSET_SIZE) {
@@ -205,10 +231,8 @@ public:
             return AddToSetResult::Failed();
         }
 
-        // The caller currently keeps track of the per-peer transaction announcements, so it
-        // should not attempt to add same tx to the set twice. However, if that happens, we will
-        // simply ignore it.
         if (peer_state->m_local_set.insert(wtxid).second) {
+            peer_state->m_short_id_mapping.emplace(short_id, wtxid);
             LogPrintLevel(BCLog::TXRECONCILIATION, BCLog::Level::Debug, "Added %s to the reconciliation set for peer=%d. "
                                                                         "Now the set contains %i transactions.\n",
                           wtxid.ToString(), peer_id, peer_state->m_local_set.size());
@@ -235,6 +259,7 @@ public:
 
         auto removed = peer_state->m_local_set.erase(wtxid) > 0;
         if (removed) {
+            peer_state->m_short_id_mapping.erase(peer_state->ComputeShortID(wtxid));
             LogPrintLevel(BCLog::TXRECONCILIATION, BCLog::Level::Debug, "Removed %s from the reconciliation set for peer=%d. "
                                                                         "Now the set contains %i transactions.\n",
                           wtxid.ToString(), peer_id, peer_state->m_local_set.size());
@@ -353,7 +378,13 @@ TxReconciliationTracker::~TxReconciliationTracker() = default;
 
 uint64_t TxReconciliationTracker::PreRegisterPeer(NodeId peer_id)
 {
-    return m_impl->PreRegisterPeer(peer_id);
+    const uint64_t local_salt{FastRandomContext().rand64()};
+    return m_impl->PreRegisterPeer(peer_id, local_salt);
+}
+
+void TxReconciliationTracker::PreRegisterPeerWithSalt(NodeId peer_id, uint64_t local_salt)
+{
+    m_impl->PreRegisterPeer(peer_id, local_salt);
 }
 
 ReconciliationRegisterResult TxReconciliationTracker::RegisterPeer(NodeId peer_id, bool is_peer_inbound,

src/node/txreconciliation.h

@@ -103,6 +103,11 @@ public:
      */
     uint64_t PreRegisterPeer(NodeId peer_id);
 
+    /**
+     * For testing purposes only. This SHOULD NEVER be used in production.
+    */
+    void PreRegisterPeerWithSalt(NodeId peer_id, uint64_t local_salt);
+
     /**
      * Step 0. Once the peer agreed to reconcile txs with us, generate the state required to track
      * ongoing reconciliations. Must be called only after pre-registering the peer and only once.
@@ -111,9 +116,9 @@ public:
                                               uint32_t peer_recon_version, uint64_t remote_salt);
 
     /**
-     * Step 1. Add a new transaction we want to announce to the peer to the local reconciliation set
-     * of the peer, so that it will be reconciled later, unless the set limit is reached.
-     * Returns whether the transaction appears in the set.
+     * Step 1. Add a to-be-announced transaction to the local reconciliation set of the target peer.
+     * Returns false if the set is at capacity, or if the set contains a colliding transaction (alongside
+     * the colliding wtxid). Returns true if the transaction is added to the set (or if it was already in it).
      */
     AddToSetResult AddToSet(NodeId peer_id, const Wtxid& wtxid);
 

src/test/txreconciliation_tests.cpp

@@ -116,41 +116,53 @@ BOOST_AUTO_TEST_CASE(AddToSetTest)
     BOOST_REQUIRE_EQUAL(tracker.RegisterPeer(peer_id1, true, 1, 1), ReconciliationRegisterResult::SUCCESS);
     BOOST_CHECK(tracker.IsPeerRegistered(peer_id1));
 
-    // As long as the peer is registered and the transaction is not in the set, adding it should succeed
-    for (size_t i = 0; i < MAX_RECONSET_SIZE; ++i)
-        r = tracker.AddToSet(peer_id1, Wtxid::FromUint256(frc.rand256()));
-        BOOST_REQUIRE(r.m_succeeded);
-        BOOST_REQUIRE(!r.m_collision.has_value());
+    // As long as the peer is registered, the transaction is not in the set, and there is no short id
+    // collision, adding should work
+    size_t added_txs = 0;
+    while (added_txs < MAX_RECONSET_SIZE) {
+        wtxid = Wtxid::FromUint256(frc.rand256());
+        Wtxid collision;
 
-    // Trying to add the same item twice should fail
+        r = tracker.AddToSet(peer_id1, wtxid);
+        if (r.m_succeeded) {
+            BOOST_REQUIRE(!r.m_collision.has_value());
+            ++added_txs;
+        } else {
+            BOOST_REQUIRE_EQUAL(r.m_collision.value(), collision);
+        }
+    }
+
+    // Adding one more item will fail due to the set being full
     r = tracker.AddToSet(peer_id1, Wtxid::FromUint256(frc.rand256()));
     BOOST_REQUIRE(!r.m_succeeded);
     BOOST_REQUIRE(!r.m_collision.has_value());
+
+    // Trying to add the same item twice will just bypass
+    r = tracker.AddToSet(peer_id1, wtxid);
+    BOOST_REQUIRE(r.m_succeeded);
+    BOOST_REQUIRE(!r.m_collision.has_value());
 }
 
-BOOST_AUTO_TEST_CASE(TryRemovingFromSetTest)
+BOOST_AUTO_TEST_CASE(AddToSetCollisionTest)
 {
     TxReconciliationTracker tracker(TXRECONCILIATION_VERSION);
     NodeId peer_id0 = 0;
     FastRandomContext frc{/*fDeterministic=*/true};
 
-    Wtxid wtxid{Wtxid::FromUint256(frc.rand256())};
+    // Precompute collision
+    Wtxid wtxid{Wtxid::FromUint256(uint256("c70d778bccef36a81aed8da0b819d2bd28bd8653e56a5d40903df1a0ade0b876"))};
+    Wtxid collision{Wtxid::FromUint256(uint256("ae52a6ecb8733fba1f7af6022a8b9dd327d7825054229fafcad7e03c38ae2a50"))};
 
-    BOOST_REQUIRE(!tracker.IsPeerRegistered(peer_id0));
-    BOOST_REQUIRE(!tracker.TryRemovingFromSet(peer_id0, wtxid));
-
-    tracker.PreRegisterPeer(peer_id0);
+    // Register the peer with a predefined salt so we can force the collision
+    tracker.PreRegisterPeerWithSalt(peer_id0, 2);
     BOOST_REQUIRE_EQUAL(tracker.RegisterPeer(peer_id0, true, 1, 1), ReconciliationRegisterResult::SUCCESS);
     BOOST_CHECK(tracker.IsPeerRegistered(peer_id0));
 
-    BOOST_REQUIRE(!tracker.TryRemovingFromSet(peer_id0, wtxid));
+    // Once the peer is registered, we can try to add both transactions and check
     BOOST_REQUIRE(tracker.AddToSet(peer_id0, wtxid).m_succeeded);
-    BOOST_REQUIRE(tracker.TryRemovingFromSet(peer_id0, wtxid));
-    BOOST_REQUIRE(!tracker.TryRemovingFromSet(peer_id0, wtxid));
-
-    BOOST_REQUIRE(tracker.AddToSet(peer_id0, wtxid).m_succeeded);
-    tracker.ForgetPeer(peer_id0);
-    BOOST_REQUIRE(!tracker.TryRemovingFromSet(peer_id0, wtxid));
+    auto r = tracker.AddToSet(peer_id0, collision);
+    BOOST_REQUIRE(!r.m_succeeded);
+    BOOST_REQUIRE_EQUAL(r.m_collision.value(), wtxid);
 }
 
 BOOST_AUTO_TEST_SUITE_END()