1use hkdf::Hkdf;
5use sha2::Sha256;
6use trussed::{
7 config::MAX_MEDIUM_DATA_LENGTH,
8 key::{Kind, Secrecy},
9 serde_extensions::ExtensionImpl,
10 service::ServiceResources,
11 store::{ClientKeystore, Keystore, Store},
12 types::CoreContext,
13 Platform,
14};
15use trussed_core::{
16 types::{Bytes, MediumData, ShortData},
17 Error,
18};
19use trussed_hkdf::{
20 HkdfExpandReply, HkdfExpandRequest, HkdfExtension, HkdfExtractReply, HkdfExtractRequest,
21 HkdfReply, HkdfRequest, KeyOrData, OkmId,
22};
23
24use crate::{StagingBackend, StagingContext};
25
26impl ExtensionImpl<HkdfExtension> for StagingBackend {
27 fn extension_request<P: Platform>(
28 &mut self,
29 core_ctx: &mut CoreContext,
30 _backend_ctx: &mut StagingContext,
31 request: &HkdfRequest,
32 resources: &mut ServiceResources<P>,
33 ) -> Result<HkdfReply, Error> {
34 let mut keystore = resources.keystore(core_ctx.path.clone())?;
35 Ok(match request {
36 HkdfRequest::Extract(req) => extract(req, &mut keystore)?.into(),
37 HkdfRequest::Expand(req) => expand(req, &mut keystore)?.into(),
38 })
39 }
40}
41
42fn get_mat<S: Store>(
43 req: &KeyOrData<MAX_MEDIUM_DATA_LENGTH>,
44 keystore: &mut ClientKeystore<S>,
45) -> Result<MediumData, Error> {
46 Ok(match req {
47 KeyOrData::Data(d) => d.clone(),
48 KeyOrData::Key(key_id) => {
49 let key_mat = keystore.load_key(Secrecy::Secret, None, key_id)?;
50 if !matches!(key_mat.kind, Kind::Symmetric(..) | Kind::Shared(..)) {
51 warn!("Attempt to HKDF on a private key");
52 return Err(Error::MechanismInvalid);
53 }
54 Bytes::try_from(&*key_mat.material).map_err(|_| {
55 warn!("Attempt to HKDF a too large key");
56 Error::InternalError
57 })?
58 }
59 })
60}
61
62fn extract<S: Store>(
63 req: &HkdfExtractRequest,
64 keystore: &mut ClientKeystore<S>,
65) -> Result<HkdfExtractReply, Error> {
66 let ikm = get_mat(&req.ikm, keystore)?;
67 let salt = req
68 .salt
69 .as_ref()
70 .map(|s| get_mat(s, keystore))
71 .transpose()?;
72 let salt_ref = salt.as_deref();
73 let (prk, _) = Hkdf::<Sha256>::extract(salt_ref, &ikm);
74 assert_eq!(prk.len(), 256 / 8);
75 let key_id = keystore.store_key(
76 req.storage,
77 Secrecy::Secret,
78 Kind::Symmetric(prk.len()),
79 &prk,
80 )?;
81 Ok(HkdfExtractReply { okm: OkmId(key_id) })
82}
83fn expand<S: Store>(
84 req: &HkdfExpandRequest,
85 keystore: &mut ClientKeystore<S>,
86) -> Result<HkdfExpandReply, Error> {
87 let prk = keystore.load_key(Secrecy::Secret, None, &req.prk.0)?;
88 if !matches!(prk.kind, Kind::Symmetric(32)) {
89 error!("Attempt to use wrong key for HKDF expand");
90 return Err(Error::ObjectHandleInvalid);
91 }
92
93 let hkdf = Hkdf::<Sha256>::from_prk(&prk.material).map_err(|_| {
94 warn!("Failed to create HKDF");
95 Error::InternalError
96 })?;
97 let mut okm = ShortData::new();
98 okm.resize_zero(req.len).map_err(|_| {
99 error!("Attempt to run HKDF with too large output");
100 Error::WrongMessageLength
101 })?;
102 hkdf.expand(&req.info, &mut okm).map_err(|_| {
103 warn!("Bad HKDF expand length");
104 Error::WrongMessageLength
105 })?;
106
107 let key = keystore.store_key(
108 req.storage,
109 Secrecy::Secret,
110 Kind::Symmetric(okm.len()),
111 &okm,
112 )?;
113
114 Ok(HkdfExpandReply { key })
115}