Merge pull request #15 from tarcieri/apdu-cleanups

Clean up APDU construction with builder API

src/apdu.rs

@@ -1,7 +1,13 @@
 //! Application Protocol Data Unit (APDU)
 
 use std::fmt::{self, Debug};
-use zeroize::Zeroize;
+use zeroize::{Zeroize, Zeroizing};
+
+/// Size of a serialized APDU (5 byte header + 255 bytes data)
+pub const APDU_SIZE: usize = 260;
+
+/// Buffer type (self-zeroizing byte vector)
+pub(crate) type Buffer = Zeroizing<Vec<u8>>;
 
 /// Application Protocol Data Unit (APDU).
 ///
@@ -10,30 +16,83 @@ use zeroize::Zeroize;
 #[derive(Clone)]
 pub struct APDU {
     /// Instruction class - indicates the type of command, e.g. interindustry or proprietary
-    pub cla: u8,
+    cla: u8,
 
     /// Instruction code - indicates the specific command, e.g. "write data"
-    pub ins: u8,
+    ins: u8,
 
     /// Instruction parameter 1 for the command, e.g. offset into file at which to write the data
-    pub p1: u8,
+    p1: u8,
 
     /// Instruction parameter 2 for the command
-    pub p2: u8,
+    p2: u8,
 
     /// Length of command - encodes the number of bytes of command data to follow
-    pub lc: u8,
+    lc: u8,
 
     /// Command data
-    pub data: [u8; 255],
+    data: [u8; 255],
 }
 
 impl APDU {
-    /// Get a mut pointer to this APDU
-    // TODO(tarcieri): eliminate pointers and use all safe references
-    pub(crate) fn as_mut_ptr(&mut self) -> *mut APDU {
+    /// Create a new APDU with the given instruction code
+    pub fn new(ins: u8) -> Self {
+        let mut apdu = Self::default();
+        apdu.ins = ins;
+        apdu
+    }
+
+    /// Set this APDU's class
+    pub fn cla(&mut self, value: u8) -> &mut Self {
+        self.cla = value;
+        self
+    }
+
+    /// Set this APDU's first parameter only
+    pub fn p1(&mut self, value: u8) -> &mut Self {
+        self.p1 = value;
+        self
+    }
+
+    /// Set both parameters for this APDU
+    pub fn params(&mut self, p1: u8, p2: u8) -> &mut Self {
+        self.p1 = p1;
+        self.p2 = p2;
         self
     }
+
+    /// Set the command data for this APDU.
+    ///
+    /// Panics if the byte slice is more than 255 bytes!
+    pub fn data(&mut self, bytes: impl AsRef<[u8]>) -> &mut Self {
+        assert_eq!(self.lc, 0, "APDU command already set!");
+
+        let bytes = bytes.as_ref();
+
+        assert!(
+            bytes.len() <= self.data.len(),
+            "APDU command data too large: {}-bytes",
+            bytes.len()
+        );
+
+        self.lc = bytes.len() as u8;
+        self.data[..bytes.len()].copy_from_slice(bytes);
+
+        self
+    }
+
+    /// Consume this APDU and return a self-zeroizing buffer
+    pub fn to_bytes(&self) -> Buffer {
+        let mut bytes = Vec::with_capacity(APDU_SIZE);
+        bytes.push(self.cla);
+        bytes.push(self.ins);
+        bytes.push(self.p1);
+        bytes.push(self.p2);
+        bytes.push(self.lc);
+        bytes.extend_from_slice(self.data.as_ref());
+
+        Zeroizing::new(bytes)
+    }
 }
 
 impl Debug for APDU {
@@ -64,6 +123,12 @@ impl Default for APDU {
     }
 }
 
+impl Drop for APDU {
+    fn drop(&mut self) {
+        self.zeroize();
+    }
+}
+
 impl Zeroize for APDU {
     fn zeroize(&mut self) {
         self.cla.zeroize();

src/util.rs

@@ -1607,12 +1607,13 @@ pub unsafe fn ykpiv_util_set_protected_mgm(
         if ykrc.is_err() {
             // if set_mgmkey fails with KeyError, it means the generated key is weak
             // otherwise, log a warning, since the device mgm key is corrupt or we're in
-            // a yubikey where we can't set the mgm key
+            // a state where we can't set the mgm key
             if Err(Error::KeyError) != ykrc {
                 error!(
                     "could not set new derived mgm key, err = {}",
                     ykrc.as_ref().unwrap_err()
                 );
+
                 let _ = yubikey._ykpiv_end_transaction();
                 return ykrc;
             }