flamenet-messenger-ios/Packages/FlamenetFiles/Tests/FilesEngineTests/InteropTests.swift

import CryptoKit
import Foundation
import XCTest

@testable import FilesEngine

/// Interop against the published reference implementation.
final class InteropTests: XCTestCase {
    private struct Vectors: Decodable {
        let version: Int
        let cases: Cases

        struct Cases: Decodable {
            let labelBytes: [LabelCase]
            let kdf: KDFCase
            let hierarchy: Hierarchy
            let nameDigest: [DigestCase]
            let recovery: Recovery
        }
        struct LabelCase: Decodable {
            let classification: Int16
            let compartments: [Int]
            let expected: String
        }
        struct KDFCase: Decodable {
            let passphrase: String
            let salt: String
            let iterations: Int
            let authSecret: String
        }
        struct DigestCase: Decodable {
            let folderKey: String
            let name: String
            let expected: String
        }
        struct Recovery: Decodable {
            let code: String
            let entropy: String
        }
        struct Hierarchy: Decodable {
            let label: LabelSpec
            let identityPublicKey: String
            let identityRecord: Identity.Record
            let labelKey: String
            let nodeKey: String
            let contentKey: String
            let labelKeyGrant: String
            let wrappedNodeKey: String
            let wrappedContentKey: String
            let nameDigest: String
            let nameCiphertext: String
            let expectedName: String
            let content: String
            let expectedContent: String
            let contentChunkSize: Int
            struct LabelSpec: Decodable {
                let classification: Int16
                let compartments: [Int]
            }
        }
    }

    private var vectors: Vectors!

    override func setUpWithError() throws {
        let url = try XCTUnwrap(Bundle.module.url(forResource: "vectors/v1", withExtension: "json"))
        vectors = try JSONDecoder().decode(Vectors.self, from: try Data(contentsOf: url))
    }

    private func b64(_ s: String) throws -> Data { try XCTUnwrap(Data(base64Encoded: s)) }
    private var hierarchyLabel: Label {
        Label(
            classification: vectors.cases.hierarchy.label.classification,
            compartments: Set(vectors.cases.hierarchy.label.compartments))
    }

    // MARK: §1

    func testLabelEncodingMatchesTheReference() throws {
        for entry in vectors.cases.labelBytes {
            let label = Label(
                classification: entry.classification, compartments: Set(entry.compartments))
            XCTAssertEqual(
                label.bytes, try b64(entry.expected),
                "classification \(entry.classification) compartments \(entry.compartments)")
        }
    }

    // MARK: §2

    func testAuthSecretDerivationMatches() throws {
        let c = vectors.cases.kdf
        let derived = try KDF.authSecret(
            passphrase: c.passphrase, salt: try b64(c.salt),
            parameters: .init(id: 1, iterations: c.iterations)
        )
        XCTAssertEqual(derived, try b64(c.authSecret))
    }

    // MARK: §6

    func testNameDigestsMatch() throws {
        for entry in vectors.cases.nameDigest {
            XCTAssertEqual(
                Names.digest(folderKey: try b64(entry.folderKey), name: entry.name),
                try b64(entry.expected), entry.name)
        }
    }

    // MARK: §3-§7, the whole chain

    func testUnlocksAnIdentityWrittenByTheBrowser() throws {
        let h = vectors.cases.hierarchy
        let identity = try Identity.unlock(
            h.identityRecord, passphrase: vectors.cases.kdf.passphrase)
        // The public half must match what the browser published for it.
        XCTAssertEqual(identity.publicKey.rawRepresentation, try b64(h.identityPublicKey))
    }

    func testAcceptsALabelKeyGrantFromTheBrowser() throws {
        let h = vectors.cases.hierarchy
        let identity = try Identity.unlock(
            h.identityRecord, passphrase: vectors.cases.kdf.passphrase)
        let labelKey = try Keys.acceptLabelKey(
            try b64(h.labelKeyGrant), label: hierarchyLabel,
            privateKey: identity)
        XCTAssertEqual(labelKey, try b64(h.labelKey))
    }

    func testUnwrapsTheKeyHierarchy() throws {
        let h = vectors.cases.hierarchy
        let nodeKey = try Keys.unwrapNodeKey(
            try b64(h.wrappedNodeKey),
            labelKey: try b64(h.labelKey), label: hierarchyLabel)
        XCTAssertEqual(nodeKey, try b64(h.nodeKey))

        let contentKey = try Keys.unwrapContentKey(
            try b64(h.wrappedContentKey),
            nodeKey: nodeKey, label: hierarchyLabel)
        XCTAssertEqual(contentKey, try b64(h.contentKey))
    }

    func testDecryptsAFilenameFromTheBrowser() throws {
        let h = vectors.cases.hierarchy
        let name = try Names.decrypt(
            folderKey: try b64(h.nodeKey), blob: try b64(h.nameCiphertext),
            digest: try b64(h.nameDigest), label: hierarchyLabel)
        XCTAssertEqual(name, h.expectedName)
    }

    func testDecryptsContentFromTheBrowser() throws {
        let h = vectors.cases.hierarchy
        let plaintext = try Content.decrypt(
            try b64(h.content), contentKey: try b64(h.contentKey),
            label: hierarchyLabel)
        XCTAssertEqual(plaintext, try b64(h.expectedContent))
    }

    /// The label is authenticated everywhere, so a relabelled row must fail.
    func testContentDoesNotOpenUnderADifferentLabel() throws {
        let h = vectors.cases.hierarchy
        let wrong = Label(classification: hierarchyLabel.classification, compartments: [])
        XCTAssertThrowsError(
            try Content.decrypt(
                try b64(h.content),
                contentKey: try b64(h.contentKey), label: wrong))
    }
}

/// Recovery codes and identity protection.
final class RecoveryTests: XCTestCase {
    private struct RecoveryVectors: Decodable {
        struct Case: Decodable {
            let entropy: String
            let code: String
            let kdfSecret: String
        }
        let recovery: [Case]
        let recoveryWithQ: Case
    }

    private var vectors: RecoveryVectors!

    override func setUpWithError() throws {
        let url = try XCTUnwrap(Bundle.module.url(forResource: "vectors/v1", withExtension: "json"))
        vectors = try JSONDecoder().decode(RecoveryVectors.self, from: try Data(contentsOf: url))
    }

    /// The one that matters: codes written by the browser must decode here, and
    /// this engine must produce the identical string from the same entropy.
    func testAgreesWithTheBrowser() throws {
        for expected in vectors.recovery + [vectors.recoveryWithQ] {
            let entropy = try XCTUnwrap(Data(base64Encoded: expected.entropy))
            XCTAssertEqual(try Recovery.encode(entropy), expected.code)
            XCTAssertEqual(try Recovery.decode(expected.code), entropy)
            XCTAssertEqual(Recovery.kdfSecret(entropy), expected.kdfSecret)
        }
    }

    /// A code containing Q, decoded from paper: lower-cased, spaces instead of
    /// dashes.
    func testTheQCaseSurvivesTranscription() throws {
        let expected = vectors.recoveryWithQ
        XCTAssertTrue(expected.code.contains("Q"), "vector should contain a Q")
        let mangled = expected.code.lowercased().replacingOccurrences(of: "-", with: " ")
        XCTAssertEqual(
            try Recovery.decode(mangled),
            try XCTUnwrap(Data(base64Encoded: expected.entropy)))
    }

    func testCodeRoundTrips() throws {
        for _ in 0..<200 {
            let (entropy, code) = Recovery.generate()
            XCTAssertEqual(code.replacingOccurrences(of: "-", with: "").count, 28)
            XCTAssertEqual(try Recovery.decode(code), entropy)
        }
    }

    /// The failure that shipped once: `Q` is a valid Crockford character, and
    /// folding it to `0` corrupted roughly a third of all codes.
    func testQIsNotFolded() throws {
        XCTAssertEqual(Recovery.normalize("Q"), "Q")
        XCTAssertEqual(Recovery.normalize("O"), "0")
        XCTAssertEqual(Recovery.normalize("I"), "1")
        XCTAssertEqual(Recovery.normalize("L"), "1")
        // A code containing Q must still decode after a round trip through
        // normalisation, which is what the fold used to break.
        var found = false
        for _ in 0..<500 where !found {
            let (entropy, code) = Recovery.generate()
            if code.contains("Q") {
                found = true
                XCTAssertEqual(try Recovery.decode(code.lowercased()), entropy)
            }
        }
        XCTAssertTrue(found, "expected at least one code containing Q in 500 draws")
    }

    func testFormattingDoesNotChangeTheKey() throws {
        let (entropy, code) = Recovery.generate()
        let mangled = code.lowercased().replacingOccurrences(of: "-", with: " ")
        XCTAssertEqual(try Recovery.decode(mangled), entropy)
    }

    func testChecksumCatchesATypo() throws {
        let (_, code) = Recovery.generate()
        var characters = Array(Recovery.normalize(code))
        // Change one character to a different valid one.
        characters[3] = characters[3] == "7" ? "8" : "7"
        XCTAssertThrowsError(try Recovery.decode(String(characters)))
    }

    func testIdentityRecordRoundTrips() throws {
        // A low iteration count: this asserts the format, not the work factor,
        // and 650,000 rounds per case would make the suite unbearable.
        let fast = KDF.Parameters(id: 1, iterations: 1_000)
        let key = Identity.generate()
        let record = try Identity.protect(key, passphrase: "correct horse", parameters: fast)
        let opened = try Identity.unlock(record, passphrase: "correct horse")
        XCTAssertEqual(opened.rawRepresentation, key.rawRepresentation)
    }

    func testTheWrongPassphraseDoesNotOpenIt() throws {
        let fast = KDF.Parameters(id: 1, iterations: 1_000)
        let record = try Identity.protect(Identity.generate(), passphrase: "one", parameters: fast)
        XCTAssertThrowsError(try Identity.unlock(record, passphrase: "two"))
    }

    /// The passphrase and the recovery code wrap the same key independently, so
    /// changing one must not disturb the other.
    func testRecoveryOpensTheSameKey() throws {
        let fast = KDF.Parameters(id: 1, iterations: 1_000)
        let key = Identity.generate()
        let (entropy, _) = Recovery.generate()
        let byRecovery = try Identity.protectWithRecovery(key, entropy: entropy, parameters: fast)
        let opened = try Identity.unlockWithRecovery(byRecovery, entropy: entropy)
        XCTAssertEqual(opened.rawRepresentation, key.rawRepresentation)
    }

    /// Different contexts, so a record made for one secret cannot be opened as
    /// the other's even when the derived material would otherwise line up.
    func testRecoveryRecordIsNotAPassphraseRecord() throws {
        let fast = KDF.Parameters(id: 1, iterations: 1_000)
        let key = Identity.generate()
        let (entropy, _) = Recovery.generate()
        let byRecovery = try Identity.protectWithRecovery(key, entropy: entropy, parameters: fast)
        XCTAssertEqual(byRecovery.context, KDF.Context.filesRecovery)
        // Same secret material, wrong context: must not open.
        XCTAssertThrowsError(
            try Identity.unlock(
                Identity.Record(
                    v: 1, kdf: byRecovery.kdf, iterations: byRecovery.iterations,
                    context: KDF.Context.filesIdentity,
                    salt: byRecovery.salt, nonce: byRecovery.nonce,
                    privateKey: byRecovery.privateKey),
                passphrase: Recovery.kdfSecret(entropy)))
    }
}