service_heartbeat.rb in contrast-agent-5.2.0

- old
+ new

@@ -1,10 +1,11 @@
 # Copyright (c) 2022 Contrast Security, Inc. See https://www.contrastsecurity.com/enduser-terms-0317a for more details.
 # frozen_string_literal: true
 
 require 'contrast/components/logger'
 require 'contrast/agent/worker_thread'
+require 'contrast/agent/reporting/report'
 
 module Contrast
   module Agent
     # The ServiceHeartbeat functions to keep the Contrast Service alive and
     # ensure that it maintains this Agent's ApplicationContext.
@@ -12,24 +13,66 @@
       include Contrast::Components::Logger::InstanceMethods
 
       # Spec recommends 30 seconds, we're going with 15.
       REFRESH_INTERVAL_SEC = 15
 
+      # check if we can report to TS
+      #
+      # @return[Boolean] true if bypass is enabled, or false if bypass disabled
+      def enabled?
+        @_enabled = Contrast::CONTRAST_SERVICE.use_agent_communication? if @_enabled.nil?
+        @_enabled
+      end
+
+      def client
+        @_client ||= Contrast::Agent::Reporting::ReporterClient.new
+      end
+
+      def connection
+        @_connection ||= client.initialize_connection
+      end
+
       def start_thread!
         return if running?
 
+        report_from_reporter = check_report_provider
+
         @_thread = Contrast::Agent::Thread.new do
           logger.info('Starting heartbeat thread.')
           loop do
-            Contrast::Agent.messaging_queue.send_event_eventually(poll_message)
+            if report_from_reporter
+              client.send_event(poll_message, connection)
+            else
+              Contrast::Agent.messaging_queue.send_event_eventually(poll_message)
+            end
 
             sleep REFRESH_INTERVAL_SEC
           end
         end
       end
 
       def poll_message
-        @_poll_message ||= Contrast::Api::Dtm::Poll.new
+        @_poll_message ||= if enabled? && client
+                             Contrast::Agent::Reporting::Poll.new
+                           else
+                             Contrast::Api::Dtm::Poll.new
+                           end
+      end
+
+      def check_report_provider
+        return false unless enabled?
+        return false unless client && connection
+
+        client.startup!(connection)
+        true
+      end
+
+      def send_event provider
+        if provider
+          client.send_event(poll_message, connection)
+          return
+        end
+        Contrast::Agent.messaging_queue.send_event_eventually(poll_message)
       end
     end
   end
 end