Memory leak optimizations (#153)

* probable optimizations for leak * fixing probable leaking in to_json * make headers nil after connection * nil check for event processor and disable security on websocket close * clearing to_json residual and csec agent main thread * remobed health check from debug log * probable fix for padrino helpers ErubisTemplate * probable fix * probable fix new * probable fix for stack level too deep error * probable fix for stack level too deep error new * minor change * minor change * try * try diff mechanism for hooking * revert last change * revert * debug logs * fix for stack level too deep * removed debug log
newrelic · Jan 8, 2025 · ff5a3dd · ff5a3dd
1 parent 5bd510d
commit ff5a3dd
Show file tree

Hide file tree

Showing 9 changed files with 57 additions and 34 deletions.
diff --git a/Gemfile_test b/Gemfile_test
@@ -14,6 +14,7 @@ gem 'rubocop-minitest'
 gem 'rubocop-rake'
 gem 'simplecov'
 gem 'railties'
+gem 'ffi', '1.15.5' if RUBY_VERSION < '3.1.0'
 if RUBY_VERSION >= '3.0.0'
   gem 'rails', '~>6.0.0'
 elsif RUBY_VERSION < '2.5.0'
@@ -23,6 +24,7 @@ else
 end
 gem 'loofah', '~> 2.19.0'
 gem 'sinatra'
+gem 'tilt', '~> 2.4.0' if RUBY_VERSION < '2.5.0'
 gem 'padrino'
 gem 'grape'
 gem 'roda'

diff --git a/lib/newrelic_security/agent/agent.rb b/lib/newrelic_security/agent/agent.rb
@@ -82,7 +82,7 @@ def stop_websocket_client_if_open
       end
 
       def start_event_processor
-        @event_processor&.event_dequeue_thread&.kill 
+        @event_processor&.event_dequeue_threads&.each { |t| t&.kill }
         @event_processor&.healthcheck_thread&.kill
         @event_processor = nil
         @event_processor = NewRelic::Security::Agent::Control::EventProcessor.new

diff --git a/lib/newrelic_security/agent/control/collector.rb b/lib/newrelic_security/agent/control/collector.rb
@@ -60,7 +60,7 @@ def collect(case_type, args, event_category = nil, **keyword_args)
             event.lineNumber = stk[0].lineno
           end
           event.stacktrace = stk[0..user_frame_index].map(&:to_s)
-          NewRelic::Security::Agent.agent.event_processor.send_event(event)
+          NewRelic::Security::Agent.agent.event_processor&.send_event(event)
           if event.httpRequest[:headers].key?(NR_CSEC_FUZZ_REQUEST_ID) && event.apiId == event.httpRequest[:headers][NR_CSEC_FUZZ_REQUEST_ID].split(COLON_IAST_COLON)[0] && NewRelic::Security::Agent.agent.iast_client.completed_requests[event.parentId]
             NewRelic::Security::Agent.agent.iast_client.completed_requests[event.parentId] << event.id
           end
@@ -73,6 +73,10 @@ def collect(case_type, args, event_category = nil, **keyword_args)
           else
             NewRelic::Security::Agent.agent.rasp_event_stats.error_count.increment
           end
+        ensure
+          event = nil
+          stk = nil
+          route = nil
         end
 
         private

diff --git a/lib/newrelic_security/agent/control/event_processor.rb b/lib/newrelic_security/agent/control/event_processor.rb
@@ -9,7 +9,7 @@ module Control
 
       class EventProcessor
 
-        attr_accessor :eventQ, :event_dequeue_thread, :healthcheck_thread
+        attr_accessor :eventQ, :event_dequeue_threads, :healthcheck_thread
 
         def initialize
           @first_event = true
@@ -24,18 +24,22 @@ def send_app_info
           NewRelic::Security::Agent.init_logger.info "[STEP-3] => Gathering information about the application"
           app_info = NewRelic::Security::Agent::Control::AppInfo.new
           app_info.update_app_info
-          NewRelic::Security::Agent.logger.info "Sending application info : #{app_info.to_json}"
-          NewRelic::Security::Agent.init_logger.info "Sending application info : #{app_info.to_json}"
+          app_info_json = app_info.to_json
+          NewRelic::Security::Agent.logger.info "Sending application info : #{app_info_json}"
+          NewRelic::Security::Agent.init_logger.info "Sending application info : #{app_info_json}"
           enqueue(app_info)
           app_info = nil
+          app_info_json = nil
         end
 
         def send_application_url_mappings
           application_url_mappings = NewRelic::Security::Agent::Control::ApplicationURLMappings.new
           application_url_mappings.update_application_url_mappings
-          NewRelic::Security::Agent.logger.info "Sending application URL Mappings : #{application_url_mappings.to_json}"
+          application_url_mappings_json = application_url_mappings.to_json
+          NewRelic::Security::Agent.logger.info "Sending application URL Mappings : #{application_url_mappings_json}"
           enqueue(application_url_mappings)
           application_url_mappings = nil
+          application_url_mappings_json = nil
         end
 
         def send_event(event)
@@ -87,15 +91,17 @@ def send_iast_data_transfer_request(iast_data_transfer_request)
         private
 
         def create_dequeue_threads
-          # TODO: Create 3 or more consumers for event sending
-          @event_dequeue_thread = Thread.new do
-            Thread.current.name = "newrelic_security_event_thread"
-            loop do
-              begin
-                data_to_be_sent = @eventQ.pop
-                NewRelic::Security::Agent::Control::WebsocketClient.instance.send(data_to_be_sent)
-              rescue => exception
-                NewRelic::Security::Agent.logger.error "Exception in event pop operation : #{exception.inspect}"
+          @event_dequeue_threads = []
+          3.times do |t|
+            @event_dequeue_threads<< Thread.new do
+              Thread.current.name = "newrelic_security_event_thread-#{t}"
+              loop do
+                begin
+                  data_to_be_sent = @eventQ.pop
+                  NewRelic::Security::Agent::Control::WebsocketClient.instance.send(data_to_be_sent)
+                rescue => exception
+                  NewRelic::Security::Agent.logger.error "Exception in event pop operation : #{exception.inspect}"
+                end
               end
             end
           end

diff --git a/lib/newrelic_security/agent/control/event_subscriber.rb b/lib/newrelic_security/agent/control/event_subscriber.rb
@@ -8,7 +8,11 @@ def initialize
               NewRelic::Security::Agent.init_logger.info "NewRelic server_source_configuration_added for pid : #{Process.pid}, Parent Pid : #{Process.ppid}"
               NewRelic::Security::Agent.config.update_server_config
               if NewRelic::Security::Agent.config[:'security.enabled'] && !NewRelic::Security::Agent.config[:high_security]
-                Thread.new { NewRelic::Security::Agent.agent.scan_scheduler.init_via_scan_scheduler }
+                NewRelic::Security::Agent.agent.event_processor&.event_dequeue_threads&.each { |t| t&.kill }
+                NewRelic::Security::Agent.agent.event_processor = nil
+                @csec_agent_main_thread&.kill
+                @csec_agent_main_thread = nil
+                @csec_agent_main_thread = Thread.new { NewRelic::Security::Agent.agent.scan_scheduler.init_via_scan_scheduler }
               else
                 NewRelic::Security::Agent.logger.info "New Relic Security is disabled by one of the user provided config `security.enabled` or `high_security`."
                 NewRelic::Security::Agent.init_logger.info "New Relic Security is disabled by one of the user provided config `security.enabled` or `high_security`."

diff --git a/lib/newrelic_security/agent/control/websocket_client.rb b/lib/newrelic_security/agent/control/websocket_client.rb
@@ -63,6 +63,7 @@ def connect()
             @ws = connection
 
             connection.on :open do
+              headers = nil
               NewRelic::Security::Agent.logger.debug "Websocket connected with IC, AgentEventMachine #{NewRelic::Security::Agent::Utils.filtered_log(connection.inspect)}"
               NewRelic::Security::Agent.init_logger.info "[STEP-4] => Web socket connection to SaaS validator established successfully"
               NewRelic::Security::Agent.agent.event_processor.send_app_info
@@ -116,25 +117,31 @@ def connect()
         end
 
         def send(message)
-          message_json = message.to_json
-          NewRelic::Security::Agent.logger.debug "Sending #{message.jsonName} : #{message_json}"
-          res = @ws.send(message_json)
-          if res && message.jsonName == :Event
-            NewRelic::Security::Agent.agent.event_sent_count.increment
-            if NewRelic::Security::Agent::Utils.is_IAST_request?(message.httpRequest[:headers])
-              NewRelic::Security::Agent.agent.iast_event_stats.sent.increment
-            else
-              NewRelic::Security::Agent.agent.rasp_event_stats.sent.increment
+          message_json = nil
+          begin
+            message_json = message.to_json
+            NewRelic::Security::Agent.logger.debug "Sending #{message.jsonName} : #{message_json}"
+            res = @ws.send(message_json)
+            if res && message.jsonName == :Event
+              NewRelic::Security::Agent.agent.event_sent_count.increment
+              if NewRelic::Security::Agent::Utils.is_IAST_request?(message.httpRequest[:headers])
+                NewRelic::Security::Agent.agent.iast_event_stats.sent.increment
+              else
+                NewRelic::Security::Agent.agent.rasp_event_stats.sent.increment
+              end
             end
+            NewRelic::Security::Agent.agent.exit_event_stats.sent.increment if res && message.jsonName == :'exit-event'
+          rescue Exception => exception
+            NewRelic::Security::Agent.logger.error "Exception in sending message : #{exception.inspect} #{exception.backtrace}"
+            NewRelic::Security::Agent.agent.event_drop_count.increment if message.jsonName == :Event
+            NewRelic::Security::Agent.agent.event_processor.send_critical_message(exception.message, "SEVERE", caller_locations[0].to_s, Thread.current.name, exception)
+          ensure
+            message_json = nil
           end
-          NewRelic::Security::Agent.agent.exit_event_stats.sent.increment if res && message.jsonName == :'exit-event'
-        rescue Exception => exception
-          NewRelic::Security::Agent.logger.error "Exception in sending message : #{exception.inspect} #{exception.backtrace}"
-          NewRelic::Security::Agent.agent.event_drop_count.increment if message.jsonName == :Event
-          NewRelic::Security::Agent.agent.event_processor.send_critical_message(exception.message, "SEVERE", caller_locations[0].to_s, Thread.current.name, exception)
         end
 
         def close(reconnect = true)
+          NewRelic::Security::Agent.config.disable_security
           NewRelic::Security::Agent.logger.info "Flushing eventQ (#{NewRelic::Security::Agent.agent.event_processor.eventQ.size} events) and closing websocket connection"
           NewRelic::Security::Agent.agent.event_processor&.eventQ&.clear
           @iast_client&.iast_data_transfer_request_processor_thread&.kill

diff --git a/lib/newrelic_security/instrumentation-security/io/chain.rb b/lib/newrelic_security/instrumentation-security/io/chain.rb
@@ -99,9 +99,9 @@ def binwrite(*var)
 
               alias_method :popen_without_security, :popen
 
-              def popen(*var)
+              def popen(*var, &block)
                 retval = nil
-                event = popen_on_enter(*var) { retval = popen_without_security(*var) }
+                event = popen_on_enter(*var) { retval = popen_without_security(*var, &block) }
                 popen_on_exit(event) { return retval }
               end
             end

diff --git a/lib/newrelic_security/instrumentation-security/io/prepend.rb b/lib/newrelic_security/instrumentation-security/io/prepend.rb
@@ -75,7 +75,7 @@ def binwrite(*var)
           binwrite_on_exit(event, retval) { return retval }
         end
 
-        def popen(*var)
+        def popen(*var, &block)
           retval = nil
           event = popen_on_enter(*var) { retval = super }
           popen_on_exit(event) { return retval }

diff --git a/test/newrelic_security/instrumentation-security/io/io_test.rb b/test/newrelic_security/instrumentation-security/io/io_test.rb
@@ -264,7 +264,7 @@ def test_binwrite_arg
                 def test_popen
                     current_path = __dir__
                     cmd = "ls " + current_path
-                    f = IO.popen(cmd)
+                    f = ::IO.popen(cmd)
                     output = f.read
                     #puts output
                     f.close