From d44adc9899800381a628e01df7afd61abcacccfa Mon Sep 17 00:00:00 2001
From: Jeremy Lorelli <lorelli@slac.stanford.edu>
Date: Tue, 19 Dec 2023 14:52:38 -0800
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/remoteClient/clientContextImpl.cpp | 44 ++++++++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/src/remoteClient/clientContextImpl.cpp b/src/remoteClient/clientContextImpl.cpp
index eff79d4a..d090fc58 100644
--- a/src/remoteClient/clientContextImpl.cpp
+++ b/src/remoteClient/clientContextImpl.cpp
@@ -20,6 +20,7 @@
 #include <pv/bitSetUtil.h>
 #include <pv/standardPVField.h>
 #include <pv/reftrack.h>
+#include <pv/sharedPtr.h>
 
 #define epicsExportSharedSymbols
 #include <pv/pvAccess.h>
@@ -49,6 +50,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 {
 
@@ -4338,6 +4342,36 @@ class InternalClientContextImpl :
             m_channelSearchManager->newServerDetected();
     }
 
+    /**
+     * Handles cleanup of old beacons.
+     */
+    class BeaconCleanupCallback : public TimerCallback
+    {
+    public:
+        typedef std::tr1::shared_ptr<BeaconCleanupCallback> shared_pointer;
+
+        BeaconCleanupCallback(InternalClientContextImpl& impl, osiSockAddr addr) :
+            m_from(addr),
+            m_impl(impl)
+        {
+        }
+
+        virtual void callback() OVERRIDE
+        {
+            Lock guard(m_impl.m_beaconMapMutex);
+            m_impl.m_beaconHandlers.erase(m_from);
+        }
+
+        virtual void timerStopped() OVERRIDE
+        {
+            this->callback(); // Remove the beacon if the timer is closed.
+        }
+
+    private:
+        osiSockAddr m_from;
+        InternalClientContextImpl& m_impl;
+    };
+
     /**
      * Get (and if necessary create) beacon handler.
      * @param protocol the protocol.
@@ -4351,8 +4385,18 @@ 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 BeaconHandler::shared_pointer();
+            }
+
             // stores weak_ptr
             handler.reset(new BeaconHandler(internal_from_this(), responseFrom));
+            m_timer->scheduleAfterDelay(BeaconCleanupCallback::shared_pointer(new BeaconCleanupCallback(*this, *responseFrom)), maxBeaconLifetime);
             m_beaconHandlers[*responseFrom] = handler;
         }
         else