Kyber

fn Kyber(comptime p: Params) type

File

Code

fn Kyber(comptime p: Params) type {
    return struct {
        // Size of a ciphertext, in bytes.
        pub const ciphertext_length = Poly.compressedSize(p.du) * p.k + Poly.compressedSize(p.dv);

        const Self = @This();
        const V = PolyVec(p.k);
        const M = Mat(p.k);

        /// Length (in bytes) of a shared secret.
        pub const shared_length = common_shared_key_size;
        /// Length (in bytes) of a seed for deterministic encapsulation.
        pub const encaps_seed_length = common_encaps_seed_length;
        /// Length (in bytes) of a seed for key generation.
        pub const seed_length: usize = inner_seed_length + shared_length;
        /// Algorithm name.
        pub const name = p.name;

        /// A shared secret, and an encapsulated (encrypted) representation of it.
        pub const EncapsulatedSecret = struct {
            shared_secret: [shared_length]u8,
            ciphertext: [ciphertext_length]u8,
        };

        /// A Kyber public key.
        pub const PublicKey = struct {
            pk: InnerPk,

            // Cached
            hpk: [h_length]u8, // H(pk)

            /// Size of a serialized representation of the key, in bytes.
            pub const encoded_length = InnerPk.encoded_length;

            /// Generates a shared secret, encapsulated for the public key,
            /// using random bytes.
            ///
            /// This is recommended over `encapsDeterministic`.
            pub fn encaps(pk: PublicKey, io: std.Io) EncapsulatedSecret {
                var m: [inner_plaintext_length]u8 = undefined;
                io.random(&m);
                return encapsInner(pk, &m);
            }

            /// Generates a shared secret, encapsulated for the public key,
            /// using the provided seed.
            ///
            /// Calling `encaps` instead is recommended.
            pub fn encapsDeterministic(pk: PublicKey, seed: *const [encaps_seed_length]u8) EncapsulatedSecret {
                var m: [inner_plaintext_length]u8 = undefined;
                if (p.ml_kem) {
                    @memcpy(&m, seed);
                } else {
                    // m = H(seed)
                    sha3.Sha3_256.hash(seed, &m, .{});
                }
                return encapsInner(pk, &m);
            }

            fn encapsInner(pk: PublicKey, m: *[inner_plaintext_length]u8) EncapsulatedSecret {
                // (K', r) = G(m ‖ H(pk))
                var kr: [inner_plaintext_length + h_length]u8 = undefined;
                var g = sha3.Sha3_512.init(.{});
                g.update(m);
                g.update(&pk.hpk);
                g.final(&kr);

                // c = innerEncrypt(pk, m, r)
                const ct = pk.pk.encrypt(m, kr[32..64]);

                if (p.ml_kem) {
                    return EncapsulatedSecret{
                        .shared_secret = kr[0..shared_length].*, // ML-KEM: K = K'
                        .ciphertext = ct,
                    };
                } else {
                    // Compute H(c) and put in second slot of kr, which will be (K', H(c)).
                    sha3.Sha3_256.hash(&ct, kr[32..], .{});

                    var ss: [shared_length]u8 = undefined;
                    sha3.Shake256.hash(&kr, &ss, .{});
                    return EncapsulatedSecret{
                        .shared_secret = ss, // Kyber: K = KDF(K' ‖ H(c))
                        .ciphertext = ct,
                    };
                }
            }

            /// Serializes the key into a byte array.
            pub fn toBytes(pk: PublicKey) [encoded_length]u8 {
                return pk.pk.toBytes();
            }

            /// Deserializes the key from a byte array.
            pub fn fromBytes(buf: *const [encoded_length]u8) errors.NonCanonicalError!PublicKey {
                var ret: PublicKey = undefined;
                ret.pk = try InnerPk.fromBytes(buf[0..InnerPk.encoded_length]);
                sha3.Sha3_256.hash(buf, &ret.hpk, .{});
                return ret;
            }
        };

        /// A Kyber secret key.
        pub const SecretKey = struct {
            sk: InnerSk,
            pk: InnerPk,
            hpk: [h_length]u8, // H(pk)
            z: [shared_length]u8,

            /// Size of a serialized representation of the key, in bytes.
            pub const encoded_length: usize =
                InnerSk.encoded_length + InnerPk.encoded_length + h_length + shared_length;

            /// Decapsulates the shared secret within ct using the private key.
            pub fn decaps(sk: SecretKey, ct: *const [ciphertext_length]u8) ![shared_length]u8 {
                // m' = innerDec(ct)
                const m2 = sk.sk.decrypt(ct);

                // (K'', r') = G(m' ‖ H(pk))
                var kr2: [64]u8 = undefined;
                var g = sha3.Sha3_512.init(.{});
                g.update(&m2);
                g.update(&sk.hpk);
                g.final(&kr2);

                // ct' = innerEnc(pk, m', r')
                const ct2 = sk.pk.encrypt(&m2, kr2[32..64]);

                if (p.ml_kem) {
                    // ML-KEM: K = K'' if ct == ct', else K = J(z || c) per FIPS 203
                    var k_bar: [shared_length]u8 = undefined;
                    var j = sha3.Shake256.init(.{});
                    j.update(&sk.z);
                    j.update(ct);
                    j.squeeze(&k_bar);
                    cmov(shared_length, kr2[0..shared_length], k_bar, ctneq(ciphertext_length, ct.*, ct2));
                    return kr2[0..shared_length].*;
                } else {
                    // Kyber: K = KDF(K''/z ‖ H(c))
                    sha3.Sha3_256.hash(ct, kr2[32..], .{});
                    cmov(32, kr2[0..32], sk.z, ctneq(ciphertext_length, ct.*, ct2));
                    var ss: [shared_length]u8 = undefined;
                    sha3.Shake256.hash(&kr2, &ss, .{});
                    return ss;
                }
            }

            /// Serializes the key into a byte array.
            pub fn toBytes(sk: SecretKey) [encoded_length]u8 {
                return sk.sk.toBytes() ++ sk.pk.toBytes() ++ sk.hpk ++ sk.z;
            }

            /// Deserializes the key from a byte array.
            pub fn fromBytes(buf: *const [encoded_length]u8) errors.NonCanonicalError!SecretKey {
                var ret: SecretKey = undefined;
                comptime var s: usize = 0;
                ret.sk = InnerSk.fromBytes(buf[s .. s + InnerSk.encoded_length]);
                s += InnerSk.encoded_length;
                ret.pk = try InnerPk.fromBytes(buf[s .. s + InnerPk.encoded_length]);
                s += InnerPk.encoded_length;
                ret.hpk = buf[s..][0..h_length].*;
                s += h_length;
                ret.z = buf[s..][0..shared_length].*;
                return ret;
            }
        };

        /// A Kyber key pair.
        pub const KeyPair = struct {
            secret_key: SecretKey,
            public_key: PublicKey,

            /// Deterministically derive a key pair from a cryptograpically secure secret seed.
            ///
            /// Except in tests, applications should generally call `generate()` instead of this function.
            pub fn generateDeterministic(seed: [seed_length]u8) !KeyPair {
                var ret: KeyPair = undefined;

                // Generate inner key
                innerKeyFromSeed(
                    seed[0..inner_seed_length].*,
                    &ret.public_key.pk,
                    &ret.secret_key.sk,
                );
                ret.secret_key.pk = ret.public_key.pk;

                // Copy over z from seed.
                ret.secret_key.z = seed[inner_seed_length..seed_length].*;

                // Compute H(pk)
                sha3.Sha3_256.hash(&ret.public_key.pk.toBytes(), &ret.secret_key.hpk, .{});
                ret.public_key.hpk = ret.secret_key.hpk;

                return ret;
            }

            /// Generate a new, random key pair.
            pub fn generate(io: std.Io) KeyPair {
                var random_seed: [seed_length]u8 = undefined;
                while (true) {
                    io.random(&random_seed);
                    return generateDeterministic(random_seed) catch {
                        @branchHint(.unlikely);
                        continue;
                    };
                }
            }
        };

        // Size of plaintexts of the in
        const inner_plaintext_length: usize = Poly.compressedSize(1);

        const InnerPk = struct {
            rho: [32]u8, // ρ, the seed for the matrix A
            th: V, // NTT(t), normalized

            // Cached values
            aT: M,

            const encoded_length = V.encoded_length + 32;

            fn encrypt(
                pk: InnerPk,
                pt: *const [inner_plaintext_length]u8,
                seed: *const [32]u8,
            ) [ciphertext_length]u8 {
                // Sample r, e₁ and e₂ appropriately
                const rh = V.noise(p.eta1, 0, seed).ntt().barrettReduce();
                const e1 = V.noise(eta2, p.k, seed);
                const e2 = Poly.noise(eta2, 2 * p.k, seed);

                // Next we compute u = Aᵀ r + e₁.  First Aᵀ.
                var u: V = undefined;
                for (0..p.k) |i| {
                    // Note that coefficients of r are bounded by q and those of Aᵀ
                    // are bounded by 4.5q and so their product is bounded by 2¹⁵q
                    // as required for multiplication.
                    u.ps[i] = pk.aT.rows[i].dotHat(rh);
                }

                // Aᵀ and r were not in Montgomery form, so the Montgomery
                // multiplications in the inner product added a factor R⁻¹ which
                // the InvNTT cancels out.
                u = u.barrettReduce().invNTT().add(e1).normalize();

                // Next, compute v = <t, r> + e₂ + Decompress_q(m, 1)
                const v = pk.th.dotHat(rh).barrettReduce().invNTT()
                    .add(Poly.decompress(1, pt)).add(e2).normalize();

                return u.compress(p.du) ++ v.compress(p.dv);
            }

            fn toBytes(pk: InnerPk) [encoded_length]u8 {
                return pk.th.toBytes() ++ pk.rho;
            }

            fn fromBytes(buf: *const [encoded_length]u8) errors.NonCanonicalError!InnerPk {
                var ret: InnerPk = undefined;

                const th_bytes = buf[0..V.encoded_length];
                ret.th = V.fromBytes(th_bytes).normalize();

                if (p.ml_kem) {
                    // Verify that the coefficients used a canonical representation.
                    if (!mem.eql(u8, &ret.th.toBytes(), th_bytes)) {
                        return error.NonCanonical;
                    }
                }

                ret.rho = buf[V.encoded_length..encoded_length].*;
                ret.aT = M.uniform(ret.rho, true);
                return ret;
            }
        };

        // Private key of the inner PKE
        const InnerSk = struct {
            sh: V, // NTT(s), normalized
            const encoded_length = V.encoded_length;

            fn decrypt(sk: InnerSk, ct: *const [ciphertext_length]u8) [inner_plaintext_length]u8 {
                const u = V.decompress(p.du, ct[0..comptime V.compressedSize(p.du)]);
                const v = Poly.decompress(
                    p.dv,
                    ct[comptime V.compressedSize(p.du)..ciphertext_length],
                );

                // Compute m = v - <s, u>
                return v.sub(sk.sh.dotHat(u.ntt()).barrettReduce().invNTT())
                    .normalize().compress(1);
            }

            fn toBytes(sk: InnerSk) [encoded_length]u8 {
                return sk.sh.toBytes();
            }

            fn fromBytes(buf: *const [encoded_length]u8) InnerSk {
                var ret: InnerSk = undefined;
                ret.sh = V.fromBytes(buf).normalize();
                return ret;
            }
        };

        // Derives inner PKE keypair from given seed.
        fn innerKeyFromSeed(seed: [inner_seed_length]u8, pk: *InnerPk, sk: *InnerSk) void {
            var expanded_seed: [64]u8 = undefined;
            var h = sha3.Sha3_512.init(.{});
            h.update(&seed);
            if (p.ml_kem) h.update(&[1]u8{p.k});
            h.final(&expanded_seed);
            pk.rho = expanded_seed[0..32].*;
            const sigma = expanded_seed[32..64];
            pk.aT = M.uniform(pk.rho, false); // Expand ρ to A; we'll transpose later on

            // Sample secret vector s.
            sk.sh = V.noise(p.eta1, 0, sigma).ntt().normalize();

            const eh = PolyVec(p.k).noise(p.eta1, p.k, sigma).ntt(); // sample blind e.
            var th: V = undefined;

            // Next, we compute t = A s + e.
            for (0..p.k) |i| {
                // Note that coefficients of s are bounded by q and those of A
                // are bounded by 4.5q and so their product is bounded by 2¹⁵q
                // as required for multiplication.
                // A and s were not in Montgomery form, so the Montgomery
                // multiplications in the inner product added a factor R⁻¹ which
                // we'll cancel out with toMont().  This will also ensure the
                // coefficients of th are bounded in absolute value by q.
                th.ps[i] = pk.aT.rows[i].dotHat(sk.sh).toMont();
            }

            pk.th = th.add(eh).normalize(); // bounded by 8q
            pk.aT = pk.aT.transpose();
        }
    };
}