1use 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
16pub struct KeyEstablishmentBudget<'a> {
22 policy: &'a KeyEstablishmentPolicy,
23 usage: KeyEstablishmentUsage,
24}
25
26impl<'a> KeyEstablishmentBudget<'a> {
27 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 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 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 pub const fn hash_blocks(&self) -> usize {
63 self.usage.hash_blocks
64 }
65
66 pub const fn owned_bytes(&self) -> usize {
68 self.usage.owned_bytes
69 }
70
71 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 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 (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 let (algorithm, digest, blocks) = preflight(parameters, width)?;
168 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
263fn 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 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 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}