import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * UnboundRdataDuplicateTest — models unbound authzone rdata_duplicate() CWE-407.
 *
 * unbound services/authzone.c rdata_duplicate() scans all existing RRs in a
 * packed_rrset linearly for each new RR added during zone load.  For R records
 * per rrset this is O(R^2) total insertion cost.
 *
 * Affected path: apply_axfr -> az_insert_rr_decompress -> az_domain_add_rr ->
 * rdata_duplicate (called per-RR, up to 3x for RRSIGs).
 *
 * Fix: maintain an auxiliary HashMap keyed on rdata bytes alongside the rrset
 * during zone load for O(1) duplicate checks.
 *
 * Run: javac UnboundRdataDuplicateTest.java && java UnboundRdataDuplicateTest
 */
public class UnboundRdataDuplicateTest {

    /** Wrap byte[] for use as HashMap key */
    static final class RdataKey {
        final byte[] data;
        final int hash;
        RdataKey(byte[] d) {
            this.data = d;
            this.hash = Arrays.hashCode(d);
        }
        @Override public int hashCode() { return hash; }
        @Override public boolean equals(Object o) {
            return o instanceof RdataKey && Arrays.equals(data, ((RdataKey)o).data);
        }
    }

    /** O(R^2): ArrayList linear scan per insertion — models rdata_duplicate scan */
    static int buildRrsetDefect(List<byte[]> rrs, int count) {
        List<byte[]> rrset = new ArrayList<>();
        for (int i = 0; i < count; i++) {
            byte[] rdata = rrs.get(i);
            // CWE-407: O(R) scan each time, O(R^2) total
            boolean dup = false;
            for (byte[] existing : rrset) {
                if (Arrays.equals(existing, rdata)) {
                    dup = true;
                    break;
                }
            }
            if (!dup) {
                rrset.add(rdata);
            }
        }
        return rrset.size();
    }

    /** O(R): HashMap for O(1) membership — models dedup_set fix */
    static int buildRrsetFixed(List<byte[]> rrs, int count) {
        List<byte[]> rrset = new ArrayList<>();
        Map<RdataKey, Boolean> dedup = new HashMap<>();
        for (int i = 0; i < count; i++) {
            byte[] rdata = rrs.get(i);
            if (dedup.put(new RdataKey(rdata), Boolean.TRUE) == null) { // O(1)
                rrset.add(rdata);
            }
        }
        return rrset.size();
    }

    public static void main(String[] args) {
        // Model: large NSEC3 or RRSIG rrset (R unique records per zone node)
        int r = 1000;
        List<byte[]> rrs = new ArrayList<>(r);
        for (int i = 0; i < r; i++) {
            // Simulate 20-byte NSEC3 hash + 4-byte bitmap = 24 bytes rdata
            byte[] rdata = new byte[24];
            rdata[0] = (byte)(i >> 8);
            rdata[1] = (byte)(i & 0xff);
            for (int j = 2; j < 24; j++) rdata[j] = (byte)((i * 31 + j) & 0xff);
            rrs.add(rdata);
        }

        // Warm up JIT
        buildRrsetDefect(rrs, 50);
        buildRrsetFixed(rrs, 50);

        int reps = 5;

        long t0 = System.nanoTime();
        int szDefect = 0;
        for (int rep = 0; rep < reps; rep++) {
            szDefect = buildRrsetDefect(rrs, r);
        }
        long defectNs = (System.nanoTime() - t0) / reps;

        long t1 = System.nanoTime();
        int szFixed = 0;
        for (int rep = 0; rep < reps; rep++) {
            szFixed = buildRrsetFixed(rrs, r);
        }
        long fixedNs = (System.nanoTime() - t1) / reps;

        System.out.printf("R=%d RRs per rrset (NSEC3/RRSIG scenario)%n", r);
        System.out.printf("  defect (O(R^2) array scan): %7.2f ms  [size=%d]%n",
                defectNs / 1e6, szDefect);
        System.out.printf("  fixed  (O(R)   HashMap):    %7.2f ms  [size=%d]%n",
                fixedNs / 1e6, szFixed);
        double ratio = (double) defectNs / fixedNs;
        System.out.printf("  ratio: %.1fx%n", ratio);

        if (szDefect != szFixed) {
            throw new AssertionError("size mismatch: " + szDefect + " vs " + szFixed);
        }
        if (szDefect != r) {
            throw new AssertionError("expected " + r + " entries, got " + szDefect);
        }
        if (ratio < 5.0) {
            System.out.println("WARN: ratio lower than expected at this scale");
        } else {
            System.out.println("PASS");
        }
    }
}