import org.apache.logging.log4j.message.MapMessage;
import com.google.gson.stream.JsonReader;
import java.io.StringReader;

/**
 * Reproduction test for CVE-2026-49844.
 *
 * Demonstrates that Apache Log4j API MapMessage JSON serialization
 * emits bare NaN / Infinity / -Infinity tokens (invalid per RFC 8259)
 * when a logged map contains non-finite floating-point values.
 *
 * Entry points exercised (both reach MapMessageJsonFormatter.format):
 *   1. MapMessage.getFormattedMessage(new String[]{"JSON"})  (public API)
 *   2. MapMessage.asJson(StringBuilder)                       (protected, called via subclass)
 *
 * Usage:
 *   java -cp log4j-api.jar:gson.jar:. NonFiniteJsonTest <version-label>
 */
public class NonFiniteJsonTest {

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

        // Expose the protected asJson(StringBuilder) for direct testing.
        public String callAsJsonDirectly() {
            StringBuilder sb = new StringBuilder();
            asJson(sb);
            return sb.toString();
        }
    }

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

    public static void main(String[] args) {
        String label = args.length > 0 ? args[0] : "unknown";
        System.out.println("==================================================");
        System.out.println("CVE-2026-49844 reproduction  |  log4j-api version: " + label);
        System.out.println("==================================================");

        double[] nonFiniteDoubles = {Double.NaN, Double.POSITIVE_INFINITY, Double.NEGATIVE_INFINITY};
        float[]  nonFiniteFloats  = {Float.NaN,  Float.POSITIVE_INFINITY,  Float.NEGATIVE_INFINITY};

        // ---- 1. getFormattedMessage(["JSON"]) with scalar doubles ----
        System.out.println("\n--- Path 1: getFormattedMessage([\"JSON\"]) with scalar Double values ---");
        for (double d : nonFiniteDoubles) {
            TestMapMessage msg = new TestMapMessage();
            msg.with("number", d);
            String json = msg.getFormattedMessage(new String[]{"JSON"});
            analyze(label, "double-scalar", String.valueOf(d), json);
        }

        // ---- 2. getFormattedMessage(["JSON"]) with double[] arrays ----
        System.out.println("\n--- Path 1: getFormattedMessage([\"JSON\"]) with double[] arrays ---");
        for (double d : nonFiniteDoubles) {
            TestMapMessage msg = new TestMapMessage();
            msg.with("numbers", new double[]{d});
            String json = msg.getFormattedMessage(new String[]{"JSON"});
            analyze(label, "double-array", String.valueOf(d), json);
        }

        // ---- 3. getFormattedMessage(["JSON"]) with scalar floats ----
        System.out.println("\n--- Path 1: getFormattedMessage([\"JSON\"]) with scalar Float values ---");
        for (float f : nonFiniteFloats) {
            TestMapMessage msg = new TestMapMessage();
            msg.with("number", f);
            String json = msg.getFormattedMessage(new String[]{"JSON"});
            analyze(label, "float-scalar", String.valueOf(f), json);
        }

        // ---- 4. getFormattedMessage(["JSON"]) with float[] arrays ----
        System.out.println("\n--- Path 1: getFormattedMessage([\"JSON\"]) with float[] arrays ---");
        for (float f : nonFiniteFloats) {
            TestMapMessage msg = new TestMapMessage();
            msg.with("numbers", new float[]{f});
            String json = msg.getFormattedMessage(new String[]{"JSON"});
            analyze(label, "float-array", String.valueOf(f), json);
        }

        // ---- 5. Direct asJson(StringBuilder) call (protected path) ----
        System.out.println("\n--- Path 2: asJson(StringBuilder) direct call (protected) ---");
        TestMapMessage msg = new TestMapMessage();
        msg.with("number", Double.NaN);
        msg.with("inf", Double.POSITIVE_INFINITY);
        msg.with("ninf", Double.NEGATIVE_INFINITY);
        String json = msg.callAsJsonDirectly();
        analyze(label, "asJson-direct", "NaN/Inf/-Inf", json);

        // ---- Summary ----
        System.out.println("\n==================================================");
        System.out.println("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("==================================================");

        // Emit a machine-readable verdict line.
        boolean hasBareTokens = bareNonFiniteCount > 0;
        System.out.println("VERDICT|" + label + "|" + (hasBareTokens ? "VULNERABLE" : "FIXED")
                + "|bare=" + bareNonFiniteCount + "|quoted=" + quotedNonFiniteCount
                + "|parseFail=" + jsonParseFailures);
    }

    /**
     * Analyze one JSON output string:
     *  - print it
     *  - detect bare (unquoted) NaN / Infinity / -Infinity tokens
     *  - attempt strict JSON parse via Gson
     */
    private static void analyze(String label, String kind, String value, String json) {
        totalCases++;
        System.out.println("  [" + kind + "] value=" + value);
        System.out.println("    JSON output: " + json);

        // Detect bare (unquoted) non-finite tokens.
        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)");
        }

        // Strict JSON parse attempt (RFC 8259 compliant, lenient=false).
        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.
     * Scans character by character tracking whether we are inside a string literal.
     */
    static boolean hasBareNonFiniteToken(String text) {
        boolean inString = false;
        int n = text.length();
        String[] tokens = {"NaN", "Infinity", "-Infinity"};
        for (int i = 0; i < n; i++) {
            char c = text.charAt(i);
            if (inString) {
                if (c == '"') inString = false;
                continue;
            }
            if (c == '"') { inString = true; continue; }
            // Try to match a bare token at position i.
            for (String tok : tokens) {
                if (text.regionMatches(i, tok, 0, tok.length())) {
                    // Ensure it's a standalone token (not part of a larger identifier).
                    int end = i + tok.length();
                    char before = (i > 0) ? text.charAt(i - 1) : ' ';
                    char after  = (end < n) ? text.charAt(end) : ' ';
                    if (!isIdentChar(before) && !isIdentChar(after)) {
                        return true;
                    }
                }
            }
        }
        return false;
    }

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

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

    /**
     * Attempt a strict (RFC 8259) JSON parse using Gson's JsonReader with lenient=false.
     * In strict mode the reader rejects bare NaN / Infinity / -Infinity tokens,
     * which are not permitted by RFC 8259. Quoted string values ("NaN", etc.) are accepted.
     * We use JsonReader.skipValue() directly because Gson.fromJson() internally resets the
     * reader to lenient mode, which would mask the strict-mode rejection.
     */
    static boolean tryStrictJsonParse(String json) {
        try {
            JsonReader reader = new JsonReader(new StringReader(json));
            reader.setLenient(false);
            reader.skipValue();
            return true;
        } catch (Throwable t) {
            return false;
        }
    }
}