Skip to content

Commit ad1efa0

Browse files
ionut-armionut-arm
authored
Merge pull request #17 from ionut-arm/unwrap_expect
Improve usage of unwrap and expect
2 parents a63779a + 146c270 commit ad1efa0

File tree

4 files changed

+81
-78
lines changed

4 files changed

+81
-78
lines changed

‎src/abstraction/transient.rs

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ impl TransientObjectContext {
3535
if root_key_auth_size > 32 {
3636
return Err(Error::local_error(ErrorKind::WrongParamSize));
3737
}
38-
if root_key_size < 1024 {
38+
if root_key_size < 1024 || root_key_size > 4096 {
3939
return Err(Error::local_error(ErrorKind::WrongParamSize));
4040
}
4141
let mut context = Context::new(tcti)?;
@@ -50,7 +50,7 @@ impl TransientObjectContext {
5050

5151
let root_key_handle = context.create_primary_key(
5252
ESYS_TR_RH_OWNER,
53-
&get_rsa_public(true, true, false, root_key_size.try_into().unwrap()),
53+
&get_rsa_public(true, true, false, root_key_size.try_into().unwrap()), // should not fail on supported targets, given the checks above
5454
&root_key_auth,
5555
&[],
5656
&[],
@@ -83,7 +83,7 @@ impl TransientObjectContext {
8383
if auth_size > 32 {
8484
return Err(Error::local_error(ErrorKind::WrongParamSize));
8585
}
86-
if key_size < 1024 {
86+
if key_size < 1024 || key_size > 4096 {
8787
return Err(Error::local_error(ErrorKind::WrongParamSize));
8888
}
8989
let key_auth = if auth_size > 0 {
@@ -95,7 +95,7 @@ impl TransientObjectContext {
9595
self.set_session_attrs()?;
9696
let (key_priv, key_pub) = self.context.create_key(
9797
self.root_key_handle,
98-
&get_rsa_public(false, false, true, key_size.try_into().unwrap()),
98+
&get_rsa_public(false, false, true, key_size.try_into().unwrap()), // should not fail on valid targets, given the checks above
9999
&key_auth,
100100
&[],
101101
&[],
@@ -122,7 +122,7 @@ impl TransientObjectContext {
122122

123123
let pk = TPMU_PUBLIC_ID {
124124
rsa: TPM2B_PUBLIC_KEY_RSA {
125-
size: public_key.len().try_into().unwrap(),
125+
size: public_key.len().try_into().unwrap(), // should not fail on valid targets, given the checks above
126126
buffer: pk_buffer,
127127
},
128128
};
@@ -131,7 +131,7 @@ impl TransientObjectContext {
131131
false,
132132
false,
133133
true,
134-
u16::try_from(public_key.len()).unwrap() * 8u16,
134+
u16::try_from(public_key.len()).unwrap() * 8u16, // should not fail on valid targets, given the checks above
135135
);
136136
public.publicArea.unique = pk;
137137

@@ -160,7 +160,7 @@ impl TransientObjectContext {
160160
let key = match PublicIdUnion::from_public(&key_pub_id) {
161161
PublicIdUnion::Rsa(pub_key) => {
162162
let mut key = pub_key.buffer.to_vec();
163-
key.truncate(pub_key.size.try_into().unwrap());
163+
key.truncate(pub_key.size.try_into().unwrap()); // should not fail on supported targets
164164
key
165165
}
166166
_ => unimplemented!(),
@@ -257,7 +257,7 @@ mod tests {
257257
let mut ctx = TransientObjectContext::new(Tcti::Mssim, 2048, 32, &[]).unwrap();
258258
for _ in 0..4 {
259259
let (key, auth) = ctx.create_rsa_signing_key(2048, 16).unwrap();
260-
let mut signature = ctx.sign(key.clone(), &auth, &HASH).unwrap();
260+
let signature = ctx.sign(key.clone(), &auth, &HASH).unwrap();
261261
let pub_key = ctx.read_public_key(key.clone()).unwrap();
262262
let pub_key = ctx.load_external_rsa_public_key(&pub_key).unwrap();
263263
ctx.verify_signature(pub_key, &HASH, signature).unwrap();

‎src/lib.rs

Lines changed: 25 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -54,10 +54,13 @@ use utils::{Signature, TpmaSession, TpmsContext};
5454
#[macro_use]
5555
macro_rules! wrap_buffer {
5656
($buf:expr, $buf_type:ty, $buf_size:expr) => {{
57+
if $buf.len() > $buf_size {
58+
return Err(Error::local_error(ErrorKind::WrongParamSize));
59+
}
5760
let mut buffer = [0u8; $buf_size];
5861
buffer[..$buf.len()].clone_from_slice(&$buf[..$buf.len()]);
5962
let mut buf_struct: $buf_type = Default::default();
60-
buf_struct.size = $buf.len().try_into().unwrap();
63+
buf_struct.size = $buf.len().try_into().unwrap(); // should not fail since the length is checked above
6164
buf_struct.buffer = buffer;
6265
buf_struct
6366
}};
@@ -115,7 +118,7 @@ impl Context {
115118
let ret = unsafe {
116119
tss2_esys::Esys_Initialize(
117120
&mut esys_context,
118-
tcti_context.as_mut().unwrap().as_mut_ptr(),
121+
tcti_context.as_mut().unwrap().as_mut_ptr(), // will not panic as per how tcti_context is initialised
119122
null_mut(),
120123
)
121124
};
@@ -162,10 +165,6 @@ impl Context {
162165
symmetric: TPMT_SYM_DEF,
163166
auth_hash: TPMI_ALG_HASH,
164167
) -> Result<ESYS_TR> {
165-
if nonce.len() > 64 {
166-
return Err(Error::local_error(ErrorKind::WrongParamSize));
167-
}
168-
169168
let nonce_caller = wrap_buffer!(nonce, TPM2B_NONCE, 64);
170169
let mut sess = ESYS_TR_NONE;
171170

@@ -218,30 +217,26 @@ impl Context {
218217
outside_info: &[u8],
219218
creation_pcrs: &[TPMS_PCR_SELECTION],
220219
) -> Result<ESYS_TR> {
221-
if auth_value.len() > 64
222-
|| initial_data.len() > 256
223-
|| outside_info.len() > 64
224-
|| creation_pcrs.len() > 16
225-
{
226-
return Err(Error::local_error(ErrorKind::WrongParamSize));
227-
}
228-
229220
let sensitive_create = TPM2B_SENSITIVE_CREATE {
230221
size: std::mem::size_of::<TPMS_SENSITIVE_CREATE>()
231222
.try_into()
232-
.unwrap(),
223+
.unwrap(), // will not fail on targets of at least 16 bits
233224
sensitive: TPMS_SENSITIVE_CREATE {
234225
userAuth: wrap_buffer!(auth_value, TPM2B_AUTH, 64),
235226
data: wrap_buffer!(initial_data, TPM2B_SENSITIVE_DATA, 256),
236227
},
237228
};
238-
239229
let outside_info = wrap_buffer!(outside_info, TPM2B_DATA, 64);
230+
231+
if creation_pcrs.len() > 16 {
232+
return Err(Error::local_error(ErrorKind::WrongParamSize));
233+
}
234+
240235
let mut creation_pcrs_buffer = [Default::default(); 16];
241236
creation_pcrs_buffer[..creation_pcrs.len()]
242237
.clone_from_slice(&creation_pcrs[..creation_pcrs.len()]);
243238
let creation_pcrs = TPML_PCR_SELECTION {
244-
count: creation_pcrs.len().try_into().unwrap(),
239+
count: creation_pcrs.len().try_into().unwrap(), // will not fail given the len checks above
245240
pcrSelections: creation_pcrs_buffer,
246241
};
247242

@@ -297,18 +292,10 @@ impl Context {
297292
outside_info: &[u8],
298293
creation_pcrs: &[TPMS_PCR_SELECTION],
299294
) -> Result<(TPM2B_PRIVATE, TPM2B_PUBLIC)> {
300-
if auth_value.len() > 64
301-
|| initial_data.len() > 256
302-
|| outside_info.len() > 64
303-
|| creation_pcrs.len() > 16
304-
{
305-
return Err(Error::local_error(ErrorKind::WrongParamSize));
306-
}
307-
308295
let sensitive_create = TPM2B_SENSITIVE_CREATE {
309296
size: std::mem::size_of::<TPMS_SENSITIVE_CREATE>()
310297
.try_into()
311-
.unwrap(),
298+
.unwrap(), // will not fail on targets of at least 16 bits
312299
sensitive: TPMS_SENSITIVE_CREATE {
313300
userAuth: wrap_buffer!(auth_value, TPM2B_AUTH, 64),
314301
data: wrap_buffer!(initial_data, TPM2B_SENSITIVE_DATA, 256),
@@ -317,11 +304,14 @@ impl Context {
317304

318305
let outside_info = wrap_buffer!(outside_info, TPM2B_DATA, 64);
319306

307+
if creation_pcrs.len() > 16 {
308+
return Err(Error::local_error(ErrorKind::WrongParamSize));
309+
}
320310
let mut creation_pcrs_buffer = [Default::default(); 16];
321311
creation_pcrs_buffer[..creation_pcrs.len()]
322312
.clone_from_slice(&creation_pcrs[..creation_pcrs.len()]);
323313
let creation_pcrs = TPML_PCR_SELECTION {
324-
count: creation_pcrs.len().try_into().unwrap(),
314+
count: creation_pcrs.len().try_into().unwrap(), // will not fail given the len checks above
325315
pcrSelections: creation_pcrs_buffer,
326316
};
327317

@@ -403,9 +393,6 @@ impl Context {
403393
scheme: TPMT_SIG_SCHEME,
404394
validation: &TPMT_TK_HASHCHECK,
405395
) -> Result<Signature> {
406-
if digest.len() > 64 {
407-
return Err(Error::local_error(ErrorKind::WrongParamSize));
408-
}
409396
let mut signature = null_mut();
410397
let digest = wrap_buffer!(digest, TPM2B_DIGEST, 64);
411398
let ret = unsafe {
@@ -438,9 +425,6 @@ impl Context {
438425
digest: &[u8],
439426
signature: &TPMT_SIGNATURE,
440427
) -> Result<TPMT_TK_VERIFIED> {
441-
if digest.len() > 64 {
442-
return Err(Error::local_error(ErrorKind::WrongParamSize));
443-
}
444428
let mut validation = null_mut();
445429
let digest = wrap_buffer!(digest, TPM2B_DIGEST, 64);
446430
let ret = unsafe {
@@ -577,7 +561,7 @@ impl Context {
577561
let ret = Error::from_tss_rc(ret);
578562
if ret.is_success() {
579563
let context = unsafe { MBox::<TPMS_CONTEXT>::from_raw(context) };
580-
Ok((*context).into())
564+
Ok((*context).try_into()?)
581565
} else {
582566
error!("Error in saving context: {}.", ret);
583567
Err(ret)
@@ -612,7 +596,9 @@ impl Context {
612596
self.sessions.0,
613597
self.sessions.1,
614598
self.sessions.2,
615-
num_bytes.try_into().unwrap(),
599+
num_bytes
600+
.try_into()
601+
.or_else(|_| Err(Error::local_error(ErrorKind::WrongParamSize)))?,
616602
&mut buffer,
617603
)
618604
};
@@ -621,7 +607,7 @@ impl Context {
621607
if ret.is_success() {
622608
let buffer = unsafe { MBox::from_raw(buffer) };
623609
let mut random = buffer.buffer.to_vec();
624-
random.truncate(buffer.size.try_into().unwrap());
610+
random.truncate(buffer.size.try_into().unwrap()); // should not panic given the TryInto above
625611
Ok(random)
626612
} else {
627613
error!("Error in flushing context: {}.", ret);
@@ -630,10 +616,6 @@ impl Context {
630616
}
631617

632618
pub fn set_handle_auth(&mut self, handle: ESYS_TR, auth_value: &[u8]) -> Result<()> {
633-
if auth_value.len() > 64 {
634-
return Err(Error::local_error(ErrorKind::WrongParamSize));
635-
}
636-
637619
let auth = wrap_buffer!(auth_value, TPM2B_AUTH, 64);
638620
let ret = unsafe { Esys_TR_SetAuth(self.mut_context(), handle, &auth) };
639621
let ret = Error::from_tss_rc(ret);
@@ -657,7 +639,7 @@ impl Context {
657639
}
658640

659641
fn mut_context(&mut self) -> *mut ESYS_CONTEXT {
660-
self.esys_context.as_mut().unwrap().as_mut_ptr()
642+
self.esys_context.as_mut().unwrap().as_mut_ptr() // will only fail if called from Drop after .take()
661643
}
662644
}
663645

@@ -673,8 +655,8 @@ impl Drop for Context {
673655
}
674656
});
675657

676-
let esys_context = self.esys_context.take().unwrap();
677-
let tcti_context = self.tcti_context.take().unwrap();
658+
let esys_context = self.esys_context.take().unwrap(); // should not fail based on how the context is initialised/used
659+
let tcti_context = self.tcti_context.take().unwrap(); // should not fail based on how the context is initialised/used
678660

679661
// Close the TCTI context.
680662
unsafe {

‎src/response_code.rs

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -328,7 +328,7 @@ impl std::fmt::Display for Tss2ResponseCode {
328328
if kind.is_none() {
329329
return write!(f, "response code not recognized");
330330
}
331-
match self.kind().unwrap() {
331+
match self.kind().unwrap() { // should not panic, given the check above
332332
Tss2ResponseCodeKind::Success => write!(f, "success"),
333333
Tss2ResponseCodeKind::TpmVendorSpecific => write!(f, "vendor specific error: {}", self.error_number()),
334334
// Format Zero
@@ -475,6 +475,8 @@ impl std::fmt::Display for Error {
475475
#[derive(Copy, Clone, PartialEq, Debug)]
476476
pub enum WrapperErrorKind {
477477
WrongParamSize,
478+
ParamsMissing,
479+
InconsistentParams,
478480
}
479481

480482
impl std::fmt::Display for WrapperErrorKind {
@@ -483,6 +485,13 @@ impl std::fmt::Display for WrapperErrorKind {
483485
WrapperErrorKind::WrongParamSize => {
484486
write!(f, "parameter provided is of the wrong size")
485487
}
488+
WrapperErrorKind::ParamsMissing => {
489+
write!(f, "some of the required parameters were not provided")
490+
}
491+
WrapperErrorKind::InconsistentParams => write!(
492+
f,
493+
"the provided parameters have inconsistent values or variants"
494+
),
486495
}
487496
}
488497
}

0 commit comments

Comments
 (0)