From 1bcc944e3326fc291b5f698514d974b6571bd461 Mon Sep 17 00:00:00 2001
From: JJL772 <jeremy.lorelli.1337@gmail.com>
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 | 88 ++++++++++++++++++++++++++
 2 files changed, 100 insertions(+)

diff --git a/src/remote/pv/beaconHandler.h b/src/remote/pv/beaconHandler.h
index cc77b76e..50733627 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 @@ class BeaconHandler
      * First beacon flag.
      */
     bool _first;
+    /**
+     * Callback for cleaning up the beacon
+     */
+    std::shared_ptr<BeaconCleanupHandler> _callback;
 
     /**
      * Update beacon.
@@ -100,6 +110,8 @@ class BeaconHandler
                       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 eff79d4a..c22dbddf 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 = 20000;
+
 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();
+
+    /**
+     * Set the last time used to the current time
+     */
+    void touch() { epicsTimeGetCurrent(&m_lastTime); }
+
+private:
+    void remove();
+
+    std::shared_ptr<Callback> m_callback;
+    osiSockAddr m_from;
+    InternalClientContextImpl& m_impl;
+    epicsTimeStamp m_lastTime;
+};
 
 class InternalClientContextImpl :
     public ClientContextImpl,
@@ -4351,12 +4390,25 @@ class InternalClientContextImpl :
         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 = std::make_shared<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 +4608,44 @@ class InternalClientContextImpl :
     Configuration::shared_pointer m_configuration;
 
     TransportRegistry::transportVector_t m_flushTransports;
+
+    friend class BeaconCleanupHandler;
 };
 
+
+BeaconCleanupHandler::BeaconCleanupHandler(InternalClientContextImpl& impl, osiSockAddr addr) :
+    m_callback(std::make_shared<Callback>(*this)),
+    m_from(addr),
+    m_impl(impl)
+{
+    m_impl.m_timer->schedulePeriodic(m_callback, maxBeaconLifetime, maxBeaconLifetime);
+}
+
+BeaconCleanupHandler::~BeaconCleanupHandler()
+{
+    m_impl.m_timer->cancel(m_callback);
+}
+
+void BeaconCleanupHandler::Callback::callback()
+{
+    epicsTimeStamp ts;
+    epicsTimeGetCurrent(&ts);
+    if (epicsTimeDiffInSeconds(&ts, &m_handler.m_lastTime) > maxBeaconLifetime)
+        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); /* Erase the beacon entirely */
+}
+
 size_t InternalClientContextImpl::num_instances;
 size_t InternalClientContextImpl::InternalChannelImpl::num_instances;
 size_t InternalClientContextImpl::InternalChannelImpl::num_active;