From 205dc587b0ef4ff8814e52a7a97526ba93a3251d Mon Sep 17 00:00:00 2001
From: Jeremy Lorelli <lorelli@slac.stanford.edu>
Date: Mon, 10 Jun 2024 15:36:47 -0400
Subject: [PATCH] clientContextImpl: Cap the number and age of beacons

Each beacon has an associated mutex. If we allocate
too many beacons on resource constrained systems, i.e.
RTEMS, we may run out of resources and crash.
---
 src/remote/pv/beaconHandler.h          | 12 ++++
 src/remoteClient/clientContextImpl.cpp | 99 ++++++++++++++++++++++++--
 2 files changed, 106 insertions(+), 5 deletions(-)

diff --git a/src/remote/pv/beaconHandler.h b/src/remote/pv/beaconHandler.h
index cc77b76..5073362 100644
--- a/src/remote/pv/beaconHandler.h
+++ b/src/remote/pv/beaconHandler.h
@@ -26,6 +26,12 @@
 #include <pv/remote.h>
 #include <pv/pvAccess.h>
 
+namespace 
+{
+    class InternalClientContextImpl;
+    class BeaconCleanupHandler;
+}
+
 namespace epics {
 namespace pvAccess {
 
@@ -85,6 +91,10 @@ private:
      * First beacon flag.
      */
     bool _first;
+    /**
+     * Callback for cleaning up the beacon
+     */
+    std::shared_ptr<BeaconCleanupHandler> _callback;
 
     /**
      * Update beacon.
@@ -100,6 +110,8 @@ private:
                       ServerGUID const &guid,
                       epics::pvData::int16 sequentalID,
                       epics::pvData::int16 changeCount);
+
+    friend class ::InternalClientContextImpl;
 };
 
 }
diff --git a/src/remoteClient/clientContextImpl.cpp b/src/remoteClient/clientContextImpl.cpp
index 290021a..ee208fd 100644
--- a/src/remoteClient/clientContextImpl.cpp
+++ b/src/remoteClient/clientContextImpl.cpp
@@ -49,6 +49,9 @@ using std::tr1::static_pointer_cast;
 using namespace std;
 using namespace epics::pvData;
 
+static const float maxBeaconLifetime = 180.f * 2.f;
+static const int maxTrackedBeacons = 2048;
+
 namespace epics {
 namespace pvAccess {
 
@@ -3038,7 +3041,43 @@ enum ContextState {
 };
 
 
+/**
+ * Handles cleanup of old beacons.
+ */
+class BeaconCleanupHandler
+{
+public:
+    POINTER_DEFINITIONS(BeaconCleanupHandler);
 
+    class Callback : public TimerCallback
+    {
+    public:
+        Callback(BeaconCleanupHandler& handler) : m_handler(handler)
+        {
+        }
+
+        virtual void callback() OVERRIDE FINAL;
+        virtual void timerStopped() OVERRIDE FINAL;
+
+        BeaconCleanupHandler& m_handler;
+    };
+
+    BeaconCleanupHandler(InternalClientContextImpl& impl, osiSockAddr addr);
+    ~BeaconCleanupHandler();
+
+    /**
+     * Extend the lifetime of the beacon, resetting removal countdown to 0
+     */
+    void touch() { epicsAtomicSetIntT(&m_count, 0); }
+
+private:
+    void remove();
+
+    std::shared_ptr<BeaconCleanupHandler::Callback> m_callback;
+    osiSockAddr m_from;
+    InternalClientContextImpl& m_impl;
+    int m_count;
+};
 
 class InternalClientContextImpl :
     public ClientContextImpl,
@@ -4058,6 +4097,12 @@ public:
 
         m_timer->close();
 
+        // Remove all beacons
+        {
+            Lock guard(m_beaconMapMutex);
+            m_beaconHandlers.clear();
+        }
+
         m_channelSearchManager->cancel();
 
         // this will also close all PVA transports
@@ -4079,11 +4124,6 @@ public:
         while ((transportCount = m_transportRegistry.size()) && tries--)
             epicsThreadSleep(0.025);
 
-        {
-            Lock guard(m_beaconMapMutex);
-            m_beaconHandlers.clear();
-        }
-
         if (transportCount)
             LOG(logLevelDebug, "PVA client context destroyed with %u transport(s) active.", (unsigned)transportCount);
     }
@@ -4351,12 +4391,25 @@ private:
         BeaconHandler::shared_pointer handler;
         if (it == m_beaconHandlers.end())
         {
+            /* If we're tracking too many beacons, we'll just ignore this one */
+            if (m_beaconHandlers.size() >= maxTrackedBeacons)
+            {
+                char ipa[64];
+                sockAddrToDottedIP(&responseFrom->sa, ipa, sizeof(ipa));
+                LOG(logLevelDebug, "Tracked beacon limit reached (%d), ignoring %s\n", maxTrackedBeacons, ipa);
+                return nullptr;
+            }
+
             // stores weak_ptr
             handler.reset(new BeaconHandler(internal_from_this(), responseFrom));
+            handler->_callback.reset(new BeaconCleanupHandler(*this, *responseFrom));
             m_beaconHandlers[*responseFrom] = handler;
         }
         else
+        {
             handler = it->second;
+            handler->_callback->touch(); /* Update the callback's latest use time */
+        }
         return handler;
     }
 
@@ -4556,8 +4609,44 @@ private:
     Configuration::shared_pointer m_configuration;
 
     TransportRegistry::transportVector_t m_flushTransports;
+
+    friend class BeaconCleanupHandler;
 };
 
+
+BeaconCleanupHandler::BeaconCleanupHandler(InternalClientContextImpl& impl, osiSockAddr addr) :
+    m_from(addr),
+    m_impl(impl),
+    m_count(0)
+{
+    m_callback.reset(new Callback(*this));
+    m_impl.m_timer->schedulePeriodic(m_callback, maxBeaconLifetime / 4, maxBeaconLifetime / 4);
+}
+
+BeaconCleanupHandler::~BeaconCleanupHandler()
+{
+    m_impl.m_timer->cancel(m_callback);
+}
+
+void BeaconCleanupHandler::Callback::callback()
+{
+    if (epicsAtomicIncrIntT(&m_handler.m_count) >= 5) {
+        m_handler.remove();
+    }
+}
+
+void BeaconCleanupHandler::Callback::timerStopped()
+{
+    m_handler.remove();
+}
+
+void BeaconCleanupHandler::remove()
+{
+    Lock guard(m_impl.m_beaconMapMutex);
+    m_impl.m_timer->cancel(m_callback);
+    m_impl.m_beaconHandlers.erase(m_from);
+}
+
 size_t InternalClientContextImpl::num_instances;
 size_t InternalClientContextImpl::InternalChannelImpl::num_instances;
 size_t InternalClientContextImpl::InternalChannelImpl::num_active;