[ntcore] WebSocketConnection: Use weak capture (#5822)

Fixes a potential use-after-free on connection close.
This commit is contained in:
Peter Johnson
2023-10-25 10:22:43 -07:00
committed by GitHub
parent 8a8e220792
commit c7d6ad5a0b
4 changed files with 17 additions and 9 deletions

View File

@@ -401,6 +401,7 @@ void NetworkClient::WsConnected(wpi::WebSocket& ws, uv::Tcp& tcp,
m_wire =
std::make_shared<net::WebSocketConnection>(ws, connInfo.protocol_version);
m_wire->Start();
m_clientImpl = std::make_unique<net::ClientImpl>(
m_loop.Now().count(), m_inst, *m_wire, m_logger, m_timeSyncUpdated,
[this](uint32_t repeatMs) {

View File

@@ -280,6 +280,7 @@ void NetworkServer::ServerConnection4::ProcessWsUpgrade() {
INFO("CONNECTED NT4 client '{}' (from {})", dedupName, m_connInfo);
m_info.remote_id = dedupName;
m_server.AddConnection(this, m_info);
m_wire->Start();
m_websocket->closed.connect([this](uint16_t, std::string_view reason) {
auto realReason = m_wire->GetDisconnectReason();
INFO("DISCONNECTED NT4 client '{}' (from {}): {}", m_info.remote_id,

View File

@@ -84,15 +84,7 @@ void WebSocketConnection::Stream::write_impl(const char* data, size_t len) {
WebSocketConnection::WebSocketConnection(wpi::WebSocket& ws,
unsigned int version)
: m_ws{ws}, m_version{version} {
m_ws.pong.connect([this](auto data) {
if (data.size() != 8) {
return;
}
m_lastPingResponse =
wpi::support::endian::read64<wpi::support::native>(data.data());
});
}
: m_ws{ws}, m_version{version} {}
WebSocketConnection::~WebSocketConnection() {
for (auto&& buf : m_bufs) {
@@ -103,6 +95,18 @@ WebSocketConnection::~WebSocketConnection() {
}
}
void WebSocketConnection::Start() {
m_ws.pong.connect([selfweak = weak_from_this()](auto data) {
if (data.size() != 8) {
return;
}
if (auto self = selfweak.lock()) {
self->m_lastPingResponse =
wpi::support::endian::read64<wpi::support::native>(data.data());
}
});
}
void WebSocketConnection::SendPing(uint64_t time) {
auto buf = AllocBuf();
buf.len = 8;

View File

@@ -26,6 +26,8 @@ class WebSocketConnection final
WebSocketConnection(const WebSocketConnection&) = delete;
WebSocketConnection& operator=(const WebSocketConnection&) = delete;
void Start();
unsigned int GetVersion() const final { return m_version; }
void SendPing(uint64_t time) final;