ml_kem/
param.rs

1//! This module encapsulates all of the compile-time logic related to parameter-set dependent sizes
2//! of objects.  `ParameterSet` captures the parameters in the form described by the ML-KEM
3//! specification.  `EncodingSize`, `VectorEncodingSize`, and `CbdSamplingSize` are "upstream" of
4//! `ParameterSet`; they provide basic logic about the size of encoded objects.  `PkeParams` and
5//! `KemParams` are "downstream" of `ParameterSet`; they define derived parameters relevant to
6//! K-PKE and ML-KEM.
7//!
8//! While the primary purpose of these traits is to describe the sizes of objects, in order to
9//! avoid leakage of complicated trait bounds, they also need to provide any logic that needs to
10//! know any details about object sizes.  For example, `VectorEncodingSize::flatten` needs to know
11//! that the size of an encoded vector is `K` times the size of an encoded polynomial.
12
13use crate::{B32, Ciphertext, Kem, algebra::NttVector};
14use array::{
15    Array,
16    typenum::{
17        U0, U2, U3, U4, U6, U12, U32, U384,
18        operator_aliases::{Prod, Sum},
19    },
20};
21use core::{
22    fmt::Debug,
23    ops::{Add, Div, Mul, Rem, Sub},
24};
25use module_lattice::{
26    ArraySize, Encode, EncodedPolynomialSize, EncodedVectorSize, EncodingSize, VectorEncodingSize,
27};
28
29#[cfg(doc)]
30use crate::Seed;
31
32/// An integer that describes a bit length to be used in CBD sampling
33#[allow(unreachable_pub)]
34pub trait CbdSamplingSize: ArraySize {
35    type SampleSize: EncodingSize;
36}
37
38impl CbdSamplingSize for U2 {
39    type SampleSize = U4;
40}
41
42impl CbdSamplingSize for U3 {
43    type SampleSize = U6;
44}
45
46/// A `ParameterSet` captures the parameters that describe a particular instance of ML-KEM.
47///
48/// There are three variants, corresponding to three different security levels.
49pub trait ParameterSet: Default + Clone + Debug + PartialEq {
50    /// The dimensionality of vectors and arrays
51    type K: ArraySize;
52
53    /// The bit width of the centered binary distribution used when sampling random polynomials in
54    /// key generation and encryption.
55    type Eta1: CbdSamplingSize;
56
57    /// The bit width of the centered binary distribution used when sampling error vectors during
58    /// encryption.
59    type Eta2: CbdSamplingSize;
60
61    /// The bit width of encoded integers in the `u` vector in a ciphertext
62    type Du: VectorEncodingSize<Self::K>;
63
64    /// The bit width of encoded integers in the `v` polynomial in a ciphertext
65    type Dv: EncodingSize;
66}
67
68pub(crate) type EncodedUSize<P> =
69    EncodedVectorSize<<P as ParameterSet>::Du, <P as ParameterSet>::K>;
70pub(crate) type EncodedVSize<P> = EncodedPolynomialSize<<P as ParameterSet>::Dv>;
71
72type EncodedU<P> = Array<u8, EncodedUSize<P>>;
73type EncodedV<P> = Array<u8, EncodedVSize<P>>;
74
75/// Derived parameter relevant to K-PKE
76pub trait PkeParams: Kem<SharedKeySize = U32> + ParameterSet {
77    type NttVectorSize: ArraySize;
78    type EncryptionKeySize: ArraySize;
79
80    fn encode_u12(p: &NttVector<Self::K>) -> EncodedNttVector<Self>;
81    fn decode_u12(v: &EncodedNttVector<Self>) -> NttVector<Self::K>;
82
83    fn concat_ct(u: EncodedU<Self>, v: EncodedV<Self>) -> Ciphertext<Self>;
84    fn split_ct(ct: &Ciphertext<Self>) -> (&EncodedU<Self>, &EncodedV<Self>);
85
86    fn concat_ek(t_hat: EncodedNttVector<Self>, rho: B32) -> EncodedEncryptionKey<Self>;
87    fn split_ek(ek: &EncodedEncryptionKey<Self>) -> (&EncodedNttVector<Self>, &B32);
88}
89
90pub(crate) type EncodedNttVector<P> = Array<u8, <P as PkeParams>::NttVectorSize>;
91pub(crate) type EncodedDecryptionKey<P> = Array<u8, <P as PkeParams>::NttVectorSize>;
92pub(crate) type EncodedEncryptionKey<P> = Array<u8, <P as PkeParams>::EncryptionKeySize>;
93
94impl<P> PkeParams for P
95where
96    P: Kem<CiphertextSize = Sum<EncodedUSize<P>, EncodedVSize<P>>, SharedKeySize = U32>
97        + ParameterSet,
98    U384: Mul<P::K>,
99    Prod<U384, P::K>: ArraySize + Add<U32> + Div<P::K, Output = U384> + Rem<P::K, Output = U0>,
100    EncodedUSize<P>: Add<EncodedVSize<P>>,
101    Sum<EncodedUSize<P>, EncodedVSize<P>>:
102        ArraySize + Sub<EncodedUSize<P>, Output = EncodedVSize<P>>,
103    EncodedVectorSize<U12, P::K>: Add<U32>,
104    Sum<EncodedVectorSize<U12, P::K>, U32>:
105        ArraySize + Sub<EncodedVectorSize<U12, P::K>, Output = U32>,
106{
107    type NttVectorSize = EncodedVectorSize<U12, P::K>;
108    type EncryptionKeySize = Sum<Self::NttVectorSize, U32>;
109
110    fn encode_u12(p: &NttVector<Self::K>) -> EncodedNttVector<Self> {
111        Encode::<U12>::encode(p)
112    }
113
114    fn decode_u12(v: &EncodedNttVector<Self>) -> NttVector<Self::K> {
115        Encode::<U12>::decode(v)
116    }
117
118    fn concat_ct(u: EncodedU<Self>, v: EncodedV<Self>) -> Ciphertext<Self> {
119        u.concat(v)
120    }
121
122    fn split_ct(ct: &Ciphertext<Self>) -> (&EncodedU<Self>, &EncodedV<Self>) {
123        ct.split_ref()
124    }
125
126    fn concat_ek(t_hat: EncodedNttVector<Self>, rho: B32) -> EncodedEncryptionKey<Self> {
127        t_hat.concat(rho)
128    }
129
130    fn split_ek(ek: &EncodedEncryptionKey<Self>) -> (&EncodedNttVector<Self>, &B32) {
131        ek.split_ref()
132    }
133}
134
135/// Derived parameters relevant to ML-KEM
136pub trait KemParams: PkeParams {
137    type DecapsulationKeySize: ArraySize;
138
139    fn concat_dk(
140        dk: EncodedDecryptionKey<Self>,
141        ek: EncodedEncryptionKey<Self>,
142        h: B32,
143        z: B32,
144    ) -> ExpandedDecapsulationKey<Self>;
145
146    fn split_dk(
147        enc: &ExpandedDecapsulationKey<Self>,
148    ) -> (
149        &EncodedDecryptionKey<Self>,
150        &EncodedEncryptionKey<Self>,
151        &B32,
152        &B32,
153    );
154}
155
156pub(crate) type DecapsulationKeySize<P> = <P as KemParams>::DecapsulationKeySize;
157pub(crate) type EncapsulationKeySize<P> = <P as PkeParams>::EncryptionKeySize;
158
159/// Serialized decapsulation key after having been expanded from a [`Seed`].
160pub type ExpandedDecapsulationKey<P> = Array<u8, <P as KemParams>::DecapsulationKeySize>;
161
162impl<P> KemParams for P
163where
164    P: PkeParams,
165    P::NttVectorSize: Add<P::EncryptionKeySize>,
166    Sum<P::NttVectorSize, P::EncryptionKeySize>:
167        ArraySize + Add<U32> + Sub<P::NttVectorSize, Output = P::EncryptionKeySize>,
168    Sum<Sum<P::NttVectorSize, P::EncryptionKeySize>, U32>:
169        ArraySize + Add<U32> + Sub<Sum<P::NttVectorSize, P::EncryptionKeySize>, Output = U32>,
170    Sum<Sum<Sum<P::NttVectorSize, P::EncryptionKeySize>, U32>, U32>:
171        ArraySize + Sub<Sum<Sum<P::NttVectorSize, P::EncryptionKeySize>, U32>, Output = U32>,
172{
173    type DecapsulationKeySize = Sum<Sum<Sum<P::NttVectorSize, P::EncryptionKeySize>, U32>, U32>;
174
175    fn concat_dk(
176        dk: EncodedDecryptionKey<Self>,
177        ek: EncodedEncryptionKey<Self>,
178        h: B32,
179        z: B32,
180    ) -> ExpandedDecapsulationKey<Self> {
181        dk.concat(ek).concat(h).concat(z)
182    }
183
184    #[allow(clippy::similar_names)] // allow dk_pke, ek_pke, following the spec
185    fn split_dk(
186        enc: &ExpandedDecapsulationKey<Self>,
187    ) -> (
188        &EncodedDecryptionKey<Self>,
189        &EncodedEncryptionKey<Self>,
190        &B32,
191        &B32,
192    ) {
193        // We parse from right to left to make it easier to write the trait bounds above
194        let (enc, z) = enc.split_ref();
195        let (enc, h) = enc.split_ref();
196        let (dk_pke, ek_pke) = enc.split_ref();
197        (dk_pke, ek_pke, h, z)
198    }
199}
ml_kem/param.rs

ml_kem/
param.rs
