Extract TLV writing into serialization::Tlv

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 */;
+    let mut buf = [0u8; CB_OBJ_MAX];
-    req_len += set_length(&mut buf, data.len());
+    let mut offset = Tlv::write(&mut buf, TAG_CERT, data)?;
-    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();
 
     // write compression info and LRC trailer
-    buf[offset] = TAG_CERT_COMPRESS;
+    offset += Tlv::write(
-    buf[offset + 1] = 0x01;
+        &mut buf,
-    buf[offset + 2] = if certinfo == CERTINFO_GZIP {
+        TAG_CERT_COMPRESS,
-        0x01
+        if certinfo == CERTINFO_GZIP {
-    } else {
+            &[0x01]
-        0x00
+        } else {
-    };
+            &[0x00]
-    buf[offset + 3] = TAG_CERT_LRC;
+        },
-    buf[offset + 4] = 00;
+    )?;
-
+    offset += Tlv::write(&mut buf, TAG_CERT_LRC, &[])?;
-    offset += 5;
 
     txn.save_object(object_id, &buf[..offset])
 }

src/key.rs

@@ -379,11 +379,8 @@ impl From<AlgorithmId> 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 {
+    pub(crate) fn write(self, buf: &mut [u8]) -> Result<usize, Error> {
-        buf[0] = 0x80;
+        Tlv::write(buf, 0x80, &[self.into()])
-        buf[1] = 0x01;
-        buf[2] = self.into();
-        3
     }
 }
 
@@ -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;
+    let mut offset = Tlv::write_as(&mut in_data, 0xac, 3, |buf| {
-    in_data[1] = 3; // length sans this 2-byte header
+        assert_eq!(algorithm.write(buf).expect("large enough"), 3);
-    assert_eq!(algorithm.write(&mut in_data[2..]), 3);
+    })?;
-    let mut offset = 5;
 
-    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 += Tlv::write_as(
-        offset += 1;
+            &mut key_data[offset..],
-
+            param_tag + i as u8,
-        offset += set_length(&mut key_data[offset..], elem_len);
+            elem_len,
-
+            |buf| {
-        let padding = elem_len - param.len();
+                let padding = elem_len - param.len();
-        let remaining = key_data.len() - offset;
+                for b in &mut buf[..padding] {
-
+                    *b = 0;
-        if padding > remaining {
+                }
-            return Err(Error::AlgorithmError);
+                buf[padding..].copy_from_slice(param);
-        }
+            },
-
+        )?;
-        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 += pin_policy.write(&mut key_data[offset..]);
+    offset += pin_policy.write(&mut key_data[offset..])?;
-    offset += touch_policy.write(&mut key_data[offset..]);
+    offset += touch_policy.write(&mut key_data[offset..])?;
 
     let txn = yubikey.begin_transaction()?;
 

src/metadata.rs

@@ -95,20 +95,7 @@ pub(crate) fn set_item(
         }
 
         // We did not find an existing tag, append
-        offset = *pcb_data;
+        *pcb_data += Tlv::write(&mut data[*pcb_data..], tag, p_item)?;
-        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;
 
         return Ok(());
     }
@@ -171,8 +158,6 @@ pub(crate) fn read(txn: &Transaction<'_>, tag: u8) -> Result<Buffer, Error> {
 /// 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 buf = Zeroizing::new(vec![0u8; CB_OBJ_MAX]);
-    let mut offset = set_length(&mut buf[1..], data.len());
+    let len = Tlv::write(&mut buf, tag, data)?;
-    buf[offset..(offset + data.len())].copy_from_slice(data);
-    offset += data.len();
 
-    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

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])
     }
 

src/msroots.rs

@@ -130,16 +130,15 @@ impl MsRoots {
                 data_len - data_offset
             };
 
-            buf[offset] = if i == n_objs - 1 {
+            offset += Tlv::write(
-                TAG_MSROOTS_END
+                &mut buf,
-            } else {
+                if i == n_objs - 1 {
-                TAG_MSROOTS_MID
+                    TAG_MSROOTS_END
-            };
+                } else {
-
+                    TAG_MSROOTS_MID
-            offset += 1;
+                },
-            offset += set_length(&mut buf[offset..], data_chunk);
+                &data[data_offset..(data_offset + data_chunk)],
-            buf[offset..].copy_from_slice(&data[data_offset..(data_offset + data_chunk)]);
+            )?;
-            offset += data_chunk;
 
             txn.save_object(OBJ_MSROOTS1 + i as u32, &buf[..offset])?;
 

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<usize, Error> {
         match self {
-            PinPolicy::Default => 0,
+            PinPolicy::Default => Ok(0),
-            _ => {
+            _ => Tlv::write(buf, 0xaa, &[self.into()]),
-                buf[0] = 0xaa;
-                buf[1] = 0x01;
-                buf[2] = self.into();
-                3
-            }
         }
     }
 }
@@ -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<usize, Error> {
         match self {
-            TouchPolicy::Default => 0,
+            TouchPolicy::Default => Ok(0),
-            _ => {
+            _ => Tlv::write(buf, 0xab, &[self.into()]),
-                buf[0] = 0xab;
-                buf[1] = 0x01;
-                buf[2] = self.into();
-                3
-            }
         }
     }
 }

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<usize, Error> {
+        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<Gen>(
+        buffer: &mut [u8],
+        tag: u8,
+        length: usize,
+        value: Gen,
+    ) -> Result<usize, Error>
+    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

src/transaction.rs

@@ -299,19 +299,21 @@ impl<'tx> Transaction<'tx> {
             3
         };
 
-        indata[0] = 0x7c;
+        let offset = Tlv::write_as(&mut indata, 0x7c, in_len + bytes + 3, |buf| {
-        let mut offset = 1 + set_length(&mut indata[1..], in_len + bytes + 3);
+            assert_eq!(Tlv::write(buf, 0x82, &[]).expect("large enough"), 2);
-        indata[offset] = 0x82;
+            assert_eq!(
-        indata[offset + 1] = 0x00;
+                Tlv::write(
-        indata[offset + 2] = match (algorithm, decipher) {
+                    &mut buf[2..],
-            (AlgorithmId::EccP256, true) | (AlgorithmId::EccP384, true) => 0x85,
+                    match (algorithm, decipher) {
-            _ => 0x81,
+                        (AlgorithmId::EccP256, true) | (AlgorithmId::EccP384, true) => 0x85,
-        };
+                        _ => 0x81,
-
+                    },
-        offset += 3;
+                    sign_in
-        offset += set_length(&mut indata[offset..], in_len);
+                )
-        indata[offset..(offset + in_len)].copy_from_slice(sign_in);
+                .expect("large enough"),
-        offset += in_len;
+                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;
+        let offset = Tlv::write(data_remaining, 0x53, indata)?;
-        data_remaining = &mut data_remaining[1..];
-
-        let offset = set_length(data_remaining, indata.len());
         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