Skip to main content

xml_sec/xmlenc/
key_establishment_budget.rs

1//! Operation-wide KDF preflight. No derived work precedes this gate.
2
3use crate::policy::{
4    KeyAgreementAlgorithm, KeyDerivationAlgorithm, KeyEstablishmentPolicy, PolicyViolation,
5};
6use crate::provider::{
7    CryptoProvider, KdfContext, KdfParameters, KeyAgreementKey, KeyAgreementParameters,
8    ProviderCapability, ProviderError, ProviderInputError, ProviderOperation,
9};
10use crate::xmldsig::{DigestAlgorithm, SignatureAlgorithm};
11
12use super::XmlEncError;
13pub(super) use crate::key_establishment::KeyEstablishmentUsage;
14use crate::key_establishment::reserve;
15
16/// Cumulative KDF allowance borrowed from one immutable operation policy.
17///
18/// Share this instance across every recipient, derived key and retry. Reservations
19/// are monotonic, including failed provider calls. Input slices remain borrowed;
20/// only the provider's returned key buffer is owned and zeroized by this API.
21pub struct KeyEstablishmentBudget<'a> {
22    policy: &'a KeyEstablishmentPolicy,
23    usage: KeyEstablishmentUsage,
24}
25
26impl<'a> KeyEstablishmentBudget<'a> {
27    /// Validate the snapshot before compiling or executing key-establishment work.
28    pub fn new(policy: &'a KeyEstablishmentPolicy) -> Result<Self, PolicyViolation> {
29        policy.validate()?;
30        Ok(Self {
31            policy,
32            usage: KeyEstablishmentUsage::default(),
33        })
34    }
35
36    /// Derive through the selected provider only after syntax, permission,
37    /// capability and aggregate allocation/CPU reservations have succeeded.
38    pub fn derive_key(
39        &mut self,
40        provider: &dyn CryptoProvider,
41        parameters: &KdfParameters<'_>,
42        secret: &[u8],
43    ) -> Result<zeroize::Zeroizing<Vec<u8>>, XmlEncError> {
44        self.usage
45            .derive_key(self.policy, provider, parameters, secret)
46    }
47
48    /// Agree and derive using the same monotonic operation allowance. Permission
49    /// and both workspace/output reservations precede either provider callback.
50    pub fn agree_and_derive(
51        &mut self,
52        provider: &dyn CryptoProvider,
53        key: &dyn KeyAgreementKey,
54        agreement: &KeyAgreementParameters<'_>,
55        parameters: &KdfParameters<'_>,
56    ) -> Result<zeroize::Zeroizing<Vec<u8>>, XmlEncError> {
57        self.usage
58            .agree_and_derive(self.policy, provider, key, agreement, parameters)
59    }
60
61    /// Cumulative compression work reserved, including failed attempts.
62    pub const fn hash_blocks(&self) -> usize {
63        self.usage.hash_blocks
64    }
65
66    /// Cumulative output allocations reserved, including failed attempts.
67    pub const fn owned_bytes(&self) -> usize {
68        self.usage.owned_bytes
69    }
70
71    /// Cumulative finite-field validation and exponentiation work reserved.
72    pub const fn modular_work(&self) -> usize {
73        self.usage.modular_work
74    }
75
76    pub(crate) fn reserve_dh_import(
77        &mut self,
78        p_bits: usize,
79        q_bits: usize,
80        p_twos: usize,
81        q_twos: usize,
82    ) -> Result<(), XmlEncError> {
83        self.policy.validate()?;
84        if self
85            .policy
86            .check_agreement(KeyAgreementAlgorithm::DhEs)
87            .is_err()
88        {
89            self.policy
90                .check_agreement(KeyAgreementAlgorithm::LegacyDh)?;
91        }
92        self.policy.check_dh_domain(p_bits, q_bits)?;
93        reserve(
94            crate::policy::resource_name::DH_MODULUS_BITS,
95            0,
96            p_bits as u128,
97            self.policy.max_dh_modulus_bits,
98        )?;
99        // crypto-primes 0.7 MillerRabin holds its parameters across 64 rounds.
100        // Each round allocates a 16-entry power table plus conversion buffers,
101        // then at most v2(n-1) squares. Include boxed headers and public-domain
102        // import/division/Montgomery setup rather than charging payload alone.
103        let bytes = |bits: usize, twos: usize| {
104            let slot = bits.div_ceil(64) as u128 * 8 + 64;
105            slot * (512 + 64 * (64 + twos as u128))
106        };
107        self.usage.commit_dh_reservation(
108            self.policy,
109            modular_cost(p_bits, p_bits, 64)
110                + modular_cost(q_bits, q_bits, 64)
111                + modular_cost(p_bits, p_bits, 2),
112            bytes(p_bits, p_twos) + bytes(q_bits, q_twos),
113        )
114    }
115}
116
117fn modular_cost(bits: usize, exponent_bits: usize, count: usize) -> u128 {
118    let limbs = bits.div_ceil(usize::BITS as usize) as u128;
119    // Bound native-word product passes for the fixed-window exponentiation,
120    // Miller-Rabin's trailing squares and reduction, including window setup.
121    // Counting native words also preserves the bound on 32-bit targets.
122    (limbs * limbs).div_ceil(64) * (exponent_bits as u128 + 64) * 8 * count as u128
123}
124
125impl KeyEstablishmentUsage {
126    pub(super) fn agree_and_derive(
127        &mut self,
128        policy: &KeyEstablishmentPolicy,
129        provider: &dyn CryptoProvider,
130        key: &dyn KeyAgreementKey,
131        agreement: &KeyAgreementParameters<'_>,
132        parameters: &KdfParameters<'_>,
133    ) -> Result<zeroize::Zeroizing<Vec<u8>>, XmlEncError> {
134        policy.validate()?;
135        let mechanism = KeyAgreementAlgorithm::from_uri(agreement.algorithm)
136            .ok_or_else(|| XmlEncError::UnsupportedAlgorithm(agreement.algorithm.into()))?;
137        policy.check_agreement(mechanism)?;
138        let width = match (mechanism, agreement.peer_public_key) {
139            (KeyAgreementAlgorithm::X25519, bytes) if bytes.len() == 32 => 32,
140            (KeyAgreementAlgorithm::X448, bytes) if bytes.len() == 56 => 56,
141            (KeyAgreementAlgorithm::EcdhEs, [4, tail @ ..])
142                if matches!(tail.len(), 64 | 96 | 132) =>
143            {
144                tail.len() / 2
145            }
146            (KeyAgreementAlgorithm::EcdhEs, [2 | 3, tail @ ..])
147                if matches!(tail.len(), 32 | 48 | 66) =>
148            {
149                tail.len()
150            }
151            (KeyAgreementAlgorithm::DhEs | KeyAgreementAlgorithm::LegacyDh, bytes) => {
152                let (p_bits, q_bits) = key
153                    .dh_domain_bits()
154                    .ok_or_else(|| invalid(ProviderInputError::DhKey))?;
155                policy.check_dh_domain(p_bits, q_bits)?;
156                let width = p_bits.div_ceil(8);
157                if bytes.is_empty() || bytes.len() > width || width == 0 {
158                    return Err(invalid(ProviderInputError::DhKey));
159                }
160                width
161            }
162            _ => return Err(invalid(ProviderInputError::EcdhKey)),
163        };
164        // XMLEnc 1.1 §5.6.4 fixes ZZ to the field-element octet width. Reserve
165        // both ZZ and the consuming key before either provider callback; a
166        // denied KDF must not initiate otherwise permitted scalar multiplication.
167        // https://www.w3.org/TR/2013/REC-xmlenc-core1-20130411/#sec-ECDH-ES
168        let (algorithm, digest, blocks) = preflight(parameters, width)?;
169        // The legacy agreement URI fixes its implicit KDF. DH-ES permits an
170        // explicit KDF, but replacing legacy derivation silently changes the
171        // wire algorithm, even if both mechanisms are separately permitted.
172        // XMLEnc 1.1 §§5.6.2.1–5.6.2.2:
173        // https://www.w3.org/TR/2013/REC-xmlenc-core1-20130411/#sec-DHKeyAgreementLegacyKDF
174        if mechanism == KeyAgreementAlgorithm::LegacyDh
175            && algorithm != KeyDerivationAlgorithm::LegacyDh
176        {
177            return Err(invalid(ProviderInputError::LegacyDhKdfParameters));
178        }
179        policy.check_derivation(algorithm, digest)?;
180        provider.require_capability(ProviderCapability::KeyAgreement(agreement))?;
181        provider.require_capability(ProviderCapability::Kdf(parameters))?;
182        let dh = matches!(
183            mechanism,
184            KeyAgreementAlgorithm::DhEs | KeyAgreementAlgorithm::LegacyDh
185        );
186        let modular = if dh {
187            modular_cost(width * 8, width * 8, 2)
188        } else {
189            0
190        };
191        let workspace = if dh { (width as u128 + 64) * 256 } else { 0 };
192        self.commit_all(
193            policy,
194            blocks,
195            modular,
196            width as u128 + parameters.output_len as u128 + workspace,
197        )?;
198        let secret = zeroize::Zeroizing::new(provider.agree_key(key, agreement)?);
199        check_width(&secret, width)?;
200        let output = zeroize::Zeroizing::new(provider.derive_key(parameters, &secret)?);
201        check_width(&output, parameters.output_len)?;
202        Ok(output)
203    }
204
205    pub(super) fn derive_key(
206        &mut self,
207        policy: &KeyEstablishmentPolicy,
208        provider: &dyn CryptoProvider,
209        parameters: &KdfParameters<'_>,
210        secret: &[u8],
211    ) -> Result<zeroize::Zeroizing<Vec<u8>>, XmlEncError> {
212        policy.validate()?;
213        let (algorithm, digest, blocks) = preflight(parameters, secret.len())?;
214        policy.check_derivation(algorithm, digest)?;
215        provider.require_capability(ProviderCapability::Kdf(parameters))?;
216        self.commit_reservation(policy, blocks, parameters.output_len as u128)?;
217        let output = zeroize::Zeroizing::new(provider.derive_key(parameters, secret)?);
218        check_width(&output, parameters.output_len)?;
219        Ok(output)
220    }
221
222    fn commit_reservation(
223        &mut self,
224        policy: &KeyEstablishmentPolicy,
225        blocks: u128,
226        allocations: u128,
227    ) -> Result<(), XmlEncError> {
228        Ok(self.commit_all(policy, blocks, 0, allocations)?)
229    }
230
231    fn commit_dh_reservation(
232        &mut self,
233        policy: &KeyEstablishmentPolicy,
234        modular: u128,
235        allocations: u128,
236    ) -> Result<(), XmlEncError> {
237        Ok(self.commit_all(policy, 0, modular, allocations)?)
238    }
239}
240
241fn check_width(bytes: &[u8], expected: usize) -> Result<(), XmlEncError> {
242    if bytes.len() != expected {
243        return Err(ProviderError::InvalidKeySize {
244            expected,
245            actual: bytes.len(),
246        }
247        .into());
248    }
249    Ok(())
250}
251
252fn invalid(kind: ProviderInputError) -> XmlEncError {
253    ProviderError::InvalidInput(kind).into()
254}
255
256fn unsupported(uri: &str) -> XmlEncError {
257    ProviderError::Unsupported {
258        operation: ProviderOperation::Kdf,
259        algorithm: Some(uri.into()),
260    }
261    .into()
262}
263
264// SHA compression-block costs include padding and HMAC's inner/outer hashes,
265// not merely the user-supplied PBKDF2 iteration count. The bound follows the
266// provider's streaming framing without allocating concatenated input strings.
267// RFC 8018 §5.2, RFC 5869 §§2.2–2.3, XMLEnc 1.1 §§5.4.1 and 5.6.2.2:
268// https://www.rfc-editor.org/rfc/rfc8018.html#section-5.2
269// https://www.rfc-editor.org/rfc/rfc5869.html#section-2
270// https://www.w3.org/TR/2013/REC-xmlenc-core1-20130411/#sec-ConcatKDF
271fn preflight(
272    parameters: &KdfParameters<'_>,
273    secret_len: usize,
274) -> Result<(KeyDerivationAlgorithm, DigestAlgorithm, u128), XmlEncError> {
275    let algorithm = KeyDerivationAlgorithm::from_uri(parameters.algorithm)
276        .ok_or_else(|| unsupported(parameters.algorithm))?;
277    let uri = parameters
278        .digest
279        .ok_or_else(|| unsupported(parameters.algorithm))?;
280    let digest = match algorithm {
281        KeyDerivationAlgorithm::Hkdf | KeyDerivationAlgorithm::Pbkdf2 => {
282            match SignatureAlgorithm::from_uri(uri) {
283                Some(SignatureAlgorithm::HmacSha1) => DigestAlgorithm::Sha1,
284                Some(SignatureAlgorithm::HmacSha224) => DigestAlgorithm::Sha224,
285                Some(SignatureAlgorithm::HmacSha256) => DigestAlgorithm::Sha256,
286                Some(SignatureAlgorithm::HmacSha384) => DigestAlgorithm::Sha384,
287                Some(SignatureAlgorithm::HmacSha512) => DigestAlgorithm::Sha512,
288                _ => return Err(unsupported(uri)),
289            }
290        }
291        _ => match DigestAlgorithm::from_uri(uri) {
292            Some(
293                value @ (DigestAlgorithm::Sha1
294                | DigestAlgorithm::Sha224
295                | DigestAlgorithm::Sha256
296                | DigestAlgorithm::Sha384
297                | DigestAlgorithm::Sha512
298                | DigestAlgorithm::Sha3_224
299                | DigestAlgorithm::Sha3_256
300                | DigestAlgorithm::Sha3_384
301                | DigestAlgorithm::Sha3_512),
302            ) => value,
303            #[cfg(feature = "legacy-algorithms")]
304            Some(value @ (DigestAlgorithm::Md5 | DigestAlgorithm::Ripemd160)) => value,
305            _ => return Err(unsupported(uri)),
306        },
307    };
308    let width = digest.output_len();
309    // FIPS 202 §§5.1, 6.1: SHA-3's rate depends on output width, not
310    // SHA-2's 512/1024-bit compression blocks. Round partial octets upward;
311    // one extra octet bounds its domain and pad10*1, including a rate crossing.
312    // https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.202.pdf
313    let (block, padding) = match digest {
314        DigestAlgorithm::Sha3_224 => (144u128, 1u128),
315        DigestAlgorithm::Sha3_256 => (136, 1),
316        DigestAlgorithm::Sha3_384 => (104, 1),
317        DigestAlgorithm::Sha3_512 => (72, 1),
318        _ if width <= 32 => (64, 9),
319        _ => (128, 17),
320    };
321    let hash = |bytes: u128| (bytes + padding).div_ceil(block);
322    let hmac_key = |bytes: usize| {
323        if bytes as u128 > block {
324            hash(bytes as u128) + 2
325        } else {
326            2
327        }
328    };
329    let output_blocks = parameters.output_len.div_ceil(width) as u128;
330    let work = match algorithm {
331        KeyDerivationAlgorithm::Pbkdf2 => {
332            if parameters.iterations == 0
333                || parameters.output_len == 0
334                || octets(parameters.info).is_none_or(|info| !info.is_empty())
335                || output_blocks > u32::MAX as u128
336            {
337                return Err(invalid(ProviderInputError::Pbkdf2Parameters));
338            }
339            // The PRF and salt prefix are prepared once, then cloned. Account
340            // their actual compression work once, not once per output block.
341            let salt = parameters.salt.len() as u128;
342            hmac_key(secret_len)
343                + salt / block
344                + output_blocks
345                    * (hash(salt % block + 4) + 1 + 2 * (parameters.iterations as u128 - 1))
346        }
347        KeyDerivationAlgorithm::Hkdf => {
348            let info = octets(parameters.info)
349                .ok_or_else(|| invalid(ProviderInputError::HkdfParameters))?;
350            if parameters.iterations != 0 || output_blocks > 255 {
351                return Err(invalid(ProviderInputError::HkdfParameters));
352            }
353            // Conservatively include a previous full digest in every T_i,
354            // including T_1, which actually uses an empty previous value.
355            hmac_key(parameters.salt.len())
356                + hash(secret_len as u128)
357                + 1
358                + 2
359                + output_blocks * (hash(width as u128 + info.len() as u128 + 1) + 1)
360        }
361        KeyDerivationAlgorithm::ConcatKdf => {
362            if parameters.iterations != 0
363                || !parameters.salt.is_empty()
364                || parameters.output_len == 0
365                || output_blocks > u32::MAX as u128
366            {
367                return Err(invalid(ProviderInputError::ConcatKdfParameters));
368            }
369            let (bytes, bits) = match parameters.info {
370                KdfContext::Octets(bytes) => (bytes, bytes.len() as u128 * 8),
371                KdfContext::Bits { bytes, bit_len } => (bytes, bit_len as u128),
372                _ => return Err(invalid(ProviderInputError::ConcatKdfParameters)),
373            };
374            let used = (bits % 8) as usize;
375            if bits.div_ceil(8) != bytes.len() as u128
376                || (used != 0 && bytes[bytes.len() - 1] & (0xff >> used) != 0)
377            {
378                return Err(invalid(ProviderInputError::ConcatKdfParameters));
379            }
380            output_blocks * hash(4 + secret_len as u128 + bits.div_ceil(8))
381        }
382        KeyDerivationAlgorithm::LegacyDh => {
383            let KdfContext::LegacyDh {
384                encryption_algorithm,
385                nonce,
386            } = parameters.info
387            else {
388                return Err(invalid(ProviderInputError::LegacyDhKdfParameters));
389            };
390            if parameters.iterations != 0
391                || !parameters.salt.is_empty()
392                || parameters.output_len == 0
393                || output_blocks > 255
394                || parameters.output_len.checked_mul(8).is_none()
395            {
396                return Err(invalid(ProviderInputError::LegacyDhKdfParameters));
397            }
398            output_blocks
399                * hash(
400                    secret_len as u128
401                        + 2
402                        + encryption_algorithm.len() as u128
403                        + nonce.len() as u128
404                        + 40,
405                )
406        }
407    };
408    Ok((algorithm, digest, work))
409}
410
411fn octets(context: KdfContext<'_>) -> Option<&[u8]> {
412    match context {
413        KdfContext::Octets(bytes) => Some(bytes),
414        KdfContext::Bits { bytes, bit_len }
415            if bit_len.is_multiple_of(8) && bit_len / 8 == bytes.len() =>
416        {
417            Some(bytes)
418        }
419        _ => None,
420    }
421}