open-standards/wai/wai-sink-swift/Tests/WaiSinkTests/ModelPolicyTests.swift

// A sink that holds a model policy (native-sink §14), through the same C interface every native
// integration links: the corpus live-conformance/model-sink/, staged by scripts/build-xcframework.sh
// with the files its cases read. Every decode case and every object of every frame case decoded by a
// decoder that holds its configuration, at its instant, held to the outcome the corpus states (a
// decoded case gives the bytes of a decoder with no policy; a refused case gives no bytes, no unit
// and the policy's code); every configuration case refused with its code; the run reports the corpus
// states signed with the device key, each holding exactly the decodes of the cases it names; and the
// bundles of the two periods verified as `wai model verify` verifies them. The kiosks (`roster`), whose
// class names them by a roster (model-governance §4.5): every configuration case refused with its
// code, each kiosk's decodes their outcomes, each report carrying the membership proof its
// configuration gives for the period's policy, and both periods' bundles accepted.
import CryptoKit
import XCTest
import WaiSink

final class ModelPolicyTests: XCTestCase {
    /// The codes the policy's check refuses with (native-sink §14); any other is the decode's own.
    static let policyCodes: Set<String> = ["model_not_in_policy", "model_revoked", "policy_expired", "policy_not_yet_valid",
                                           "model_tier_below_policy", "device_not_in_class", "model_runs_full", "record_overlap"]

    func dir() -> URL {
        #if SWIFT_PACKAGE
        return Bundle.module.resourceURL!.appendingPathComponent("Resources/model")
        #else
        return Bundle(for: TestBundleMarker.self).resourceURL!.appendingPathComponent("model")
        #endif
    }

    /// A file the corpus names by its path relative to wai/, as staged.
    func file(_ rel: String) throws -> Data {
        try Data(contentsOf: dir().appendingPathComponent(rel.replacingOccurrences(of: "/", with: "_")))
    }

    func json(_ v: Any) throws -> Data {
        try JSONSerialization.data(withJSONObject: v, options: [.fragmentsAllowed])
    }

    func sha256(_ d: Data) -> String {
        SHA256.hash(data: d).map { String(format: "%02x", $0) }.joined()
    }

    func seed(_ hex: String) -> Data {
        var out = Data()
        var i = hex.startIndex
        while i < hex.endIndex {
            let j = hex.index(i, offsetBy: 2)
            out.append(UInt8(hex[i..<j], radix: 16)!)
            i = j
        }
        return out
    }

    func check(_ name: String, _ o: [String: Any], _ got: WaiIntDecoded, _ want: WaiIntDecoded, _ sha: String?) {
        if let code = o["refused"] as? String {
            XCTAssertEqual(got.refused, code, "\(name): the code")
            XCTAssertTrue(got.bytes.isEmpty, "\(name): a refused decode gave bytes")
            XCTAssertNotEqual(got.model?["kept"] as? Bool, true, "\(name): a refused decode's run was kept")
            if Self.policyCodes.contains(code) {
                XCTAssertNil(got.unit, "\(name): a decode the policy refused ran on a unit")
                XCTAssertTrue(got.dispatch.isEmpty, "\(name): a decode the policy refused was dispatched")
            }
            return
        }
        XCTAssertNil(got.refused, "\(name): refused")
        XCTAssertFalse(got.bytes.isEmpty, name)
        XCTAssertEqual(got.bytes, want.bytes, "\(name): not the bytes of a decoder with no policy")
        if o["kept"] as? Bool == true {
            XCTAssertEqual(got.model?["kept"] as? Bool, true, "\(name): kept")
            XCTAssertEqual((got.model?["set"] as? [String: Any])?["sha256"] as? String, sha, "\(name): the set its run names")
        } else {
            XCTAssertTrue(got.model == nil || got.model?["kept"] as? Bool == false, "\(name): a run kept")
        }
    }

    func testEveryCaseOfTheModelPolicyCorpusGetsItsOutcome() throws {
        let doc = try JSONSerialization.jsonObject(with: try file("cases.json")) as! [String: Any]
        let sets = Dictionary(uniqueKeysWithValues: (doc["sets"] as! [[String: Any]]).map { ($0["name"] as! String, $0) })
        func files(_ name: Any?) throws -> [String: Data] {
            guard let n = name as? String else { return [:] }
            var out: [String: Data] = [:]
            for (m, p) in sets[n]!["files"] as! [String: String] { out[m] = try file(p) }
            return out
        }
        func sha(_ name: Any?) -> String? { (name as? String).flatMap { sets[$0]?["sha256"] as? String } }
        func decode(_ d: WaiIntDecoder, _ dec: [String: Any], _ set: [String: Data]) throws -> WaiIntDecoded {
            if let e = dec["envelope"] as? String { return try d.decode(envelope: try file(e), set: set) }
            return try d.decode(capability: dec["capability"] as! String, payload: try file(dec["payload"] as! String), set: set)
        }
        let configs = doc["configs"] as! [String: Any]

        // Every configuration case is refused, with its code.
        let refuser = try WaiIntDecoder(gpu: .off)
        for c in doc["config_cases"] as! [[String: Any]] {
            XCTAssertThrowsError(try refuser.holdPolicy(try json(c["config"]!)), c["name"] as! String) { e in
                let code = (c["outcome"] as! [String: Any])["refused"] as! String
                XCTAssertTrue("\(e)".contains("\(code): "), "\(c["name"]!): \(e)")
            }
        }

        let plain = try WaiIntDecoder(gpu: .off)
        var decoders: [String: WaiIntDecoder] = [:]
        var decoded = 0
        var refused = 0
        for config in ["set-top", "signage"] {
            let d = try WaiIntDecoder(gpu: .auto)
            try d.holdPolicy(try json(configs[config]!))
            for c in (doc["decode_cases"] as! [[String: Any]]) where c["config"] as? String == config {
                let dec = c["decode"] as! [String: Any]
                let set = try files(dec["set"])
                try d.setClock(c["at"] as? String)
                let o = c["outcome"] as! [String: Any]
                check(c["name"] as! String, o, try decode(d, dec, set), try decode(plain, dec, set), sha(dec["set"]))
                if o["refused"] != nil { refused += 1 } else { decoded += 1 }
            }
            for c in (doc["frame_cases"] as! [[String: Any]]) where c["config"] as? String == config {
                let f = c["frames"] as! [String: Any]
                let set = try files(f["set"])
                // The stream's objects pin the set in the form `wai.prior.set`: the host holds its listing.
                let listing = Data(set.keys.sorted().map { "\(sha256(set[$0]!))  \($0)\n" }.joined().utf8)
                try d.setClock(c["at"] as? String)
                let s = try WaiIntFrameStream(capability: f["capability"] as! String, set: set, listing: listing)
                let p = try WaiIntFrameStream(capability: f["capability"] as! String, set: set, listing: listing)
                for (i, o) in (f["objects"] as! [String]).enumerated() {
                    let e = try file(o)
                    let got = try d.decode(frame: e, stream: s)
                    check("\(c["name"]!) object \(i)", c["outcome"] as! [String: Any], got, try plain.decode(frame: e, stream: p), sha(f["set"]))
                }
            }
            decoders[config] = d
        }

        // The reports, in the corpus's order, each holding exactly the decodes of the cases it names.
        var reports: [Any] = []
        let seeds = doc["seeds"] as! [String: String]
        for r in doc["reports"] as! [[String: Any]] {
            let config = r["config"] as! String
            let to = r["to"] as! String
            let text = try decoders[config]!.runReport(from: r["from"] as! String, to: to, issuedAt: to, deviceSeed: seed(seeds[config]!))
            let rep = try JSONSerialization.jsonObject(with: text) as! [String: Any]
            var want: [[String]] = []
            for n in r["decodes"] as! [String] {
                if let c = (doc["decode_cases"] as! [[String: Any]]).first(where: { $0["name"] as? String == n }) {
                    let dec = c["decode"] as! [String: Any]
                    want.append([c["at"] as! String, sha(dec["set"])!, dec["capability"] as! String])
                } else {
                    let c = (doc["frame_cases"] as! [[String: Any]]).first(where: { $0["name"] as? String == n })!
                    let f = c["frames"] as! [String: Any]
                    for _ in f["objects"] as! [String] { want.append([c["at"] as! String, sha(f["set"])!, f["capability"] as! String]) }
                }
            }
            let got: [[String]] = (rep["decodes"] as! [[String: Any]]).map {
                [$0["at"] as! String, ($0["set"] as! [String: Any])["sha256"] as! String, ($0["dispatch"] as! [String: Any])["capability"] as! String]
            }
            XCTAssertEqual(got, want, "\(config) \(r["from"]!): the decodes the report states")
            reports.append(rep)
        }
        XCTAssertThrowsError(try decoders["set-top"]!.runReport(from: "2026-11-14T00:00:00Z", to: "2026-11-20T00:00:00Z", deviceSeed: seed(seeds["set-top"]!))) { e in
            XCTAssertTrue("\(e)".contains("record_overlap: "), "\(e)")
        }
        // The bundles of the two periods, as `wai model verify` verifies them.
        for rs in [[reports[0], reports[1]], [reports[2]]] {
            let bundle: [String: Any] = ["kind": "wai-model-bundle", "policies": doc["policies"]!, "key_log": doc["key_log"]!,
                                         "attestations": [doc["attestation"]!], "presentations": [Any](), "reports": rs, "rollup": NSNull()]
            let o = try WaiIntDecoder.verifyModelBundle(try json(bundle))
            XCTAssertEqual(o["outcome"] as? String, "accept", "\(o)")
        }
        print("model policy: \(decoded) decodes and \(refused) refusals as the corpus states, \(reports.count) run reports, both bundles accepted")
    }

    func testKiosksNamedByARosterCarryTheirProofs() throws {
        let doc = try JSONSerialization.jsonObject(with: try file("cases.json")) as! [String: Any]
        let k = doc["roster"] as! [String: Any]
        let sets = Dictionary(uniqueKeysWithValues: (doc["sets"] as! [[String: Any]]).map { ($0["name"] as! String, $0) })
        func files(_ name: String) throws -> [String: Data] {
            var out: [String: Data] = [:]
            for (m, p) in sets[name]!["files"] as! [String: String] { out[m] = try file(p) }
            return out
        }
        let configs = k["configs"] as! [String: Any]
        let cases = k["decode_cases"] as! [[String: Any]]

        let refuser = try WaiIntDecoder(gpu: .off)
        for c in k["config_cases"] as! [[String: Any]] {
            XCTAssertThrowsError(try refuser.holdPolicy(try json(c["config"]!)), c["name"] as! String) { e in
                let code = (c["outcome"] as! [String: Any])["refused"] as! String
                XCTAssertTrue("\(e)".contains("\(code): "), "\(c["name"]!): \(e)")
            }
        }

        let plain = try WaiIntDecoder(gpu: .off)
        var decoders: [String: WaiIntDecoder] = [:]
        for config in configs.keys.sorted() {
            let d = try WaiIntDecoder(gpu: .off)
            try d.holdPolicy(try json(configs[config]!))
            for c in cases where c["config"] as? String == config {
                let dec = c["decode"] as! [String: Any]
                let set = try files(dec["set"] as! String)
                let payload = try file(dec["payload"] as! String)
                try d.setClock(c["at"] as? String)
                let got = try d.decode(capability: dec["capability"] as! String, payload: payload, set: set)
                let want = try plain.decode(capability: dec["capability"] as! String, payload: payload, set: set)
                check(c["name"] as! String, c["outcome"] as! [String: Any], got, want, sets[dec["set"] as! String]?["sha256"] as? String)
            }
            decoders[config] = d
        }

        let seeds = k["seeds"] as! [String: String]
        var reports: [[String: Any]] = []
        for r in k["reports"] as! [[String: Any]] {
            let config = r["config"] as! String
            let p = r["policy"] as! Int
            let to = r["to"] as! String
            let text = try decoders[config]!.runReport(from: r["from"] as! String, to: to, issuedAt: to, deviceSeed: seed(seeds[config]!))
            let rep = try JSONSerialization.jsonObject(with: text) as! [String: Any]
            let membership = (configs[config] as! [String: Any])["membership"] as! [Any]
            XCTAssertEqual(rep["membership"] as? NSDictionary, membership[p] as? NSDictionary, "\(config) \(r["from"]!): the proof the report carries")
            XCTAssertNotNil(rep["membership"], "\(config) \(r["from"]!): a kiosk's report carries its proof")
            XCTAssertEqual((rep["decodes"] as! [Any]).count, (r["decodes"] as! [Any]).count, "\(config) \(r["from"]!): the decodes the report states")
            reports.append(rep)
        }
        for rs in [[reports[0], reports[1]], [reports[2]]] {
            let bundle: [String: Any] = ["kind": "wai-model-bundle", "policies": k["policies"]!, "key_log": NSNull(), "attestations": [Any](),
                                         "presentations": [Any](), "reports": rs, "rollup": NSNull()]
            let o = try WaiIntDecoder.verifyModelBundle(try json(bundle))
            XCTAssertEqual(o["outcome"] as? String, "accept", "\(o)")
        }
        print("model policy, kiosks: \(reports.count) run reports carrying their proofs, both bundles accepted")
    }
}