package unit;
import java.util.*;

/**
 * OpensmtpdTest — CWE-407 benchmark for opensmtpd-0001
 *
 * Models ruleset_match() in usr.sbin/smtpd/ruleset.c:
 *   SLOW: TAILQ_FOREACH over R rules × M recipient envelopes = O(R×M)
 *   FAST: HashMap dispatch from dest domain → candidate rule = O(M)
 *
 * Run: javac -d . OpensmtpdTest.java && java -ea unit.OpensmtpdTest
 */
public class OpensmtpdTest {

    // ── Rule model ────────────────────────────────────────────────────────────

    static class Rule {
        final int id;
        final String forDomain;   // null = match all (no flag_for)
        final String fromDomain;  // null = match all
        Rule(int id, String forDomain, String fromDomain) {
            this.id = id;
            this.forDomain = forDomain;
            this.fromDomain = fromDomain;
        }
        boolean matches(String destDomain, String srcDomain) {
            if (forDomain != null && !forDomain.equalsIgnoreCase(destDomain)) return false;
            if (fromDomain != null && !fromDomain.equalsIgnoreCase(srcDomain)) return false;
            return true;
        }
    }

    static class Envelope {
        final String destDomain;
        final String srcDomain;
        Envelope(String destDomain, String srcDomain) {
            this.destDomain = destDomain;
            this.srcDomain = srcDomain;
        }
    }

    // ── Build test data ───────────────────────────────────────────────────────

    /**
     * Build R rules: all have specific forDomain (no catch-all).
     * This models a large ISP config where each hosted domain has its own
     * accept rule. The matching rule for domain "dest-K" is at index K in
     * the TAILQ, forcing O(K) scan to reach it.
     */
    static List<Rule> buildRules(int totalRules) {
        List<Rule> rules = new ArrayList<>(totalRules);
        for (int i = 0; i < totalRules; i++)
            rules.add(new Rule(i, "dest-" + i + ".example.com", null));
        return rules;
    }

    // ── SLOW: TAILQ_FOREACH — O(R) per envelope ───────────────────────────────

    /**
     * Models ruleset_match(): scan all rules from head until a match.
     * Returns total number of rule.matches() evaluations (each = O(1) here
     * but represents one full pass through the rule's sub-matchers in C).
     */
    static long rulesetMatchSlow(List<Rule> rules, List<Envelope> envelopes) {
        long ops = 0;
        for (Envelope evp : envelopes) {
            for (Rule r : rules) {
                ops++;  // one rule evaluation
                if (r.matches(evp.destDomain, evp.srcDomain)) break;
            }
        }
        return ops;
    }

    // ── FAST: HashMap dispatch — O(1) candidate lookup per envelope ───────────

    /**
     * Models the CWE-407 fix: build a domain→rule index at startup,
     * then do O(1) dict_get(destDomain) per envelope.
     *
     * Match-all rules are always appended to a fallback list and checked
     * only if the domain-specific candidate does not match.
     */
    static long rulesetMatchFast(List<Rule> rules, List<Envelope> envelopes) {
        // Build index: destDomain → first matching rule (simplified)
        Map<String, Rule> domainIndex = new HashMap<>();
        List<Rule> catchAll = new ArrayList<>();
        for (Rule r : rules) {
            if (r.forDomain != null) {
                domainIndex.putIfAbsent(r.forDomain, r);
            } else {
                catchAll.add(r);
            }
        }

        long ops = 0;
        for (Envelope evp : envelopes) {
            ops++;  // O(1) hash lookup
            Rule candidate = domainIndex.get(evp.destDomain);
            if (candidate != null && candidate.matches(evp.destDomain, evp.srcDomain)) {
                // matched — done
            } else {
                // fallback to catch-all rules (typically 1-2)
                for (Rule r : catchAll) {
                    ops++;
                    if (r.matches(evp.destDomain, evp.srcDomain)) break;
                }
            }
        }
        return ops;
    }

    // ── bench harness ─────────────────────────────────────────────────────────

    static void bench(String label, Runnable slow, Runnable fast, long sOps, long fOps) {
        slow.run(); fast.run();
        long t0 = System.nanoTime(); slow.run(); long sMs = (System.nanoTime() - t0) / 1_000_000;
        long t1 = System.nanoTime(); fast.run(); long fMs = (System.nanoTime() - t1) / 1_000_000;
        double r = fOps > 0 ? (double) sOps / fOps : 0;
        System.out.printf("  %-56s  slow:%4dms (%,d ops)  fast:%4dms (%,d ops)  speedup:%.0fx%n",
                label, sMs, sOps, fMs, fOps, r);
    }

    // ── main ──────────────────────────────────────────────────────────────────

    public static void main(String[] args) {
        final int R = 300;   // policy rules in smtpd.conf
        final int M = 2000;  // recipient envelopes per message burst

        List<Rule> rules = buildRules(R);

        // Worst-case scenario: envelopes are spread across all R domains,
        // so the average rule scan depth is R/2. With R=300 rules and M=2000
        // envelopes, the TAILQ scan does O(R/2 × M) = 300,000 evaluations.
        // The hash dispatch does O(1) per envelope = M evaluations.
        List<Envelope> envelopes = new ArrayList<>(M);
        for (int i = 0; i < M; i++)
            envelopes.add(new Envelope(
                    "dest-" + (i % R) + ".example.com",  // uniformly distributed across all rules
                    "sender.example.com"));

        System.out.println("OpensmtpdTest — CWE-407");
        System.out.println();
        System.out.println("opensmtpd-0001: ruleset_match() TAILQ scan");

        final long[] sOps = new long[1], fOps = new long[1];
        bench(String.format("ruleset_match  R=%d rules, M=%d envelopes", R, M),
                () -> { sOps[0] = rulesetMatchSlow(rules, envelopes); },
                () -> { fOps[0] = rulesetMatchFast(rules, envelopes); },
                rulesetMatchSlow(rules, envelopes),
                rulesetMatchFast(rules, envelopes));

        assert sOps[0] > fOps[0] * 10 :
                "opensmtpd-0001: expected >10x more ops slow vs fast, got slow="
                + sOps[0] + " fast=" + fOps[0];

        System.out.println();
        System.out.println("All assertions passed.");
    }
}