diff --git a/src/certificate.rs b/src/certificate.rs index e920b6a..3f70452 100644 --- a/src/certificate.rs +++ b/src/certificate.rs @@ -246,9 +246,6 @@ pub(crate) fn write_certificate( data: Option<&[u8]>, certinfo: u8, ) -> Result<(), Error> { - let mut buf = [0u8; CB_OBJ_MAX]; - let mut offset = 0; - let object_id = slot.object_id(); if data.is_none() { @@ -257,34 +254,20 @@ pub(crate) fn write_certificate( let data = data.unwrap(); - let mut req_len = 1 /* cert tag */ + 3 /* compression tag + data*/ + 2 /* lrc */; - req_len += set_length(&mut buf, data.len()); - req_len += data.len(); - - if req_len < data.len() || req_len > CB_OBJ_MAX { - return Err(Error::SizeError); - } - - buf[offset] = TAG_CERT; - offset += 1; - offset += set_length(&mut buf[offset..], data.len()); - - buf[offset..(offset + data.len())].copy_from_slice(&data); - - offset += data.len(); + let mut buf = [0u8; CB_OBJ_MAX]; + let mut offset = Tlv::write(&mut buf, TAG_CERT, data)?; // write compression info and LRC trailer - buf[offset] = TAG_CERT_COMPRESS; - buf[offset + 1] = 0x01; - buf[offset + 2] = if certinfo == CERTINFO_GZIP { - 0x01 - } else { - 0x00 - }; - buf[offset + 3] = TAG_CERT_LRC; - buf[offset + 4] = 00; - - offset += 5; + offset += Tlv::write( + &mut buf, + TAG_CERT_COMPRESS, + if certinfo == CERTINFO_GZIP { + &[0x01] + } else { + &[0x00] + }, + )?; + offset += Tlv::write(&mut buf, TAG_CERT_LRC, &[])?; txn.save_object(object_id, &buf[..offset]) } diff --git a/src/key.rs b/src/key.rs index 8eaf67b..84826d2 100644 --- a/src/key.rs +++ b/src/key.rs @@ -379,11 +379,8 @@ impl From for u8 { #[cfg(feature = "untested")] impl AlgorithmId { /// Writes the `AlgorithmId` in the format the YubiKey expects during key generation. - pub(crate) fn write(self, buf: &mut [u8]) -> usize { - buf[0] = 0x80; - buf[1] = 0x01; - buf[2] = self.into(); - 3 + pub(crate) fn write(self, buf: &mut [u8]) -> Result { + Tlv::write(buf, 0x80, &[self.into()]) } } @@ -539,16 +536,15 @@ pub fn generate( let templ = [0, Ins::GenerateAsymmetric.code(), 0, slot.into()]; let mut in_data = [0u8; 11]; - in_data[0] = 0xac; - in_data[1] = 3; // length sans this 2-byte header - assert_eq!(algorithm.write(&mut in_data[2..]), 3); - let mut offset = 5; + let mut offset = Tlv::write_as(&mut in_data, 0xac, 3, |buf| { + assert_eq!(algorithm.write(buf).expect("large enough"), 3); + })?; - let pin_len = pin_policy.write(&mut in_data[offset..]); + let pin_len = pin_policy.write(&mut in_data[offset..])?; in_data[1] += pin_len as u8; offset += pin_len; - let touch_len = touch_policy.write(&mut in_data[offset..]); + let touch_len = touch_policy.write(&mut in_data[offset..])?; in_data[1] += touch_len as u8; offset += touch_len; @@ -695,28 +691,22 @@ pub fn import( let mut offset = 0; for (i, param) in params.into_iter().enumerate() { - key_data[offset] = param_tag + i as u8; - offset += 1; - - offset += set_length(&mut key_data[offset..], elem_len); - - let padding = elem_len - param.len(); - let remaining = key_data.len() - offset; - - if padding > remaining { - return Err(Error::AlgorithmError); - } - - for b in &mut key_data[offset..offset + padding] { - *b = 0; - } - offset += padding; - key_data[offset..offset + param.len()].copy_from_slice(param); - offset += param.len(); + offset += Tlv::write_as( + &mut key_data[offset..], + param_tag + i as u8, + elem_len, + |buf| { + let padding = elem_len - param.len(); + for b in &mut buf[..padding] { + *b = 0; + } + buf[padding..].copy_from_slice(param); + }, + )?; } - offset += pin_policy.write(&mut key_data[offset..]); - offset += touch_policy.write(&mut key_data[offset..]); + offset += pin_policy.write(&mut key_data[offset..])?; + offset += touch_policy.write(&mut key_data[offset..])?; let txn = yubikey.begin_transaction()?; diff --git a/src/metadata.rs b/src/metadata.rs index 641ac08..e2cc3ef 100644 --- a/src/metadata.rs +++ b/src/metadata.rs @@ -95,20 +95,7 @@ pub(crate) fn set_item( } // We did not find an existing tag, append - offset = *pcb_data; - cb_len = get_length_size(cb_item); - - // If length would cause buffer overflow, return error - if (*pcb_data + cb_len + cb_item) > cb_data_max { - return Err(Error::GenericError); - } - - data[offset] = tag; - offset += 1; - offset += set_length(&mut data[offset..], cb_item); - data[offset..offset + cb_item].copy_from_slice(p_item); - - *pcb_data += 1 + cb_len + cb_item; + *pcb_data += Tlv::write(&mut data[*pcb_data..], tag, p_item)?; return Ok(()); } @@ -171,8 +158,6 @@ pub(crate) fn read(txn: &Transaction<'_>, tag: u8) -> Result { /// Write metadata #[cfg(feature = "untested")] pub(crate) fn write(txn: &Transaction<'_>, tag: u8, data: &[u8]) -> Result<(), Error> { - let mut buf = Zeroizing::new(vec![0u8; CB_OBJ_MAX]); - if data.len() > CB_OBJ_MAX - CB_OBJ_TAG_MAX { return Err(Error::GenericError); } @@ -188,12 +173,10 @@ pub(crate) fn write(txn: &Transaction<'_>, tag: u8, data: &[u8]) -> Result<(), E return txn.save_object(obj_id, &[]); } - buf[0] = tag; - let mut offset = set_length(&mut buf[1..], data.len()); - buf[offset..(offset + data.len())].copy_from_slice(data); - offset += data.len(); + let mut buf = Zeroizing::new(vec![0u8; CB_OBJ_MAX]); + let len = Tlv::write(&mut buf, tag, data)?; - txn.save_object(obj_id, &buf[..offset]) + txn.save_object(obj_id, &buf[..len]) } /// Get the size of a length tag for the given length diff --git a/src/mscmap.rs b/src/mscmap.rs index dc58c90..e8ec870 100644 --- a/src/mscmap.rs +++ b/src/mscmap.rs @@ -107,8 +107,6 @@ impl Container { /// Write MS Container Map records. pub fn write_mscmap(yubikey: &mut YubiKey, containers: &[Self]) -> Result<(), Error> { - let mut buf = [0u8; CB_OBJ_MAX]; - let mut offset = 0; let n_containers = containers.len(); let data_len = n_containers * CONTAINER_REC_LEN; @@ -118,24 +116,13 @@ impl Container { return txn.save_object(OBJ_MSCMAP, &[]); } - let req_len = 1 + set_length(&mut buf, data_len) + data_len; + let mut buf = [0u8; CB_OBJ_MAX]; + let offset = Tlv::write_as(&mut buf, TAG_MSCMAP, data_len, |buf| { + for (i, chunk) in buf.chunks_exact_mut(CONTAINER_REC_LEN).enumerate() { + chunk.copy_from_slice(&containers[i].to_bytes()); + } + })?; - if req_len > CB_OBJ_MAX { - return Err(Error::SizeError); - } - - buf[offset] = TAG_MSCMAP; - offset += 1; - offset += set_length(&mut buf[offset..], data_len); - - for (i, chunk) in buf[..data_len] - .chunks_exact_mut(CONTAINER_REC_LEN) - .enumerate() - { - chunk.copy_from_slice(&containers[i].to_bytes()); - } - - offset += data_len; txn.save_object(OBJ_MSCMAP, &buf[..offset]) } diff --git a/src/msroots.rs b/src/msroots.rs index 39c04e5..b7d65b7 100644 --- a/src/msroots.rs +++ b/src/msroots.rs @@ -130,16 +130,15 @@ impl MsRoots { data_len - data_offset }; - buf[offset] = if i == n_objs - 1 { - TAG_MSROOTS_END - } else { - TAG_MSROOTS_MID - }; - - offset += 1; - offset += set_length(&mut buf[offset..], data_chunk); - buf[offset..].copy_from_slice(&data[data_offset..(data_offset + data_chunk)]); - offset += data_chunk; + offset += Tlv::write( + &mut buf, + if i == n_objs - 1 { + TAG_MSROOTS_END + } else { + TAG_MSROOTS_MID + }, + &data[data_offset..(data_offset + data_chunk)], + )?; txn.save_object(OBJ_MSROOTS1 + i as u32, &buf[..offset])?; diff --git a/src/policy.rs b/src/policy.rs index e82c180..a5b2c8d 100644 --- a/src/policy.rs +++ b/src/policy.rs @@ -1,5 +1,7 @@ //! Enums representing key policies. +use crate::{error::Error, serialization::Tlv}; + /// Specifies how often the PIN needs to be entered for access to the credential in a /// given slot. This policy must be set upon key generation or importation, and cannot be /// changed later. @@ -36,15 +38,10 @@ impl PinPolicy { /// Writes the `PinPolicy` in the format the YubiKey expects during key generation or /// importation. #[cfg(feature = "untested")] - pub(crate) fn write(self, buf: &mut [u8]) -> usize { + pub(crate) fn write(self, buf: &mut [u8]) -> Result { match self { - PinPolicy::Default => 0, - _ => { - buf[0] = 0xaa; - buf[1] = 0x01; - buf[2] = self.into(); - 3 - } + PinPolicy::Default => Ok(0), + _ => Tlv::write(buf, 0xaa, &[self.into()]), } } } @@ -86,15 +83,10 @@ impl TouchPolicy { /// Writes the `TouchPolicy` in the format the YubiKey expects during key generation /// or importation. #[cfg(feature = "untested")] - pub(crate) fn write(self, buf: &mut [u8]) -> usize { + pub(crate) fn write(self, buf: &mut [u8]) -> Result { match self { - TouchPolicy::Default => 0, - _ => { - buf[0] = 0xab; - buf[1] = 0x01; - buf[2] = self.into(); - 3 - } + TouchPolicy::Default => Ok(0), + _ => Tlv::write(buf, 0xab, &[self.into()]), } } } diff --git a/src/serialization.rs b/src/serialization.rs index 794a4e3..55fc699 100644 --- a/src/serialization.rs +++ b/src/serialization.rs @@ -77,6 +77,52 @@ impl<'a> Tlv<'a> { buffer.truncate(len); Ok(buffer) } + + /// Writes a TLV to the given buffer. + pub(crate) fn write(buffer: &mut [u8], tag: u8, value: &[u8]) -> Result { + if buffer.len() < CB_OBJ_TAG_MIN { + return Err(Error::SizeError); + } + buffer[0] = tag; + + // TODO: Raise error + let offset = 1 + set_length(&mut buffer[1..], value.len()); + + if buffer.len() < offset + value.len() { + return Err(Error::SizeError); + } + buffer[offset..offset + value.len()].copy_from_slice(value); + + Ok(offset + value.len()) + } + + /// Writes a TLV to the given buffer. + /// + /// `value` is guaranteed to be called with a mutable slice of length `length`. + pub(crate) fn write_as( + buffer: &mut [u8], + tag: u8, + length: usize, + value: Gen, + ) -> Result + where + Gen: FnOnce(&mut [u8]), + { + if buffer.len() < CB_OBJ_TAG_MIN { + return Err(Error::SizeError); + } + buffer[0] = tag; + + // TODO: Raise error + let offset = 1 + set_length(&mut buffer[1..], length); + + if buffer.len() < offset + length { + return Err(Error::SizeError); + } + value(&mut buffer[offset..offset + length]); + + Ok(offset + length) + } } /// Set length diff --git a/src/transaction.rs b/src/transaction.rs index fa465aa..d2dd9d9 100644 --- a/src/transaction.rs +++ b/src/transaction.rs @@ -299,19 +299,21 @@ impl<'tx> Transaction<'tx> { 3 }; - indata[0] = 0x7c; - let mut offset = 1 + set_length(&mut indata[1..], in_len + bytes + 3); - indata[offset] = 0x82; - indata[offset + 1] = 0x00; - indata[offset + 2] = match (algorithm, decipher) { - (AlgorithmId::EccP256, true) | (AlgorithmId::EccP384, true) => 0x85, - _ => 0x81, - }; - - offset += 3; - offset += set_length(&mut indata[offset..], in_len); - indata[offset..(offset + in_len)].copy_from_slice(sign_in); - offset += in_len; + let offset = Tlv::write_as(&mut indata, 0x7c, in_len + bytes + 3, |buf| { + assert_eq!(Tlv::write(buf, 0x82, &[]).expect("large enough"), 2); + assert_eq!( + Tlv::write( + &mut buf[2..], + match (algorithm, decipher) { + (AlgorithmId::EccP256, true) | (AlgorithmId::EccP384, true) => 0x85, + _ => 0x81, + }, + sign_in + ) + .expect("large enough"), + 1 + bytes + in_len + ); + })?; let response = self .transfer_data(&templ, &indata[..offset], 1024) @@ -484,14 +486,8 @@ impl<'tx> Transaction<'tx> { let mut len = data.len(); let mut data_remaining = set_object(object_id, &mut data); - data_remaining[0] = 0x53; - data_remaining = &mut data_remaining[1..]; - - let offset = set_length(data_remaining, indata.len()); + let offset = Tlv::write(data_remaining, 0x53, indata)?; data_remaining = &mut data_remaining[offset..]; - data_remaining[..indata.len()].copy_from_slice(indata); - - data_remaining = &mut data_remaining[indata.len()..]; len -= data_remaining.len(); let status_words = self