From f3ba58b8ff40bddc6ace36a3f0631e380fc9e73a Mon Sep 17 00:00:00 2001
From: Peter Johnson <johnson.peter@gmail.com>
Date: Wed, 8 Mar 2023 21:44:30 -0800
Subject: [PATCH] [ntcore] NT4 client: close timed-out connections

Pings are sent every 3 seconds.  If we don't see a response to the last
ping before it's time to send the next one, close the connection.
---
 ntcore/src/main/native/cpp/NetworkClient.cpp  |  2 +-
 ntcore/src/main/native/cpp/net/ClientImpl.cpp | 19 +++++++++++++++----
 ntcore/src/main/native/cpp/net/ClientImpl.h   |  2 +-
 3 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/ntcore/src/main/native/cpp/NetworkClient.cpp b/ntcore/src/main/native/cpp/NetworkClient.cpp
index 9f6a4c7138f..4391d733eda 100644
--- a/ntcore/src/main/native/cpp/NetworkClient.cpp
+++ b/ntcore/src/main/native/cpp/NetworkClient.cpp
@@ -495,7 +495,7 @@ void NCImpl4::WsConnected(wpi::WebSocket& ws, uv::Tcp& tcp) {
   });
   ws.binary.connect([this](std::span<const uint8_t> data, bool) {
     if (m_clientImpl) {
-      m_clientImpl->ProcessIncomingBinary(data);
+      m_clientImpl->ProcessIncomingBinary(m_loop.Now().count(), data);
     }
   });
 }
diff --git a/ntcore/src/main/native/cpp/net/ClientImpl.cpp b/ntcore/src/main/native/cpp/net/ClientImpl.cpp
index ed9e2d04caa..17ec807edb1 100644
--- a/ntcore/src/main/native/cpp/net/ClientImpl.cpp
+++ b/ntcore/src/main/native/cpp/net/ClientImpl.cpp
@@ -53,7 +53,7 @@ class CImpl : public ServerMessageHandler {
             timeSyncUpdated,
         std::function<void(uint32_t repeatMs)> setPeriodic);
 
-  void ProcessIncomingBinary(std::span<const uint8_t> data);
+  void ProcessIncomingBinary(uint64_t curTimeMs, std::span<const uint8_t> data);
   void HandleLocal(std::vector<ClientMessage>&& msgs);
   bool SendControl(uint64_t curTimeMs);
   void SendValues(uint64_t curTimeMs, bool flush);
@@ -91,6 +91,7 @@ class CImpl : public ServerMessageHandler {
   // timestamp handling
   static constexpr uint32_t kPingIntervalMs = 3000;
   uint64_t m_nextPingTimeMs{0};
+  uint64_t m_pongTimeMs{0};
   uint32_t m_rtt2Us{UINT32_MAX};
   bool m_haveTimeOffset{false};
   int64_t m_serverTimeOffsetUs{0};
@@ -125,7 +126,8 @@ CImpl::CImpl(
   m_setPeriodic(m_periodMs);
 }
 
-void CImpl::ProcessIncomingBinary(std::span<const uint8_t> data) {
+void CImpl::ProcessIncomingBinary(uint64_t curTimeMs,
+                                  std::span<const uint8_t> data) {
   for (;;) {
     if (data.empty()) {
       break;
@@ -150,6 +152,7 @@ void CImpl::ProcessIncomingBinary(std::span<const uint8_t> data) {
       }
       DEBUG4("RTT ping response time {} value {}", value.time(),
              value.GetInteger());
+      m_pongTimeMs = curTimeMs;
       int64_t now = wpi::Now();
       int64_t rtt2 = (now - value.GetInteger()) / 2;
       if (rtt2 < m_rtt2Us) {
@@ -207,6 +210,12 @@ bool CImpl::SendControl(uint64_t curTimeMs) {
 
   // start a timestamp RTT ping if it's time to do one
   if (curTimeMs >= m_nextPingTimeMs) {
+    // if we didn't receive a response to our last ping, disconnect
+    if (m_nextPingTimeMs != 0 && m_pongTimeMs == 0) {
+      m_wire.Disconnect("timed out");
+      return false;
+    }
+
     if (!CheckNetworkReady(curTimeMs)) {
       return false;
     }
@@ -215,6 +224,7 @@ bool CImpl::SendControl(uint64_t curTimeMs) {
     WireEncodeBinary(m_wire.SendBinary().Add(), -1, 0, Value::MakeInteger(now));
     // drift isn't critical here, so just go from current time
     m_nextPingTimeMs = curTimeMs + kPingIntervalMs;
+    m_pongTimeMs = 0;
   }
 
   if (!m_outgoing.empty()) {
@@ -465,8 +475,9 @@ void ClientImpl::ProcessIncomingText(std::string_view data) {
   WireDecodeText(data, *m_impl, m_impl->m_logger);
 }
 
-void ClientImpl::ProcessIncomingBinary(std::span<const uint8_t> data) {
-  m_impl->ProcessIncomingBinary(data);
+void ClientImpl::ProcessIncomingBinary(uint64_t curTimeMs,
+                                       std::span<const uint8_t> data) {
+  m_impl->ProcessIncomingBinary(curTimeMs, data);
 }
 
 void ClientImpl::HandleLocal(std::vector<ClientMessage>&& msgs) {
diff --git a/ntcore/src/main/native/cpp/net/ClientImpl.h b/ntcore/src/main/native/cpp/net/ClientImpl.h
index 98ed32870fd..d6eb1d939ce 100644
--- a/ntcore/src/main/native/cpp/net/ClientImpl.h
+++ b/ntcore/src/main/native/cpp/net/ClientImpl.h
@@ -40,7 +40,7 @@ class ClientImpl {
   ~ClientImpl();
 
   void ProcessIncomingText(std::string_view data);
-  void ProcessIncomingBinary(std::span<const uint8_t> data);
+  void ProcessIncomingBinary(uint64_t curTimeMs, std::span<const uint8_t> data);
   void HandleLocal(std::vector<ClientMessage>&& msgs);
 
   void SendControl(uint64_t curTimeMs);