import org.apache.logging.log4j.message.MapMessage;
import org.apache.logging.log4j.util.StringBuilderFormattable;
import com.google.gson.stream.JsonReader;
import java.io.StringReader;
import java.math.BigInteger;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * Variant analysis harness for CVE-2026-49844.
 *
 * The original repro (NonFiniteJsonTest) exercised only the *direct* scalar and
 * primitive-array paths into {@code MapMessageJsonFormatter.formatNumber} /
 * {@code formatDoubleArray} / {@code formatFloatArray}:
 *   - scalar Double / Float
 *   - double[] / float[]
 *
 * This harness exercises the *alternate* data paths through the
 * {@code MapMessageJsonFormatter.format(sb, object, depth)} dispatcher that were
 * NOT covered by the original repro.  Every value below reaches the same sink
 * ({@code MapMessageJsonFormatter}) via a materially different dispatcher branch:
 *
 *   A. nested java.util.Map value            -> formatMap  -> format(value) -> formatNumber
 *   B. java.util.List value                  -> formatList -> format(item)  -> formatNumber
 *   C. Object[] value                        -> formatObjectArray -> format(item) -> formatNumber
 *   D. non-List Collection (Set) value       -> formatCollection -> format(item) -> formatNumber
 *   E. custom Number subclass (else branch)  -> formatNumber else-branch -> sb.append(doubleValue) / formatDouble
 *   F. BigInteger whose doubleValue() = Inf  -> formatNumber else-branch (BigInteger not special-cased)
 *   G. StringBuilderFormattable emitting NaN -> formatFormattable (wraps in quotes, always quoted)
 *
 * The goal is to determine whether the PR #4163 fix (formatDouble/formatFloat
 * gating) is COMPLETE across all these alternate data paths, or whether any path
 * still emits a bare NaN/Infinity token on the FIXED versions (a bypass).
 *
 * Usage:
 *   java -cp log4j-api.jar:gson.jar:. VariantTest <version-label>
 */
public class VariantTest {

    /** Public concrete subclass of the generic MapMessage (values are Object). */
    public static class TestMapMessage extends MapMessage<TestMapMessage, Object> {
        private static final long serialVersionUID = 1L;
    }

    /**
     * A custom java.lang.Number subclass that is NOT Double/Float/BigDecimal/etc.
     * Its doubleValue() can be non-finite.  This exercises the {@code else}
     * branch of {@code MapMessageJsonFormatter.formatNumber}, which the original
     * repro never reached (it only used boxed Double/Float).
     */
    public static final class CustomNumber extends Number {
        private static final long serialVersionUID = 1L;
        private final double value;
        public CustomNumber(final double value) { this.value = value; }
        @Override public int intValue() { return (int) value; }
        @Override public long longValue() { return (long) value; }
        @Override public float floatValue() { return (float) value; }
        @Override public double doubleValue() { return value; }
        @Override public String toString() { return Double.toString(value); }
    }

    /**
     * A StringBuilderFormattable that deliberately emits a bare "NaN" token from
     * its formatTo().  {@code formatFormattable} wraps the output in quotes and
     * JSON-escapes it, so this should always be a valid JSON string in BOTH the
     * vulnerable and fixed versions.  It is a control case to prove the
     * formattable path is not itself a bypass.
     */
    public static final class NaNEmittingFormattable implements StringBuilderFormattable {
        @Override public void formatTo(final StringBuilder buffer) { buffer.append("NaN"); }
    }

    // --- Result bookkeeping ---
    static int bareNonFiniteCount = 0;
    static int quotedNonFiniteCount = 0;
    static int jsonParseFailures = 0;
    static int totalCases = 0;

    public static void main(final String[] args) {
        final String label = args.length > 0 ? args[0] : "unknown";
        System.out.println("==================================================");
        System.out.println("CVE-2026-49844 VARIANT analysis  |  log4j-api: " + label);
        System.out.println("  (alternate data paths into MapMessageJsonFormatter)");
        System.out.println("==================================================");

        final double[] nonFinite = {Double.NaN, Double.POSITIVE_INFINITY, Double.NEGATIVE_INFINITY};
        final String[] names = {"NaN", "Infinity", "-Infinity"};

        // ---- A. Nested java.util.Map value ----
        System.out.println("\n--- A. nested java.util.Map value (formatMap -> formatNumber) ---");
        for (int i = 0; i < nonFinite.length; i++) {
            final Map<String, Object> inner = new LinkedHashMap<>();
            inner.put("v", nonFinite[i]);
            final TestMapMessage msg = new TestMapMessage();
            msg.with("data", inner);
            analyze(label, "nested-map", names[i], msg.getFormattedMessage(new String[]{"JSON"}));
        }

        // ---- B. java.util.List value ----
        System.out.println("\n--- B. java.util.List value (formatList -> formatNumber) ---");
        for (int i = 0; i < nonFinite.length; i++) {
            final List<Object> list = new ArrayList<>();
            list.add(nonFinite[i]);
            final TestMapMessage msg = new TestMapMessage();
            msg.with("items", list);
            analyze(label, "list", names[i], msg.getFormattedMessage(new String[]{"JSON"}));
        }

        // ---- C. Object[] value ----
        System.out.println("\n--- C. Object[] value (formatObjectArray -> formatNumber) ---");
        for (int i = 0; i < nonFinite.length; i++) {
            final Object[] arr = {nonFinite[i]};
            final TestMapMessage msg = new TestMapMessage();
            msg.with("arr", arr);
            analyze(label, "object-array", names[i], msg.getFormattedMessage(new String[]{"JSON"}));
        }

        // ---- D. non-List Collection (Set) value ----
        System.out.println("\n--- D. non-List Collection (Set) value (formatCollection -> formatNumber) ---");
        for (int i = 0; i < nonFinite.length; i++) {
            final Set<Object> set = new LinkedHashSet<>();
            set.add(nonFinite[i]);
            final TestMapMessage msg = new TestMapMessage();
            msg.with("set", set);
            analyze(label, "set-collection", names[i], msg.getFormattedMessage(new String[]{"JSON"}));
        }
        // Also exercise ArrayDeque (another non-List Collection) for breadth.
        {
            final ArrayDeque<Object> deque = new ArrayDeque<>();
            deque.add(Double.NaN);
            final TestMapMessage msg = new TestMapMessage();
            msg.with("deque", deque);
            analyze(label, "deque-collection", "NaN", msg.getFormattedMessage(new String[]{"JSON"}));
        }

        // ---- E. Custom Number subclass (formatNumber else-branch) ----
        System.out.println("\n--- E. custom Number subclass (formatNumber else-branch) ---");
        for (int i = 0; i < nonFinite.length; i++) {
            final TestMapMessage msg = new TestMapMessage();
            msg.with("custom", new CustomNumber(nonFinite[i]));
            analyze(label, "custom-number", names[i], msg.getFormattedMessage(new String[]{"JSON"}));
        }

        // ---- F. BigInteger whose doubleValue() overflows to Infinity ----
        // BigInteger is NOT special-cased in MapMessageJsonFormatter.formatNumber
        // (unlike JsonWriter, which has an instanceof BigInteger branch).  It falls
        // into the else branch; for a sufficiently large BigInteger, doubleValue()
        // returns Infinity.
        System.out.println("\n--- F. BigInteger overflow -> Infinity (formatNumber else-branch) ---");
        {
            final BigInteger huge = new BigInteger("10").pow(400); // 10^400 >> Double.MAX_VALUE
            final double dv = huge.doubleValue();
            System.out.println("  (sanity) huge.doubleValue() = " + dv
                    + "  (expect Infinity; isFinite=" + Double.isFinite(dv) + ")");
            final TestMapMessage msg = new TestMapMessage();
            msg.with("huge", huge);
            analyze(label, "bigint-overflow", "Infinity", msg.getFormattedMessage(new String[]{"JSON"}));
        }

        // ---- G. StringBuilderFormattable emitting bare NaN (control: should be quoted) ----
        System.out.println("\n--- G. StringBuilderFormattable emitting bare NaN (control: always quoted) ---");
        {
            final TestMapMessage msg = new TestMapMessage();
            msg.with("formattable", new NaNEmittingFormattable());
            analyze(label, "formattable-nan", "NaN", msg.getFormattedMessage(new String[]{"JSON"}));
        }

        // ---- Summary ----
        System.out.println("\n==================================================");
        System.out.println("VARIANT SUMMARY (" + label + ")");
        System.out.println("  total cases tested        : " + totalCases);
        System.out.println("  bare NaN/Infinity tokens  : " + bareNonFiniteCount);
        System.out.println("  quoted NaN/Infinity tokens: " + quotedNonFiniteCount);
        System.out.println("  JSON parse failures       : " + jsonParseFailures);
        System.out.println("==================================================");

        final boolean hasBareTokens = bareNonFiniteCount > 0;
        System.out.println("VERDICT|" + label + "|"
                + (hasBareTokens ? "VULNERABLE" : "FIXED")
                + "|bare=" + bareNonFiniteCount
                + "|quoted=" + quotedNonFiniteCount
                + "|parseFail=" + jsonParseFailures);
    }

    private static void analyze(final String label, final String kind, final String value, final String json) {
        totalCases++;
        System.out.println("  [" + kind + "] value=" + value);
        System.out.println("    JSON output: " + json);

        final boolean hasBare = hasBareNonFiniteToken(json);
        if (hasBare) {
            bareNonFiniteCount++;
            System.out.println("    >> BARE non-finite token detected (invalid JSON per RFC 8259)");
        } else if (containsQuotedNonFinite(json)) {
            quotedNonFiniteCount++;
            System.out.println("    >> non-finite value is QUOTED (valid JSON string)");
        }

        final boolean parsedOk = tryStrictJsonParse(json);
        if (!parsedOk) {
            jsonParseFailures++;
            System.out.println("    >> strict JSON parse FAILED (invalid per RFC 8259)");
        } else {
            System.out.println("    >> strict JSON parse OK (valid RFC 8259 JSON)");
        }
    }

    /**
     * Returns true if the JSON text contains a bare (unquoted) NaN, Infinity, or
     * -Infinity token, scanning character by character tracking string literals.
     */
    static boolean hasBareNonFiniteToken(final String text) {
        boolean inString = false;
        final int n = text.length();
        final String[] tokens = {"NaN", "Infinity", "-Infinity"};
        for (int i = 0; i < n; i++) {
            final char c = text.charAt(i);
            if (inString) {
                if (c == '"') {
                    inString = false;
                }
                continue;
            }
            if (c == '"') {
                inString = true;
                continue;
            }
            for (final String tok : tokens) {
                if (text.regionMatches(i, tok, 0, tok.length())) {
                    final int end = i + tok.length();
                    final char before = (i > 0) ? text.charAt(i - 1) : ' ';
                    final char after = (end < n) ? text.charAt(end) : ' ';
                    if (!isIdentChar(before) && !isIdentChar(after)) {
                        return true;
                    }
                }
            }
        }
        return false;
    }

    static boolean containsQuotedNonFinite(final String text) {
        return text.contains("\"NaN\"") || text.contains("\"Infinity\"") || text.contains("\"-Infinity\"");
    }

    static boolean isIdentChar(final char c) {
        return Character.isLetterOrDigit(c) || c == '_' || c == '$';
    }

    /**
     * Strict (RFC 8259) JSON parse via Gson JsonReader with lenient=false.  Bare
     * NaN / Infinity / -Infinity are rejected; quoted string values are accepted.
     * Uses skipValue() directly because fromJson() internally resets to lenient.
     */
    static boolean tryStrictJsonParse(final String json) {
        try {
            final JsonReader reader = new JsonReader(new StringReader(json));
            reader.setLenient(false);
            reader.skipValue();
            return true;
        } catch (final Throwable t) {
            return false;
        }
    }
}