From ed70c162244b7b42119e2011bb5cf90061160629 Mon Sep 17 00:00:00 2001
From: Lee *!* Clagett <code@leeclagett.com>
Date: Tue, 8 Oct 2024 20:27:09 -0400
Subject: Some cleanup in span/connection_context + few more checks

---
 src/cryptonote_basic/connection_context.cpp | 20 ++++++++++++++++++++
 src/cryptonote_basic/connection_context.h   | 10 +++++++++-
 2 files changed, 29 insertions(+), 1 deletion(-)

(limited to 'src/cryptonote_basic')

diff --git a/src/cryptonote_basic/connection_context.cpp b/src/cryptonote_basic/connection_context.cpp
index 4395bad9f..642749f3c 100644
--- a/src/cryptonote_basic/connection_context.cpp
+++ b/src/cryptonote_basic/connection_context.cpp
@@ -29,6 +29,7 @@
 
 #include "connection_context.h"
 
+#include <boost/optional/optional.hpp>
 #include "cryptonote_protocol/cryptonote_protocol_defs.h"
 #include "p2p/p2p_protocol_defs.h"
 
@@ -69,4 +70,23 @@ namespace cryptonote
     };
     return std::numeric_limits<size_t>::max();
   }
+
+  void cryptonote_connection_context::set_state_normal()
+  {
+    m_state = state_normal;
+    m_expected_heights_start = 0;
+    m_needed_objects.clear();
+    m_needed_objects.shrink_to_fit();
+    m_expected_heights.clear();
+    m_expected_heights.shrink_to_fit();
+    m_requested_objects.clear();
+  }
+
+  boost::optional<crypto::hash> cryptonote_connection_context::get_expected_hash(const uint64_t height) const
+  {
+    const auto difference = height - m_expected_heights_start;
+    if (height < m_expected_heights_start || m_expected_heights.size() < difference)
+      return boost::none;
+    return m_expected_heights[difference];
+  }
 } // cryptonote
diff --git a/src/cryptonote_basic/connection_context.h b/src/cryptonote_basic/connection_context.h
index 818999a60..db24e28fa 100644
--- a/src/cryptonote_basic/connection_context.h
+++ b/src/cryptonote_basic/connection_context.h
@@ -34,6 +34,7 @@
 #include <atomic>
 #include <algorithm>
 #include <boost/date_time/posix_time/posix_time.hpp>
+#include <boost/optional/optional_fwd.hpp>
 #include "net/net_utils_base.h"
 #include "crypto/hash.h"
 
@@ -42,7 +43,7 @@ namespace cryptonote
   struct cryptonote_connection_context: public epee::net_utils::connection_context_base
   {
     cryptonote_connection_context(): m_state(state_before_handshake), m_remote_blockchain_height(0), m_last_response_height(0),
-        m_last_request_time(boost::date_time::not_a_date_time), m_callback_request_count(0),
+        m_expected_heights_start(0), m_last_request_time(boost::date_time::not_a_date_time), m_callback_request_count(0),
         m_last_known_hash(crypto::null_hash), m_pruning_seed(0), m_rpc_port(0), m_rpc_credits_per_hash(0), m_anchor(false), m_score(0),
         m_expect_response(0), m_expect_height(0), m_num_requested(0) {}
 
@@ -92,11 +93,18 @@ namespace cryptonote
     //! \return Maximum number of bytes permissible for `command`.
     static size_t get_max_bytes(int command) noexcept;
 
+    //! Use this instead of `m_state = state_normal`.
+    void set_state_normal();
+
+    boost::optional<crypto::hash> get_expected_hash(uint64_t height) const;
+
     state m_state;
     std::vector<std::pair<crypto::hash, uint64_t>> m_needed_objects;
+    std::vector<crypto::hash> m_expected_heights;
     std::unordered_set<crypto::hash> m_requested_objects;
     uint64_t m_remote_blockchain_height;
     uint64_t m_last_response_height;
+    uint64_t m_expected_heights_start;
     boost::posix_time::ptime m_last_request_time;
     copyable_atomic m_callback_request_count; //in debug purpose: problem with double callback rise
     crypto::hash m_last_known_hash;
-- 
cgit v1.2.3