test

Author: iceseer

include/libp2p/muxer/yamux/yamux_stream.hpp

@@ -133,7 +133,7 @@ namespace libp2p::connection {
     closeCompleted();
 
     /// Underlying connection (secured)
-    std::shared_ptr<connection::SecureConnection> connection_;
+    std::weak_ptr<connection::SecureConnection> connection_;
 
     /// Yamux-specific interface of connection
     YamuxStreamFeedback &feedback_;

include/libp2p/muxer/yamux/yamuxed_connection.hpp

@@ -78,6 +78,8 @@ namespace libp2p::connection {
                            ReadCallbackFunc cb) override;
     void deferWriteCallback(std::error_code ec, WriteCallbackFunc cb) override;
 
+    void markAsRegistered();
+
    private:
     using Streams = std::unordered_map<StreamId, std::shared_ptr<YamuxStream>>;
 
@@ -241,6 +243,8 @@ namespace libp2p::connection {
 
     bool close_after_write_ = false;
 
+    bool registered_in_manager_ = false;
+
    public:
     LIBP2P_METRICS_INSTANCE_COUNT_IF_ENABLED(
         libp2p::connection::YamuxedConnection);

src/muxer/yamux/yamux_stream.cpp

@@ -30,14 +30,14 @@ namespace libp2p::connection {
       uint32_t stream_id,
       size_t maximum_window_size,
       size_t write_queue_limit)
-      : connection_(std::move(connection)),
+      : connection_(connection),
         feedback_(feedback),
         stream_id_(stream_id),
         window_size_(YamuxFrame::kInitialWindowSize),
         peers_window_size_(YamuxFrame::kInitialWindowSize),
         maximum_window_size_(maximum_window_size),
         write_queue_(write_queue_limit) {
-    assert(connection_);
+    assert(connection);
     assert(stream_id_ > 0);
     assert(window_size_ <= maximum_window_size_);
     assert(peers_window_size_ <= maximum_window_size_);
@@ -151,19 +151,31 @@ namespace libp2p::connection {
   }
 
   outcome::result<peer::PeerId> YamuxStream::remotePeerId() const {
-    return connection_->remotePeer();
+    if (auto conn = connection_.lock()) {
+      return conn->remotePeer();
+    }
+    return Error::STREAM_RESET_BY_HOST;
   }
 
   outcome::result<bool> YamuxStream::isInitiator() const {
-    return connection_->isInitiator();
+    if (auto conn = connection_.lock()) {
+      return conn->isInitiator();
+    }
+    return Error::STREAM_RESET_BY_HOST;
   }
 
   outcome::result<multi::Multiaddress> YamuxStream::localMultiaddr() const {
-    return connection_->localMultiaddr();
+    if (auto conn = connection_.lock()) {
+      return conn->localMultiaddr();
+    }
+    return Error::STREAM_RESET_BY_HOST;
   }
 
   outcome::result<multi::Multiaddress> YamuxStream::remoteMultiaddr() const {
-    return connection_->remoteMultiaddr();
+    if (auto conn = connection_.lock()) {
+      return conn->remoteMultiaddr();
+    }
+    return Error::STREAM_RESET_BY_HOST;
   }
 
   void YamuxStream::increaseSendWindow(size_t delta) {

src/muxer/yamux/yamuxed_connection.cpp

@@ -553,7 +553,7 @@ namespace libp2p::connection {
       cb(notify_streams_code);
     }
 
-    if (closed_callback_) {
+    if (closed_callback_ && registered_in_manager_) {
       closed_callback_(remote_peer_, shared_from_this());
     }
 
@@ -779,4 +779,8 @@ namespace libp2p::connection {
         },
         config_.ping_interval);
   }
+
+  void YamuxedConnection::markAsRegistered() {
+    registered_in_manager_ = true;
+  }
 }  // namespace libp2p::connection

src/network/impl/connection_manager_impl.cpp

@@ -148,13 +148,14 @@ namespace libp2p::network {
     }
     auto it = connections_.find(peer_id);
     if (it == connections_.end()) {
-      log()->error("inconsistency in onConnectionClosed, peer not found");
+      log()->warn("onConnectionClosed called for peer {} that was not in connection manager (connection may have been closed before being added)", peer_id.toBase58());
       return;
     }
 
     [[maybe_unused]] auto erased = it->second.erase(conn);
     if (erased == 0) {
-      log()->error("inconsistency in onConnectionClosed, connection not found");
+      log()->warn("onConnectionClosed called for connection to peer {} that was not in connection manager (connection may have been closed before being added)", peer_id.toBase58());
+      return;
     }
 
     if (it->second.empty()) {

src/network/impl/listener_manager_impl.cpp

@@ -7,6 +7,7 @@
 #include <libp2p/network/impl/listener_manager_impl.hpp>
 
 #include <libp2p/log/logger.hpp>
+#include <libp2p/muxer/yamux/yamuxed_connection.hpp>
 
 namespace libp2p::network {
 
@@ -240,6 +241,9 @@ namespace libp2p::network {
         });
 
     // store connection
+    if (auto yamux_conn = std::dynamic_pointer_cast<connection::YamuxedConnection>(conn)) {
+      yamux_conn->markAsRegistered();
+    }
     this->cmgr_->addConnectionToPeer(id, conn);
   }