Preserve and sanitise imported heartbeat strings

Co-authored-by: Amp <amp@ampcode.com>

Author: skyfallwastaken

app/services/heartbeat_import_service.rb

@@ -17,15 +17,15 @@ def self.import_from_file(file_content, user, on_progress: nil, progress_interva
       heartbeat_batch.clear
     end
 
-    handler = HeartbeatSaxHandler.new do |hb|
+    handler = HeartbeatStreamHandler.new do |hb|
       total_count += 1
       on_progress&.call(total_count) if progress_interval.positive? && (total_count % progress_interval).zero?
 
       heartbeat_batch << hb
       flush.call if heartbeat_batch.size >= BATCH_SIZE
     end
 
-    Oj.saj_parse(handler, file_content)
+    Oj.sc_parse(handler, file_content)
     on_progress&.call(total_count)
 
     raise StandardError, "Expected a heartbeat export JSON file." if total_count.zero?
@@ -40,54 +40,56 @@ def self.import_from_file(file_content, user, on_progress: nil, progress_interva
       skipped_count: total_count - imported_count, errors: errors + [ e.message ] }
   end
 
-  class HeartbeatSaxHandler < Oj::Saj
+  # Retain only the current heartbeat, not the surrounding dump or day arrays.
+  # Unlike SAJ, ScHandler preserves embedded NULs for ingestion to sanitize.
+  class HeartbeatStreamHandler < Oj::ScHandler
     def initialize(&block)
       @block = block
       @depth = 0
-      @current_heartbeat = nil
       @heartbeat_array_depths = []
-      @field_array_stack = []
     end
 
-    def hash_start(key)
-      @current_heartbeat = {} if inside_heartbeat_array? && @depth == @heartbeat_array_depths.last + 1
+    def hash_start
+      @key = nil
       @depth += 1
+      {} if @heartbeat_array_depths.any?
     end
 
-    def hash_end(key)
+    def hash_end
       @depth -= 1
-      if inside_heartbeat_array? && @depth == @heartbeat_array_depths.last + 1 && @current_heartbeat
-        @block.call(@current_heartbeat)
-        @current_heartbeat = nil
-      end
     end
 
-    def array_start(key)
-      @heartbeat_array_depths << @depth if key == "heartbeats"
-      if @current_heartbeat && key.present?
-        @current_heartbeat[key] = []
-        @field_array_stack << key
+    def hash_key(key)
+      @key = key
+    end
+
+    def hash_set(hash, key, value)
+      hash[key] = value if hash
+    end
+
+    def array_start
+      container = if @key == "heartbeats" && @heartbeat_array_depths.empty?
+        @heartbeat_array_depths << @depth
+        :heartbeats
+      elsif @heartbeat_array_depths.any?
+        []
       end
+      @key = nil
       @depth += 1
+      container
     end
 
-    def array_end(key)
+    def array_end
       @depth -= 1
-      @heartbeat_array_depths.pop if key == "heartbeats" && @heartbeat_array_depths.last == @depth
-      @field_array_stack.pop if @field_array_stack.last == key
+      @heartbeat_array_depths.pop if @heartbeat_array_depths.last == @depth
     end
 
-    def add_value(value, key)
-      return unless @current_heartbeat
-      if key
-        @current_heartbeat[key] = value
-      elsif @field_array_stack.any?
-        @current_heartbeat[@field_array_stack.last] << value
+    def array_append(array, value)
+      if array == :heartbeats
+        @block.call(value)
+      elsif array
+        array << value
       end
     end
-
-    private
-
-    def inside_heartbeat_array? = @heartbeat_array_depths.any?
   end
 end

app/services/heartbeat_ingest.rb

@@ -237,7 +237,8 @@ def ingest_import
 
   def normalize_imported_heartbeat(heartbeat, placeholder_state: { contexts: {}, last_project: nil })
     hb = heartbeat.respond_to?(:with_indifferent_access) ? heartbeat.with_indifferent_access : heartbeat.to_h.with_indifferent_access
-    user_agent_info = (@user_agents_by_id[hb[:user_agent_id].to_s] || {}).with_indifferent_access
+    hb = strip_null_bytes(hb)
+    user_agent_info = strip_null_bytes((@user_agents_by_id[hb[:user_agent_id].to_s] || {}).with_indifferent_access)
     resolved_user_agent = hb[:user_agent].presence || user_agent_info[:value].presence || hb[:user_agent_id].presence
     parsed_user_agent = parse_user_agent(resolved_user_agent, category: hb[:category])
     derived_ai_editor = parsed_user_agent[:editor].presence if parsed_user_agent[:ai_model].present?

test/services/heartbeat_import_service_test.rb

@@ -1,6 +1,29 @@
 require "test_helper"
 
 class HeartbeatImportServiceTest < ActiveSupport::TestCase
+  test "sanitizes null bytes without losing valid batch rows and deduplicates replays" do
+    user = create(:user)
+    rows = [
+      { entity: "first.rb", project: "api", time: 1_700_000_000.0, type: "file" },
+      { entity: "sec\0ond.rb", project: "api", branch: "ma\0in", dependencies: [ "ra\0ils" ], time: 1_700_000_060.0, type: "file" }
+    ]
+
+    result = HeartbeatImportService.import_from_file(StringIO.new({ heartbeats: rows }.to_json), user)
+    assert result[:success], result[:error]
+    assert_equal 2, result[:imported_count]
+    heartbeat = user.heartbeats.find_by!(entity: "second.rb")
+    assert_equal "main", heartbeat.branch
+    assert_equal [ "rails" ], heartbeat.dependencies
+
+    sanitized = [ rows.first, rows.second.merge(entity: "second.rb", branch: "main", dependencies: [ "rails" ]) ]
+    [ rows, sanitized ].each do |replay|
+      result = HeartbeatImportService.import_from_file({ heartbeats: replay }.to_json, user)
+      assert result[:success], result[:error]
+      assert_equal 0, result[:imported_count]
+    end
+    assert_equal 2, user.heartbeats.count
+  end
+
   test "deduplicates imported heartbeats by fields hash" do
     user = create(:user)
     file_content = {