diff --git a/app/services/heartbeat_import_service.rb b/app/services/heartbeat_import_service.rb index d5a1c7d77..f6ee233d8 100644 --- a/app/services/heartbeat_import_service.rb +++ b/app/services/heartbeat_import_service.rb @@ -17,7 +17,7 @@ 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? @@ -25,7 +25,7 @@ def self.import_from_file(file_content, user, on_progress: nil, progress_interva 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 diff --git a/app/services/heartbeat_ingest.rb b/app/services/heartbeat_ingest.rb index a165179b0..75fe7b8b7 100644 --- a/app/services/heartbeat_ingest.rb +++ b/app/services/heartbeat_ingest.rb @@ -27,7 +27,7 @@ def initialize(user:, mode:, heartbeats:, request_context: {}, user_agents_by_id @mode = mode @heartbeats = heartbeats @request_context = request_context.with_indifferent_access - @user_agents_by_id = user_agents_by_id + @user_agents_by_id = user_agents_by_id.transform_keys { |id| strip_null_bytes(id.to_s) } @schedule_rollup_refresh = schedule_rollup_refresh end @@ -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? diff --git a/test/services/heartbeat_import_service_test.rb b/test/services/heartbeat_import_service_test.rb index e2ff9d090..f5b84e706 100644 --- a/test/services/heartbeat_import_service_test.rb +++ b/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 = { diff --git a/test/services/heartbeat_ingest_test.rb b/test/services/heartbeat_ingest_test.rb index 601fd967d..4c7769239 100644 --- a/test/services/heartbeat_ingest_test.rb +++ b/test/services/heartbeat_ingest_test.rb @@ -16,6 +16,30 @@ class HeartbeatIngestTest < ActiveSupport::TestCase ActiveJob::Base.queue_adapter = @original_queue_adapter end + test "import resolves sanitised user agent identifiers and deduplicates replays" do + user = create(:user) + payload = { time: 1_700_000_000.0, entity: "main.rb", type: "file", user_agent_id: "agent\0-1" } + metadata = { value: "custom\0-client", editor: "ze\0d", os: "lin\0ux" } + agents = { "agent\0-1" => metadata } + + result = HeartbeatIngest.call(user:, mode: :import, heartbeats: [ payload ], user_agents_by_id: agents) + assert_equal 1, result.persisted_count + heartbeat = user.heartbeats.sole + assert_equal "custom-client", heartbeat.user_agent + assert_equal "zed", heartbeat.editor + assert_equal "linux", heartbeat.operating_system + assert_equal [ "agent\0-1" ], agents.keys + + [ agents, { "agent-1" => metadata } ].each do |replay_agents| + assert_no_difference("user.heartbeats.count") do + replay = HeartbeatIngest.call(user:, mode: :import, + heartbeats: [ payload.merge(user_agent_id: "agent-1") ], user_agents_by_id: replay_agents) + assert_equal 1, replay.duplicate_count + assert_equal 0, replay.failed_count + end + end + end + test "direct and imported heartbeats share authoritative rules and deduplicate corrected replays" do expected = { ".env" => "Dotenv", ".gitignore" => "Ignore List", ".rspec" => "Option List",