newrelic_security 0.2.0 → 0.4.0

data/lib/newrelic_security/agent/control/event_processor.rb

@@ -9,7 +9,7 @@ module NewRelic::Security
 
       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 @@ module NewRelic::Security
           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)
@@ -48,6 +52,7 @@ module NewRelic::Security
           enqueue(event)
           if @first_event
             NewRelic::Security::Agent.init_logger.info "[STEP-8] => First event sent for validation. Security agent started successfully : #{event.to_json}"
+            NewRelic::Security::Agent.config.traffic_start_time = current_time_millis unless NewRelic::Security::Agent.config[:traffic_start_time]
             @first_event = false
           end
           event = nil
@@ -64,7 +69,7 @@ module NewRelic::Security
           if exc
             exception = {}
             exception[:message] = exc.message
-            exception[:cause] = exc.cause
+            exception[:cause] = { :message => exc.cause }
             exception[:stackTrace] = exc.backtrace.map(&:to_s)
           end
           critical_message = NewRelic::Security::Agent::Control::CriticalMessage.new(message, level, caller, thread_name, exception)
@@ -86,15 +91,17 @@ module NewRelic::Security
         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
@@ -130,6 +137,10 @@ module NewRelic::Security
           NewRelic::Security::Agent.logger.error "Exception in health check thread, #{exception.inspect}"
         end
 
+        def current_time_millis
+          (Time.now.to_f * 1000).to_i
+        end
+
         def create_error_reporting_thread
           @error_reporting_thread = Thread.new {
             Thread.current.name = "newrelic_security_error_reporting_thread"

data/lib/newrelic_security/agent/control/event_subscriber.rb

@@ -7,19 +7,17 @@ module NewRelic::Security
               NewRelic::Security::Agent.logger.info "NewRelic server_source_configuration_added for pid : #{Process.pid}, Parent Pid : #{Process.ppid}"
               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[:enabled] && !NewRelic::Security::Agent.config[:high_security]
-                NewRelic::Security::Agent.agent.init
+              if NewRelic::Security::Agent.config[:'security.enabled'] && !NewRelic::Security::Agent.config[:high_security]
+                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`."
               end
             }
-            ::NewRelic::Agent.instance.events.subscribe(:security_policy_received) { |received_policy| 
-              NewRelic::Security::Agent.logger.info "security_policy_received pid ::::::: #{Process.pid} #{Process.ppid}, #{received_policy}"
-              NewRelic::Security::Agent.config[:policy].merge!(received_policy)
-              NewRelic::Security::Agent.init_logger.info "[STEP-7] => Received and applied policy/configuration : #{received_policy}"
-              NewRelic::Security::Agent.agent.start_iast_client if NewRelic::Security::Agent::Utils.is_IAST?
-            }
         end
       end
     end

data/lib/newrelic_security/agent/control/health_check.rb

@@ -35,6 +35,10 @@ module NewRelic::Security
           @iastEventStats = {}
           @raspEventStats = {}
           @exitEventStats = {}
+          @procStartTime = NewRelic::Security::Agent.config[:process_start_time]
+          @trafficStartedTime = NewRelic::Security::Agent.config[:traffic_start_time]
+          @scanStartTime = NewRelic::Security::Agent.config[:scan_start_time]
+          @iastTestIdentifer = NewRelic::Security::Agent.config[:'security.iast_test_identifier']
         end
 
         def as_json

data/lib/newrelic_security/agent/control/http_context.rb

@@ -12,17 +12,20 @@ module NewRelic::Security
       REQUEST_METHOD = 'REQUEST_METHOD'
       HTTP_HOST = 'HTTP_HOST'
       PATH_INFO = 'PATH_INFO'
+      QUERY_STRING = 'QUERY_STRING'
       RACK_INPUT = 'rack.input'
       CGI_VARIABLES = ::Set.new(%W[ AUTH_TYPE CONTENT_LENGTH CONTENT_TYPE GATEWAY_INTERFACE HTTPS HTTP_HOST PATH_INFO PATH_TRANSLATED REQUEST_URI QUERY_STRING REMOTE_ADDR REMOTE_HOST REMOTE_IDENT REMOTE_USER REQUEST_METHOD SCRIPT_NAME SERVER_NAME SERVER_PORT SERVER_PROTOCOL SERVER_SOFTWARE rack.url_scheme ])
+      REQUEST_BODY_LIMIT = 500 #KB
 
       class HTTPContext
         
-        attr_accessor :time_stamp, :req, :method, :headers, :params, :body, :data_truncated, :route, :cache, :fuzz_files, :event_counter, :mutex
+        attr_accessor :time_stamp, :req, :method, :headers, :params, :body, :data_truncated, :route, :cache, :fuzz_files, :event_counter, :custom_data_type, :mutex, :url
 
         def initialize(env)
           @time_stamp = current_time_millis
           @req = env.select { |key, _| CGI_VARIABLES.include? key}
           @method = @req[REQUEST_METHOD]
+          @url = "#{@req[PATH_INFO]}?#{@req[QUERY_STRING]}"
           @headers = env.select { |key, _| key.include?(HTTP_) }
           @headers = @headers.transform_keys{ |key| key[5..-1].gsub(UNDERSCORE, HYPHEN).downcase }
           request = Rack::Request.new(env) unless env.empty?
@@ -31,20 +34,21 @@ module NewRelic::Security
           strio = env[RACK_INPUT]
           if strio.instance_of?(::StringIO)
 						offset = strio.tell
-						@body = strio.read(NewRelic::Security::Agent.config[:'security.request.body_limit'] * 1024) #after read, offset changes
+						@body = strio.read(REQUEST_BODY_LIMIT * 1024) #after read, offset changes
 						strio.seek(offset)
             # In case of Grape and Roda strio.read giving empty result, added below approach to handle such cases
             @body = strio.string if @body.nil? && strio.size > 0
           elsif defined?(::Rack) && defined?(::Rack::Lint::InputWrapper) && strio.instance_of?(::Rack::Lint::InputWrapper)
-						@body = strio.read(NewRelic::Security::Agent.config[:'security.request.body_limit'] * 1024)
+						@body = strio.read(REQUEST_BODY_LIMIT * 1024)
           elsif defined?(::Protocol::Rack::Input) && defined?(::Protocol::Rack::Input) && strio.instance_of?(::Protocol::Rack::Input)
-            @body = strio.read(NewRelic::Security::Agent.config[:'security.request.body_limit'] * 1024)
+            @body = strio.read(REQUEST_BODY_LIMIT * 1024)
           elsif defined?(::PhusionPassenger::Utils::TeeInput) && strio.instance_of?(::PhusionPassenger::Utils::TeeInput)
-						@body = strio.read(NewRelic::Security::Agent.config[:'security.request.body_limit'] * 1024)
+						@body = strio.read(REQUEST_BODY_LIMIT * 1024)
           end
-          @data_truncated = @body && @body.size >= NewRelic::Security::Agent.config[:'security.request.body_limit'] * 1024
+          @data_truncated = @body && @body.size >= REQUEST_BODY_LIMIT * 1024
 					strio&.rewind
 					@body = @body.force_encoding(Encoding::UTF_8) if @body.is_a?(String)
+          @custom_data_type = {}
           @cache = Hash.new
           @fuzz_files = ::Set.new
           @event_counter = 0

data/lib/newrelic_security/agent/control/iast_client.rb

@@ -17,9 +17,8 @@ module NewRelic::Security
       IS_GRPC = 'isGrpc'
       INPUT_CLASS = 'inputClass'
       SERVER_PORT_1 = 'serverPort'
-      PROBING = 'probing'
-      INTERVAL = 'interval'
       IS_GRPC_CLIENT_STREAM = 'isGrpcClientStream'
+      PROBING_INTERVAL = 5
 
       class IASTClient
         
@@ -54,6 +53,7 @@ module NewRelic::Security
               Thread.current.name = "newrelic_security_iast_thread-#{t}"
               loop do
                 fuzz_request = @fuzzQ.deq #thread blocks when the queue is empty
+                NewRelic::Security::Agent.config.scan_start_time = current_time_millis unless NewRelic::Security::Agent.config[:scan_start_time]
                 if fuzz_request.request[IS_GRPC]
                   fire_grpc_request(fuzz_request.id, fuzz_request.request, fuzz_request.reflected_metadata)
                 else
@@ -71,7 +71,8 @@ module NewRelic::Security
           @iast_data_transfer_request_processor_thread = Thread.new do
             Thread.current.name = "newrelic_security_iast_data_transfer_request_processor"
             loop do
-              sleep NewRelic::Security::Agent.config[:policy][VULNERABILITY_SCAN][IAST_SCAN][PROBING][INTERVAL]
+              # TODO: Check & remove this probing interval if not required, earlier this was used from policy sent by SE.
+              sleep PROBING_INTERVAL
               current_timestamp = current_time_millis
               cooldown_sleep_time = @cooldown_till_timestamp - current_timestamp
               sleep cooldown_sleep_time/1000 if cooldown_sleep_time > 0
@@ -84,9 +85,21 @@ module NewRelic::Security
               if batch_size > 100 && remaining_record_capacity > batch_size
                 iast_data_transfer_request = NewRelic::Security::Agent::Control::IASTDataTransferRequest.new
                 iast_data_transfer_request.batchSize = batch_size * 2
+                # TODO: Below calculation of batch_size overrides above logic and can be removed once below one is stablises or rate limit feature is released.
+                if NewRelic::Security::Agent.config[:'security.scan_controllers.iast_scan_request_rate_limit']
+                  batch_size =
+                    if NewRelic::Security::Agent.config[:'security.scan_controllers.iast_scan_request_rate_limit'] < 12
+                      1
+                    elsif NewRelic::Security::Agent.config[:'security.scan_controllers.iast_scan_request_rate_limit'] > 3600
+                      300
+                    else
+                      NewRelic::Security::Agent.config[:'security.scan_controllers.iast_scan_request_rate_limit'] / 12
+                    end
+                  iast_data_transfer_request.batchSize = batch_size
+                end
                 iast_data_transfer_request.pendingRequestIds = pending_request_ids.to_a
                 iast_data_transfer_request.completedRequests = completed_requests
-                NewRelic::Security::Agent.agent.event_processor.send_iast_data_transfer_request(iast_data_transfer_request)
+                NewRelic::Security::Agent.agent.event_processor&.send_iast_data_transfer_request(iast_data_transfer_request) if NewRelic::Security::Agent::Control::WebsocketClient.instance.is_open?
               end
             end
           end
@@ -122,18 +135,18 @@ module NewRelic::Security
         def fire_grpc_request(fuzz_request_id, request, reflected_metadata)
           service = Object.const_get(request[METHOD].split(SLASH)[0]).superclass
           method = request[METHOD].split(SLASH)[1]
-          @stub = service.rpc_stub_class.new("localhost:#{request[SERVER_PORT_1]}", :this_channel_is_insecure) unless @stub
+          @stub ||= service.rpc_stub_class.new("localhost:#{request[SERVER_PORT_1]}", :this_channel_is_insecure)
 
-          parsed_body =  request[BODY][1..-2].split(',')
-          if reflected_metadata[IS_GRPC_CLIENT_STREAM]
-            chunks_enum = Enumerator.new do |y|
+          parsed_body = request[BODY][1..-2].split(',')
+          chunks_enum = if reflected_metadata[IS_GRPC_CLIENT_STREAM]
+            Enumerator.new do |y|
               parsed_body.each do |b|
                 y << Object.const_get(reflected_metadata[INPUT_CLASS]).decode_json(b)
               end
             end
-          else
-            chunks_enum = Object.const_get(reflected_metadata[INPUT_CLASS]).decode_json(request[BODY])
-          end
+                        else
+            Object.const_get(reflected_metadata[INPUT_CLASS]).decode_json(request[BODY])
+                        end
           response = @stub.public_send(method, chunks_enum, metadata: request[HEADERS])
           # response = @stub.send(method, JSON.parse(request['body'], object_class: OpenStruct))
           # request[HEADERS].delete(VERSION) if request[HEADERS].key?(VERSION)

data/lib/newrelic_security/agent/control/reflected_xss.rb

@@ -108,8 +108,8 @@ module NewRelic::Security
                 processed_data.add(body)
               end
             when APPLICATION_XML
-              # Unescaping of xml data is remaining
-              processed_data.add(body)
+              xml_data = ::CGI.unescapeHTML(body)
+              processed_data.add(xml_data)
             when APPLICATION_X_WWW_FORM_URLENCODED
               body = ::CGI.unescape(body, UTF_8)
               processed_data.add(body)
@@ -134,7 +134,7 @@ module NewRelic::Security
               # do while loop in java code here
               old_processed_body = processed_body
               body = ::JSON.parse(processed_body)
-              processed_data.add(body) if old_processed_body != body && body.to_s.include?(LESS_THAN)
+              processed_data.add(body.to_s) if old_processed_body != body && body.to_s.include?(LESS_THAN)
             when APPLICATION_XML
               # Unescaping of xml data is remaining
               processed_data.add(processed_data)
@@ -176,7 +176,6 @@ module NewRelic::Security
           start_pos = 0
           tmp_curr_pos = 0
           tmp_start_pos = 0
-      
           while curr_pos < data.length
             matcher = TAG_NAME_REGEX.match(data, curr_pos)
             is_attack_construct = false

data/lib/newrelic_security/agent/control/scan_scheduler.rb

@@ -0,0 +1,77 @@
+require 'newrelic_security/parse-cron/cron_parser'
+
+module NewRelic::Security
+  module Agent
+    module Control
+      class ScanScheduler
+        def init_via_scan_scheduler
+          if NewRelic::Security::Agent.config[:'security.scan_schedule.delay'].positive?
+            NewRelic::Security::Agent.logger.info "IAST delay is set to: #{NewRelic::Security::Agent.config[:'security.scan_schedule.delay']}, current time: #{Time.now}"
+            start_agent_with_delay(NewRelic::Security::Agent.config[:'security.scan_schedule.delay']*60)
+          elsif !NewRelic::Security::Agent.config[:'security.scan_schedule.schedule'].to_s.empty?
+            cron_expression_task(NewRelic::Security::Agent.config[:'security.scan_schedule.schedule'], NewRelic::Security::Agent.config[:'security.scan_schedule.duration']*60)
+          else
+            NewRelic::Security::Agent.agent.init
+            NewRelic::Security::Agent.agent.start_iast_client if NewRelic::Security::Agent::Utils.is_IAST?
+            shutdown_at_duration_reached(NewRelic::Security::Agent.config[:'security.scan_schedule.duration']*60)
+          end
+        rescue StandardError => exception
+          NewRelic::Security::Agent.logger.error "Exception in IAST scan scheduler: #{exception.inspect} #{exception.backtrace}"
+          ::NewRelic::Agent.notice_error(exception)
+        end
+
+        def start_agent_with_delay(delay)
+          NewRelic::Security::Agent.logger.info "Security Agent delay scan time is set to: #{(delay/60).ceil} minutes when always_sample_traces is #{NewRelic::Security::Agent.config[:'security.scan_schedule.always_sample_traces']}, current time: #{Time.now}"
+          if NewRelic::Security::Agent.config[:'security.scan_schedule.always_sample_traces']
+            NewRelic::Security::Agent.agent.init unless NewRelic::Security::Agent.config[:enabled]
+            sleep delay if NewRelic::Security::Agent.config[:'security.scan_schedule.always_sample_traces']
+          else
+            sleep delay
+            NewRelic::Security::Agent.agent.init
+          end
+          NewRelic::Security::Agent.agent.start_iast_client if NewRelic::Security::Agent::Utils.is_IAST?
+          shutdown_at_duration_reached(NewRelic::Security::Agent.config[:'security.scan_schedule.duration']*60)
+        end
+
+        def shutdown_at_duration_reached(duration)
+          shutdown_at = Time.now.to_i + duration
+          shut_down_time = (Time.now + duration).strftime("%a %d %b %Y %H:%M:%S")
+          return if duration <= 0
+          NewRelic::Security::Agent.logger.info "IAST Duration is set to: #{duration/60} minutes, timestamp: #{shut_down_time} time, current time: #{Time.now}"
+          @shutdown_monitor_thread = Thread.new do
+            Thread.current.name = "newrelic_security_shutdown_monitor_thread"
+            loop do
+              sleep 1
+              next if Time.now.to_i < shutdown_at
+              if NewRelic::Security::Agent.config[:'security.scan_schedule.always_sample_traces']
+                NewRelic::Security::Agent.logger.info "Shutdown IAST Data transfer request processor only as 'security.scan_schedule.always_sample_traces' is #{NewRelic::Security::Agent.config[:'security.scan_schedule.always_sample_traces']} now at current time: #{Time.now}"
+                NewRelic::Security::Agent.agent.iast_client&.iast_data_transfer_request_processor_thread&.kill
+              else
+                NewRelic::Security::Agent.logger.info "Shutdown IAST agent now at current time: #{Time.now}"
+                ::NewRelic::Agent.notice_error(StandardError.new("WS Connection closed by local"))
+                NewRelic::Security::Agent.agent.shutdown_security_agent
+              end
+              break
+            end
+          end
+        end
+
+        def cron_expression_task(schedule, duration)
+          @cron_parser = NewRelic::Security::ParseCron::CronParser.new(schedule)
+          loop do
+            next_run = @cron_parser.next(Time.now)
+            NewRelic::Security::Agent.logger.info "Next init via cron exp: #{schedule},  is scheduled at : #{next_run}"
+            delay = next_run - Time.now
+            if NewRelic::Security::Agent.agent.iast_client&.iast_data_transfer_request_processor_thread&.alive?
+              sleep delay > 2 ? delay - 2 : delay
+            else
+              start_agent_with_delay(delay) 
+            end
+            return if duration <= 0
+          end
+        end
+        
+      end
+    end
+  end
+end

data/lib/newrelic_security/agent/control/websocket_client.rb

@@ -21,6 +21,10 @@ module NewRelic::Security
       NR_CSEC_ENTITY_NAME = 'NR-CSEC-ENTITY-NAME'
       NR_CSEC_ENTITY_GUID = 'NR-CSEC-ENTITY-GUID'
       NR_CSEC_IAST_DATA_TRANSFER_MODE = 'NR-CSEC-IAST-DATA-TRANSFER-MODE'
+      NR_CSEC_IGNORED_VUL_CATEGORIES = 'NR-CSEC-IGNORED-VUL-CATEGORIES'
+      NR_CSEC_PROCESS_START_TIME = 'NR-CSEC-PROCESS-START-TIME'
+      NR_CSEC_IAST_SCAN_INSTANCE_COUNT = 'NR-CSEC-IAST-SCAN-INSTANCE-COUNT'
+      NR_CSEC_IAST_TEST_IDENTIFIER = 'NR-CSEC-IAST-TEST-IDENTIFIER'
 
       class WebsocketClient
         include Singleton
@@ -43,6 +47,13 @@ module NewRelic::Security
           headers[NR_CSEC_ENTITY_NAME] = NewRelic::Security::Agent.config[:app_name]
           headers[NR_CSEC_ENTITY_GUID] = NewRelic::Security::Agent.config[:entity_guid]
           headers[NR_CSEC_IAST_DATA_TRANSFER_MODE] = PULL
+          headers[NR_CSEC_IGNORED_VUL_CATEGORIES] = ingnored_vul_categories.join(COMMA)
+          headers[NR_CSEC_PROCESS_START_TIME] = NewRelic::Security::Agent.config[:process_start_time]
+          headers[NR_CSEC_IAST_SCAN_INSTANCE_COUNT] = NewRelic::Security::Agent.config[:'security.scan_controllers.scan_instance_count']
+          if NewRelic::Security::Agent.config[:'security.iast_test_identifier'] && !NewRelic::Security::Agent.config[:'security.iast_test_identifier'].empty?
+            headers[NR_CSEC_IAST_TEST_IDENTIFIER] = NewRelic::Security::Agent.config[:'security.iast_test_identifier']
+            headers[NR_CSEC_IAST_SCAN_INSTANCE_COUNT] = 1
+          end
 
           begin
             cert_store = ::OpenSSL::X509::Store.new
@@ -50,13 +61,16 @@ module NewRelic::Security
             NewRelic::Security::Agent.logger.info "Websocket connection URL : #{NewRelic::Security::Agent.config[:validator_service_url]}"
             connection = NewRelic::Security::WebSocket::Client::Simple.connect NewRelic::Security::Agent.config[:validator_service_url], headers: headers, cert_store: cert_store
             @ws = connection
+            @mutex = Mutex.new
 
             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
               NewRelic::Security::Agent.agent.event_processor.send_application_url_mappings
               NewRelic::Security::Agent.config.enable_security
+              NewRelic::Security::Agent::Control::WebsocketClient.instance.start_ping_thread
             end
         
             connection.on :message do |msg|
@@ -71,13 +85,14 @@ module NewRelic::Security
             connection.on :close do |e|
               NewRelic::Security::Agent.logger.info "Closing websocket connection : #{e.inspect}\n"
               NewRelic::Security::Agent.config.disable_security
-              Thread.new { NewRelic::Security::Agent.agent.reconnect(0) } if e
+              reconnect_interval = e.instance_of?(TrueClass) ? 0 : 15
+              Thread.new { NewRelic::Security::Agent.agent.reconnect(reconnect_interval) } if e
             end
 
             connection.on :error do |e|
               NewRelic::Security::Agent.logger.error "Error in websocket connection : #{e.inspect} #{e.backtrace}"
               ::NewRelic::Agent.notice_error(e)
-              Thread.new { NewRelic::Security::Agent::Control::WebsocketClient.instance.close(true) }
+              Thread.new { NewRelic::Security::Agent::Control::WebsocketClient.instance.close(e) }
             end
           rescue Errno::EPIPE => exception
             NewRelic::Security::Agent.logger.error "Unable to connect to validator_service: #{exception.inspect}"
@@ -103,25 +118,37 @@ module NewRelic::Security
         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}"
+            @mutex.synchronize do
+              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'
             end
+          rescue Exception => exception
+            NewRelic::Security::Agent.logger.error "Exception in sending message : #{exception.inspect} #{exception.backtrace}, message: #{message_json}"
+            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
+          stop_ping_thread
           @ws.close(reconnect) if @ws
         end
 
@@ -130,6 +157,34 @@ module NewRelic::Security
           false
         end
 
+        def start_ping_thread
+          @ping_thread = Thread.new do
+            loop do
+              sleep 30
+              @ws.send(EMPTY_STRING, :type => :ping)
+            end
+          end
+        end
+
+        private
+
+        def stop_ping_thread
+          @ping_thread&.kill
+          @ping_thread = nil
+        end
+
+        def ingnored_vul_categories
+          list = []
+          list << FILE_OPERATION << FILE_INTEGRITY if NewRelic::Security::Agent.config[:'security.exclude_from_iast_scan.iast_detection_category.invalid_file_access']
+          list << SQL_DB_COMMAND if NewRelic::Security::Agent.config[:'security.exclude_from_iast_scan.iast_detection_category.sql_injection']
+          list << NOSQL_DB_COMMAND if NewRelic::Security::Agent.config[:'security.exclude_from_iast_scan.iast_detection_category.nosql_injection']
+          list << LDAP if NewRelic::Security::Agent.config[:'security.exclude_from_iast_scan.iast_detection_category.ldap_injection']
+          list << SYSTEM_COMMAND if NewRelic::Security::Agent.config[:'security.exclude_from_iast_scan.iast_detection_category.command_injection']
+          list << XPATH if NewRelic::Security::Agent.config[:'security.exclude_from_iast_scan.iast_detection_category.xpath_injection']
+          list << HTTP_REQUEST if NewRelic::Security::Agent.config[:'security.exclude_from_iast_scan.iast_detection_category.ssrf']
+          list << REFLECTED_XSS if NewRelic::Security::Agent.config[:'security.exclude_from_iast_scan.iast_detection_category.rxss']
+          list
+        end
       end
     end
   end

data/lib/newrelic_security/agent/utils/agent_utils.rb

@@ -15,8 +15,7 @@ module NewRelic::Security
       ASTERISK = '*'
 
       def is_IAST?
-        return false if NewRelic::Security::Agent.config[:policy].empty?
-        return NewRelic::Security::Agent.config[:policy][VULNERABILITY_SCAN][IAST_SCAN][ENABLED] if NewRelic::Security::Agent.config[:policy][VULNERABILITY_SCAN][ENABLED]
+        return true if NewRelic::Security::Agent.config[:mode] == IAST
         false
       end
 
@@ -96,7 +95,8 @@ module NewRelic::Security
 
       def get_app_routes(framework, router = nil)
         enable_object_space_in_jruby
-        if framework == :rails
+        case framework
+        when :rails
           ::Rails.application.routes.routes.each do |route|
             if route.verb.is_a?(::Regexp)
               method = route.verb.inspect.match(/[a-zA-Z]+/)
@@ -107,33 +107,41 @@ module NewRelic::Security
               }
             end
           end
-        elsif framework == :sinatra
+        when :sinatra
           ::Sinatra::Application.routes.each do |method, routes|
             routes.map { |r| r.first.to_s }.map do |route|
               NewRelic::Security::Agent.agent.route_map << "#{method}@#{route}"
             end
           end
-        elsif framework == :grape
-          ObjectSpace.each_object(::Grape::Endpoint) { |z|
-            z.instance_variable_get(:@routes)&.each { |route|
-              http_method = route.instance_variable_get(:@request_method) || route.instance_variable_get(:@options)[:method]
-              NewRelic::Security::Agent.agent.route_map << "#{http_method}@#{route.pattern.origin}"
-            }
-          }
-        elsif framework == :padrino
+        when :grape
+          if defined?(::Grape::API)
+            ObjectSpace.each_object(Class).select { |klass| klass < ::Grape::API }.each do |api_class|
+              api_class.routes.each do |route|
+                http_method = route.request_method || route.options[:method]
+                NewRelic::Security::Agent.agent.route_map << "#{http_method}@#{route.pattern.origin}"
+              end
+            end
+          end
+        when :padrino
           if router.instance_of?(::Padrino::PathRouter::Router)
             router.instance_variable_get(:@routes).each do |route|
               NewRelic::Security::Agent.agent.route_map << "#{route.instance_variable_get(:@verb)}@#{route.matcher.instance_variable_get(:@path)}"
             end
           end
-        elsif framework == :roda
-          NewRelic::Security::Agent.logger.warn "TODO: Roda is a routing tree web toolkit, which generates route dynamically, hence route extraction is not possible."
+        when :roda
+          NewRelic::Security::Agent.logger.debug "TODO: Roda is a routing tree web toolkit, which generates route dynamically, hence route extraction is not possible."
+        when :grpc
+          router.owner.superclass.public_instance_methods(false).each do |m|
+            NewRelic::Security::Agent.agent.route_map << "*@/#{router.owner}/#{m}"
+          end
+        when :rack
+          # TODO: API enpointes(routes) extraction for rack
         else
           NewRelic::Security::Agent.logger.error "Unable to get app routes as Framework not detected"
         end
         disable_object_space_in_jruby if NewRelic::Security::Agent.config[:jruby_objectspace_enabled]
         NewRelic::Security::Agent.logger.debug "ALL ROUTES : #{NewRelic::Security::Agent.agent.route_map}"
-        NewRelic::Security::Agent.agent.event_processor&.send_application_url_mappings
+        NewRelic::Security::Agent.agent.event_processor&.send_application_url_mappings unless NewRelic::Security::Agent.agent.route_map.empty?
       rescue Exception => exception
         NewRelic::Security::Agent.logger.error "Error in get app routes : #{exception.inspect} #{exception.backtrace}"
       end
@@ -195,14 +203,14 @@ module NewRelic::Security
       end
 
       def enable_object_space_in_jruby
-        if RUBY_ENGINE == 'jruby' && !JRuby.objectspace
+        if RUBY_ENGINE == 'jruby' && JRuby.respond_to?(:objectspace) && !JRuby.objectspace
           JRuby.objectspace = true
           NewRelic::Security::Agent.config.jruby_objectspace_enabled = true
         end
       end
 
       def disable_object_space_in_jruby
-        if RUBY_ENGINE == 'jruby' && JRuby.objectspace
+        if RUBY_ENGINE == 'jruby' && JRuby.respond_to?(:objectspace) && JRuby.objectspace
           JRuby.objectspace = false
           NewRelic::Security::Agent.config.jruby_objectspace_enabled = false
         end

data/lib/newrelic_security/constants.rb

@@ -17,6 +17,7 @@ module NewRelic::Security
   NR_CSEC_FUZZ_REQUEST_ID = 'nr-csec-fuzz-request-id'
   NR_CSEC_TRACING_DATA = 'nr-csec-tracing-data'
   NR_CSEC_PARENT_ID = 'nr-csec-parent-id'
+  IAST = 'IAST'
   COLON_IAST_COLON = ':IAST:'
   NOSQL_DB_COMMAND = 'NOSQL_DB_COMMAND'
   SQL_DB_COMMAND = 'SQL_DB_COMMAND'
@@ -63,6 +64,4 @@ module NewRelic::Security
   CONTENT_TYPE1 = 'content-Type'
   PULL = 'PULL'
   SHA1 = 'sha1'
-  VULNERABILITY_SCAN = 'vulnerabilityScan'
-  IAST_SCAN = 'iastScan'
 end

data/lib/newrelic_security/instrumentation-security/async-http/instrumentation.rb

@@ -6,22 +6,11 @@ module NewRelic::Security
   module Instrumentation
     module AsyncHttp
 
-      def call_on_enter(method, url, headers, body)
+      def call_on_enter(_method, url, headers, _body)
         event = nil
         NewRelic::Security::Agent.logger.debug "OnEnter : #{self.class}.#{__method__}"
-        ob = {}
-        ob[:Method] = method
         uri = ::URI.parse url
-        ob[:scheme]  = uri.scheme
-        ob[:host]    = uri.host
-        ob[:port]    = uri.port
-        ob[:URI]     = uri.to_s
-        ob[:path]    = uri.path
-        ob[:query]   = uri.query
-        ob[:Body] = body.respond_to?(:join) ? body.join.to_s : body.to_s
-        ob[:Headers] = headers.to_h
-        ob.each { |_, value| value.dup.force_encoding(ISO_8859_1).encode(UTF_8) if value.is_a?(String) }
-        event = NewRelic::Security::Agent::Control::Collector.collect(HTTP_REQUEST, [ob])
+        event = NewRelic::Security::Agent::Control::Collector.collect(HTTP_REQUEST, [uri.to_s])
         NewRelic::Security::Instrumentation::InstrumentationUtils.append_tracing_data(headers, event) if event
         event
       rescue => exception