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::EcdhEs, [4, tail @ ..])
141                if matches!(tail.len(), 64 | 96 | 132) =>
142            {
143                tail.len() / 2
144            }
145            (KeyAgreementAlgorithm::EcdhEs, [2 | 3, tail @ ..])
146                if matches!(tail.len(), 32 | 48 | 66) =>
147            {
148                tail.len()
149            }
150            (KeyAgreementAlgorithm::DhEs | KeyAgreementAlgorithm::LegacyDh, bytes) => {
151                let (p_bits, q_bits) = key
152                    .dh_domain_bits()
153                    .ok_or_else(|| invalid(ProviderInputError::DhKey))?;
154                policy.check_dh_domain(p_bits, q_bits)?;
155                let width = p_bits.div_ceil(8);
156                if bytes.is_empty() || bytes.len() > width || width == 0 {
157                    return Err(invalid(ProviderInputError::DhKey));
158                }
159                width
160            }
161            _ => return Err(invalid(ProviderInputError::EcdhKey)),
162        };
163        // XMLEnc 1.1 §5.6.4 fixes ZZ to the field-element octet width. Reserve
164        // both ZZ and the consuming key before either provider callback; a
165        // denied KDF must not initiate otherwise permitted scalar multiplication.
166        // https://www.w3.org/TR/2013/REC-xmlenc-core1-20130411/#sec-ECDH-ES
167        let (algorithm, digest, blocks) = preflight(parameters, width)?;
168        // The legacy agreement URI fixes its implicit KDF. DH-ES permits an
169        // explicit KDF, but replacing legacy derivation silently changes the
170        // wire algorithm, even if both mechanisms are separately permitted.
171        // XMLEnc 1.1 §§5.6.2.1–5.6.2.2:
172        // https://www.w3.org/TR/2013/REC-xmlenc-core1-20130411/#sec-DHKeyAgreementLegacyKDF
173        if mechanism == KeyAgreementAlgorithm::LegacyDh
174            && algorithm != KeyDerivationAlgorithm::LegacyDh
175        {
176            return Err(invalid(ProviderInputError::LegacyDhKdfParameters));
177        }
178        policy.check_derivation(algorithm, digest)?;
179        provider.require_capability(ProviderCapability::KeyAgreement(agreement))?;
180        provider.require_capability(ProviderCapability::Kdf(parameters))?;
181        let dh = matches!(
182            mechanism,
183            KeyAgreementAlgorithm::DhEs | KeyAgreementAlgorithm::LegacyDh
184        );
185        let modular = if dh {
186            modular_cost(width * 8, width * 8, 2)
187        } else {
188            0
189        };
190        let workspace = if dh { (width as u128 + 64) * 256 } else { 0 };
191        self.commit_all(
192            policy,
193            blocks,
194            modular,
195            width as u128 + parameters.output_len as u128 + workspace,
196        )?;
197        let secret = zeroize::Zeroizing::new(provider.agree_key(key, agreement)?);
198        check_width(&secret, width)?;
199        let output = zeroize::Zeroizing::new(provider.derive_key(parameters, &secret)?);
200        check_width(&output, parameters.output_len)?;
201        Ok(output)
202    }
203
204    pub(super) fn derive_key(
205        &mut self,
206        policy: &KeyEstablishmentPolicy,
207        provider: &dyn CryptoProvider,
208        parameters: &KdfParameters<'_>,
209        secret: &[u8],
210    ) -> Result<zeroize::Zeroizing<Vec<u8>>, XmlEncError> {
211        policy.validate()?;
212        let (algorithm, digest, blocks) = preflight(parameters, secret.len())?;
213        policy.check_derivation(algorithm, digest)?;
214        provider.require_capability(ProviderCapability::Kdf(parameters))?;
215        self.commit_reservation(policy, blocks, parameters.output_len as u128)?;
216        let output = zeroize::Zeroizing::new(provider.derive_key(parameters, secret)?);
217        check_width(&output, parameters.output_len)?;
218        Ok(output)
219    }
220
221    fn commit_reservation(
222        &mut self,
223        policy: &KeyEstablishmentPolicy,
224        blocks: u128,
225        allocations: u128,
226    ) -> Result<(), XmlEncError> {
227        Ok(self.commit_all(policy, blocks, 0, allocations)?)
228    }
229
230    fn commit_dh_reservation(
231        &mut self,
232        policy: &KeyEstablishmentPolicy,
233        modular: u128,
234        allocations: u128,
235    ) -> Result<(), XmlEncError> {
236        Ok(self.commit_all(policy, 0, modular, allocations)?)
237    }
238}
239
240fn check_width(bytes: &[u8], expected: usize) -> Result<(), XmlEncError> {
241    if bytes.len() != expected {
242        return Err(ProviderError::InvalidKeySize {
243            expected,
244            actual: bytes.len(),
245        }
246        .into());
247    }
248    Ok(())
249}
250
251fn invalid(kind: ProviderInputError) -> XmlEncError {
252    ProviderError::InvalidInput(kind).into()
253}
254
255fn unsupported(uri: &str) -> XmlEncError {
256    ProviderError::Unsupported {
257        operation: ProviderOperation::Kdf,
258        algorithm: Some(uri.into()),
259    }
260    .into()
261}
262
263// SHA compression-block costs include padding and HMAC's inner/outer hashes,
264// not merely the user-supplied PBKDF2 iteration count. The bound follows the
265// provider's streaming framing without allocating concatenated input strings.
266// RFC 8018 §5.2, RFC 5869 §§2.2–2.3, XMLEnc 1.1 §§5.4.1 and 5.6.2.2:
267// https://www.rfc-editor.org/rfc/rfc8018.html#section-5.2
268// https://www.rfc-editor.org/rfc/rfc5869.html#section-2
269// https://www.w3.org/TR/2013/REC-xmlenc-core1-20130411/#sec-ConcatKDF
270fn preflight(
271    parameters: &KdfParameters<'_>,
272    secret_len: usize,
273) -> Result<(KeyDerivationAlgorithm, DigestAlgorithm, u128), XmlEncError> {
274    let algorithm = KeyDerivationAlgorithm::from_uri(parameters.algorithm)
275        .ok_or_else(|| unsupported(parameters.algorithm))?;
276    let uri = parameters
277        .digest
278        .ok_or_else(|| unsupported(parameters.algorithm))?;
279    let digest = match algorithm {
280        KeyDerivationAlgorithm::Hkdf | KeyDerivationAlgorithm::Pbkdf2 => {
281            match SignatureAlgorithm::from_uri(uri) {
282                Some(SignatureAlgorithm::HmacSha1) => DigestAlgorithm::Sha1,
283                Some(SignatureAlgorithm::HmacSha224) => DigestAlgorithm::Sha224,
284                Some(SignatureAlgorithm::HmacSha256) => DigestAlgorithm::Sha256,
285                Some(SignatureAlgorithm::HmacSha384) => DigestAlgorithm::Sha384,
286                Some(SignatureAlgorithm::HmacSha512) => DigestAlgorithm::Sha512,
287                _ => return Err(unsupported(uri)),
288            }
289        }
290        _ => match DigestAlgorithm::from_uri(uri) {
291            Some(
292                value @ (DigestAlgorithm::Sha1
293                | DigestAlgorithm::Sha224
294                | DigestAlgorithm::Sha256
295                | DigestAlgorithm::Sha384
296                | DigestAlgorithm::Sha512),
297            ) => value,
298            _ => return Err(unsupported(uri)),
299        },
300    };
301    let width = digest.output_len();
302    let block = if width <= 32 { 64u128 } else { 128 };
303    let padding = if width <= 32 { 9u128 } else { 17 };
304    let hash = |bytes: u128| (bytes + padding).div_ceil(block);
305    let hmac_key = |bytes: usize| {
306        if bytes as u128 > block {
307            hash(bytes as u128) + 2
308        } else {
309            2
310        }
311    };
312    let output_blocks = parameters.output_len.div_ceil(width) as u128;
313    let work = match algorithm {
314        KeyDerivationAlgorithm::Pbkdf2 => {
315            if parameters.iterations == 0
316                || parameters.output_len == 0
317                || octets(parameters.info).is_none_or(|info| !info.is_empty())
318                || output_blocks > u32::MAX as u128
319            {
320                return Err(invalid(ProviderInputError::Pbkdf2Parameters));
321            }
322            // The PRF and salt prefix are prepared once, then cloned. Account
323            // their actual compression work once, not once per output block.
324            let salt = parameters.salt.len() as u128;
325            hmac_key(secret_len)
326                + salt / block
327                + output_blocks
328                    * (hash(salt % block + 4) + 1 + 2 * (parameters.iterations as u128 - 1))
329        }
330        KeyDerivationAlgorithm::Hkdf => {
331            let info = octets(parameters.info)
332                .ok_or_else(|| invalid(ProviderInputError::HkdfParameters))?;
333            if parameters.iterations != 0 || output_blocks > 255 {
334                return Err(invalid(ProviderInputError::HkdfParameters));
335            }
336            // Conservatively include a previous full digest in every T_i,
337            // including T_1, which actually uses an empty previous value.
338            hmac_key(parameters.salt.len())
339                + hash(secret_len as u128)
340                + 1
341                + 2
342                + output_blocks * (hash(width as u128 + info.len() as u128 + 1) + 1)
343        }
344        KeyDerivationAlgorithm::ConcatKdf => {
345            if parameters.iterations != 0
346                || !parameters.salt.is_empty()
347                || parameters.output_len == 0
348                || output_blocks > u32::MAX as u128
349            {
350                return Err(invalid(ProviderInputError::ConcatKdfParameters));
351            }
352            let (bytes, bits) = match parameters.info {
353                KdfContext::Octets(bytes) => (bytes, bytes.len() as u128 * 8),
354                KdfContext::Bits { bytes, bit_len } => (bytes, bit_len as u128),
355                _ => return Err(invalid(ProviderInputError::ConcatKdfParameters)),
356            };
357            let used = (bits % 8) as usize;
358            if bits.div_ceil(8) != bytes.len() as u128
359                || (used != 0 && bytes[bytes.len() - 1] & (0xff >> used) != 0)
360            {
361                return Err(invalid(ProviderInputError::ConcatKdfParameters));
362            }
363            output_blocks * hash(4 + secret_len as u128 + bits.div_ceil(8))
364        }
365        KeyDerivationAlgorithm::LegacyDh => {
366            let KdfContext::LegacyDh {
367                encryption_algorithm,
368                nonce,
369            } = parameters.info
370            else {
371                return Err(invalid(ProviderInputError::LegacyDhKdfParameters));
372            };
373            if parameters.iterations != 0
374                || !parameters.salt.is_empty()
375                || parameters.output_len == 0
376                || output_blocks > 255
377                || parameters.output_len.checked_mul(8).is_none()
378            {
379                return Err(invalid(ProviderInputError::LegacyDhKdfParameters));
380            }
381            output_blocks
382                * hash(
383                    secret_len as u128
384                        + 2
385                        + encryption_algorithm.len() as u128
386                        + nonce.len() as u128
387                        + 40,
388                )
389        }
390    };
391    Ok((algorithm, digest, work))
392}
393
394fn octets(context: KdfContext<'_>) -> Option<&[u8]> {
395    match context {
396        KdfContext::Octets(bytes) => Some(bytes),
397        KdfContext::Bits { bytes, bit_len }
398            if bit_len.is_multiple_of(8) && bit_len / 8 == bytes.len() =>
399        {
400            Some(bytes)
401        }
402        _ => None,
403    }
404}