openpgp: Correctly handle padding on Features, add rountrip test.

1 files changed, 62 insertions, 4 deletions

diff --git a/openpgp/src/types/features.rs b/openpgp/src/types/features.rs
index cce547f6..ca4ac999 100644
--- a/openpgp/src/types/features.rs
+++ b/openpgp/src/types/features.rs
@@ -1,9 +1,13 @@
+use std::hash::{Hash, Hasher};
+
 /// Describes features supported by an OpenPGP implementation.
-#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+#[derive(Clone, Debug)]
 pub struct Features{
     mdc: bool,
     aead: bool,
     unknown: Box<[u8]>,
+    /// Original length, including trailing zeros.
+    pad_to: usize,
 }
 
 impl Default for Features {
@@ -12,13 +16,35 @@ impl Default for Features {
             mdc: false,
             aead: false,
             unknown: Default::default(),
+            pad_to: 0,
         }
     }
 }
 
+impl PartialEq for Features {
+    fn eq(&self, other: &Self) -> bool {
+        self.mdc == other.mdc
+            && self.aead == other.aead
+            && self.unknown == other.unknown
+    }
+}
+
+impl Eq for Features {}
+
+impl Hash for Features {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.mdc.hash(state);
+        self.aead.hash(state);
+        self.unknown.hash(state);
+    }
+}
+
 impl Features {
     /// Creates a new instance from `bits`.
-    pub fn new(bits: &[u8]) -> Self {
+    pub fn new<B: AsRef<[u8]>>(bits: B) -> Self {
+        let bits = bits.as_ref();
+        let mut pad_to = 0;
+
         let mdc = bits.get(0)
             .map(|x| x & FEATURE_FLAG_MDC != 0).unwrap_or(false);
         let aead = bits.get(0)
@@ -30,12 +56,12 @@ impl Features {
 
             cpy[0] &= (FEATURE_FLAG_MDC | FEATURE_FLAG_AEAD) ^ 0xff;
 
-            while cpy.last().cloned() == Some(0) { cpy.pop(); }
+            pad_to = crate::types::bitfield_remove_padding(&mut cpy);
             cpy.into_boxed_slice()
         };
 
         Features{
-            mdc: mdc, aead: aead, unknown: unk
+            mdc: mdc, aead: aead, unknown: unk, pad_to,
         }
     }
 
@@ -45,6 +71,7 @@ impl Features {
             mdc: true,
             aead: false,
             unknown: Default::default(),
+            pad_to: 0,
         }
     }
 
@@ -59,6 +86,18 @@ impl Features {
         if self.mdc { ret[0] |= FEATURE_FLAG_MDC; }
         if self.aead { ret[0] |= FEATURE_FLAG_AEAD; }
 
+        // Corner case: empty flag field.  We initialized ret to
+        // vec![0] for easy setting of flags.  See if any of the above
+        // was set.
+        if ret.len() == 1 && ret[0] == 0 {
+            // Nope.  Trim this byte.
+            ret.pop();
+        }
+
+        for _ in ret.len()..self.pad_to {
+            ret.push(0);
+        }
+
         ret
     }
 
@@ -91,3 +130,22 @@ const FEATURE_FLAG_MDC: u8 = 0x01;
 /// AEAD Encrypted Data Packet (packet 20) and version 5 Symmetric-Key
 /// Encrypted Session Key Packets (packet 3).
 const FEATURE_FLAG_AEAD: u8 = 0x02;
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    quickcheck! {
+        fn roundtrip(raw: Vec<u8>) -> bool {
+            let val = Features::new(&raw);
+            assert_eq!(raw, val.to_vec());
+
+            // Check that equality ignores padding.
+            let mut val_without_padding = val.clone();
+            val_without_padding.pad_to = val.unknown.len();
+            assert_eq!(val, val_without_padding);
+
+            true
+        }
+    }
+}