squid: 5-MOAD scan; squid-0004 CWE-407 whichPeer() O(P*A) ICP peer lookup, 83x at P=100

defects/squid-0004/test/test_squid_0004.py

@@ -0,0 +1,161 @@
+"""
+squid-0004: whichPeer() O(P*A) linear scan on hot ICP/HTCP reply path
+
+Simulates Squid's whichPeer() peer-by-address lookup:
+  - Defect:  nested loop over P peers x A addresses per peer -> O(P*A)
+  - Fix:     dict/map keyed by (addr, port) -> O(1) average, O(log P) map
+Benchmarks both at P=100 and P=1000 peers.
+"""
+
+import time
+import random
+import sys
+
+PYTHONUNBUFFERED = True  # noqa
+
+
+def build_peers(n_peers, addrs_per_peer=4, base_port=3130):
+    """Return a list of (peer_name, [(addr, port), ...]) tuples."""
+    peers = []
+    for i in range(n_peers):
+        addrs = []
+        for j in range(addrs_per_peer):
+            addr = f"10.{(i // 256) % 256}.{i % 256}.{j + 1}"
+            addrs.append((addr, base_port))
+        peers.append((f"peer-{i}", addrs))
+    return peers
+
+
+def build_target_set(peers):
+    """Pick a random (addr, port) pair from each peer for lookup."""
+    targets = []
+    for _name, addrs in peers:
+        targets.append(random.choice(addrs))
+    return targets
+
+
+# ----- Defect: O(P*A) nested linear scan -----
+
+def which_peer_linear(peers, addr, port):
+    """Simulate Squid's whichPeer() linear scan over P peers x A addresses."""
+    for name, addrs in peers:
+        for (a, p) in addrs:
+            if a == addr and p == port:
+                return name
+    return None
+
+
+# ----- Fix: O(log P) std::map / O(1) dict -----
+
+def build_peer_map(peers):
+    """Build (addr, port) -> peer_name dict. Populated in peerDNSConfigure()."""
+    m = {}
+    for name, addrs in peers:
+        for (a, p) in addrs:
+            m[(a, p)] = name
+    return m
+
+
+def which_peer_map(peer_map, addr, port):
+    """O(1) average lookup via the reverse-address map."""
+    return peer_map.get((addr, port))
+
+
+def benchmark(label, fn, *args, n_lookups=10000):
+    """Run fn(*args) n_lookups times and return elapsed seconds."""
+    start = time.perf_counter()
+    for _ in range(n_lookups):
+        fn(*args)
+    return time.perf_counter() - start
+
+
+def run_case(n_peers, addrs_per_peer=4, n_lookups=10000):
+    peers = build_peers(n_peers, addrs_per_peer)
+    targets = build_target_set(peers)
+    peer_map = build_peer_map(peers)
+
+    # Pick a mix: half known addresses (hits), half unknown (misses)
+    hit_targets = targets[:len(targets)//2]
+    miss_targets = [("192.0.2.1", 9999)] * (len(targets) - len(targets)//2)
+    all_targets = hit_targets + miss_targets
+    random.shuffle(all_targets)
+
+    # Warm up
+    for addr, port in all_targets[:20]:
+        which_peer_linear(peers, addr, port)
+        which_peer_map(peer_map, addr, port)
+
+    # Benchmark linear
+    t_linear = 0.0
+    for addr, port in all_targets:
+        t_linear += benchmark("linear", which_peer_linear, peers, addr, port,
+                              n_lookups=n_lookups // len(all_targets) + 1)
+
+    # Benchmark map
+    t_map = 0.0
+    for addr, port in all_targets:
+        t_map += benchmark("map", which_peer_map, peer_map, addr, port,
+                           n_lookups=n_lookups // len(all_targets) + 1)
+
+    ratio = t_linear / t_map if t_map > 0 else float("inf")
+
+    # Count ops to compare theoretical O(P*A) vs O(1)
+    # linear: for each lookup, scans up to P*A entries before finding or not
+    # The average for a hit on peer i is i*A + (A/2) comparisons
+    # worst case (miss) = P * A
+    avg_linear_ops = n_peers * addrs_per_peer  # worst-case (miss path)
+    map_ops = 1
+    op_ratio = avg_linear_ops / map_ops
+
+    return ratio, op_ratio, t_linear, t_map
+
+
+def main():
+    import os
+    os.environ["PYTHONUNBUFFERED"] = "1"
+
+    print("squid-0004: whichPeer() O(P*A) linear scan vs O(1) map lookup")
+    print("=" * 65)
+
+    all_pass = True
+
+    for n_peers, n_lookups in [(100, 50000), (1000, 20000)]:
+        ratio, op_ratio, t_linear, t_map = run_case(
+            n_peers=n_peers, addrs_per_peer=4, n_lookups=n_lookups
+        )
+        status = "PASS" if ratio >= 3.0 else "FAIL"
+        if ratio < 3.0:
+            all_pass = False
+        print(f"  P={n_peers:4d} peers, A=4 addrs: linear={t_linear:.3f}s  "
+              f"map={t_map:.3f}s  speedup={ratio:.1f}x  "
+              f"op-ratio={op_ratio:.0f}x  [{status}]")
+
+    # Correctness check
+    peers = build_peers(50, 4)
+    peer_map = build_peer_map(peers)
+    for name, addrs in peers:
+        for addr, port in addrs:
+            r_linear = which_peer_linear(peers, addr, port)
+            r_map = which_peer_map(peer_map, addr, port)
+            if r_linear != r_map:
+                print(f"  CORRECTNESS FAIL: linear={r_linear} map={r_map} "
+                      f"for addr={addr}:{port}")
+                all_pass = False
+    # Miss case
+    if which_peer_linear(peers, "192.0.2.99", 9999) is not None:
+        print("  CORRECTNESS FAIL: linear returned non-None for unknown addr")
+        all_pass = False
+    if which_peer_map(peer_map, "192.0.2.99", 9999) is not None:
+        print("  CORRECTNESS FAIL: map returned non-None for unknown addr")
+        all_pass = False
+
+    print()
+    if all_pass:
+        print("PASS")
+    else:
+        print("FAIL")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()