openpgp: Improve Cert::insert_packets.neal/smarter-insert-packets

  - The current implementation of `Cert::insert_packets` is O(n*m)
    where n is the number of signatures in the certificate and m is
    the number of signatures in the packets to insert.

  - Be smarter.  Put the signatures in a BTreeMap and search it
    instead of doing a linear scan to see if we already have a
    given packet.

4 files changed, 351 insertions, 63 deletions

diff --git a/openpgp/src/cert.rs b/openpgp/src/cert.rs
index 162b63e0..39f55d11 100644
--- a/openpgp/src/cert.rs
+++ b/openpgp/src/cert.rs
@@ -137,10 +137,14 @@
 //! [`Signature::verify_user_attribute_revocation`]: ../packet/enum.Signature.html#method.verify_user_attribute_revocation
 
 use std::io;
+use std::collections::btree_map::BTreeMap;
+use std::collections::btree_map::Entry;
+use std::collections::hash_map::DefaultHasher;
 use std::cmp;
 use std::cmp::Ordering;
 use std::convert::TryFrom;
 use std::convert::TryInto;
+use std::hash::Hasher;
 use std::path::Path;
 use std::mem;
 use std::fmt;
@@ -2353,78 +2357,127 @@ impl Cert {
     {
         let mut combined = self.into_packets().collect::<Vec<_>>();
 
-        fn replace_or_push(acc: &mut Vec<Packet>, p: Packet)
-        {
-            match p {
-                Packet::PublicKey(_) => (),
-                Packet::SecretKey(_) => (),
-                Packet::PublicSubkey(_) => (),
-                Packet::SecretSubkey(_) => (),
-                _ => unreachable!(),
-            }
+        // Hashes a packet ignoring the unhashed subpacket area and
+        // any secret key material.
+        let hash_packet = |p: &Packet| -> u64 {
+            let mut hasher = DefaultHasher::new();
+            p.normalized_hash(&mut hasher);
+            hasher.finish()
+        };
 
-            for q in acc.iter_mut() {
-                let replace = match (&p, &q) {
-                    (Packet::PublicKey(a), Packet::PublicKey(b)) =>
-                        a.public_cmp(&b) == Ordering::Equal,
-                    (Packet::SecretKey(a), Packet::SecretKey(b)) =>
-                        a.public_cmp(&b) == Ordering::Equal,
-                    (Packet::PublicKey(a), Packet::SecretKey(b)) =>
-                        a.public_cmp(&b) == Ordering::Equal,
-                    (Packet::SecretKey(a), Packet::PublicKey(b)) =>
-                        a.public_cmp(&b) == Ordering::Equal,
-
-                    (Packet::PublicSubkey(a), Packet::PublicSubkey(b)) =>
-                        a.public_cmp(&b) == Ordering::Equal,
-                    (Packet::SecretSubkey(a), Packet::SecretSubkey(b)) =>
-                        a.public_cmp(&b) == Ordering::Equal,
-                    (Packet::PublicSubkey(a), Packet::SecretSubkey(b)) =>
-                        a.public_cmp(&b) == Ordering::Equal,
-                    (Packet::SecretSubkey(a), Packet::PublicSubkey(b)) =>
-                        a.public_cmp(&b) == Ordering::Equal,
-
-                    _ => false,
-                };
-
-                if replace {
-                    *q = p;
-                    return;
+        // BTreeMap of (hash) -> Vec<index in combined>.
+        //
+        // We don't use a HashMap, because the key would be a
+        // reference to the packets in combined, which would prevent
+        // us from modifying combined.
+        //
+        // Note: we really don't want to dedup components now, because
+        // we want to keep signatures immediately after their
+        // components.
+        let mut packet_map: BTreeMap<u64, Vec<usize>> = BTreeMap::new();
+        for (i, p) in combined.iter().enumerate() {
+            match packet_map.entry(hash_packet(p)) {
+                Entry::Occupied(mut oe) => {
+                    oe.get_mut().push(i)
+                }
+                Entry::Vacant(ve) => {
+                    ve.insert(vec![ i ]);
                 }
             }
-            acc.push(p);
         }
 
-        /// Replaces or pushes a signature.
-        ///
-        /// If `sig` is equal to an existing signature modulo unhashed
-        /// subpacket area, replaces the existing signature with it.
-        /// Otherwise, `sig` is pushed to the vector.
-        fn rop_sig(acc: &mut Vec<Packet>, sig: Signature)
-        {
-            for q in acc.iter_mut() {
-                let replace = match q {
-                    Packet::Signature(s) => s.normalized_eq(&sig),
-                    _ => false,
-                };
-
-                if replace {
-                    *q = sig.into();
-                    return;
-                }
-            }
-            acc.push(sig.into());
+        enum Action {
+            Drop,
+            Overwrite(usize),
+            Insert,
         }
+        use Action::*;
 
+        // Now we merge in the new packets.
         for p in packets {
             let p = p.into();
             Cert::valid_packet(&p)?;
-            match p {
-                p @ Packet::PublicKey(_) => replace_or_push(&mut combined, p),
-                p @ Packet::SecretKey(_) => replace_or_push(&mut combined, p),
-                p @ Packet::PublicSubkey(_) => replace_or_push(&mut combined, p),
-                p @ Packet::SecretSubkey(_) => replace_or_push(&mut combined, p),
-                Packet::Signature(sig) => rop_sig(&mut combined, sig),
-                p => combined.push(p),
+
+            let hash = hash_packet(&p);
+            let mut action = Insert;
+            if let Some(combined_i) = packet_map.get(&hash) {
+                for i in combined_i {
+                    let i: usize = *i;
+                    let (same, identical) = match (&p, &combined[i]) {
+                        // For keys, only compare the public bits.  If
+                        // they match, then we keep whatever is in the
+                        // new key.
+                        (Packet::PublicKey(a), Packet::PublicKey(b)) =>
+                            (a.public_cmp(&b) == Ordering::Equal,
+                             a == b),
+                        (Packet::SecretKey(a), Packet::SecretKey(b)) =>
+                            (a.public_cmp(&b) == Ordering::Equal,
+                             a == b),
+                        (Packet::PublicKey(a), Packet::SecretKey(b)) =>
+                            (a.public_cmp(&b) == Ordering::Equal,
+                             false),
+                        (Packet::SecretKey(a), Packet::PublicKey(b)) =>
+                            (a.public_cmp(&b) == Ordering::Equal,
+                             false),
+
+                        (Packet::PublicSubkey(a), Packet::PublicSubkey(b)) =>
+                            (a.public_cmp(&b) == Ordering::Equal,
+                             a == b),
+                        (Packet::SecretSubkey(a), Packet::SecretSubkey(b)) =>
+                            (a.public_cmp(&b) == Ordering::Equal,
+                             a == b),
+                        (Packet::PublicSubkey(a), Packet::SecretSubkey(b)) =>
+                            (a.public_cmp(&b) == Ordering::Equal,
+                             false),
+                        (Packet::SecretSubkey(a), Packet::PublicSubkey(b)) =>
+                            (a.public_cmp(&b) == Ordering::Equal,
+                             false),
+
+                        // For signatures, don't compare the unhashed
+                        // subpacket areas.  If it's the same
+                        // signature, then we keep what is the new
+                        // signature's unhashed subpacket area.
+                        (Packet::Signature(a), Packet::Signature(b)) =>
+                            (a.normalized_eq(b),
+                             a == b),
+
+                        (a, b) => {
+                            let identical = a == b;
+                            (identical, identical)
+                        }
+                    };
+
+                    if same {
+                        if identical {
+                            action = Drop;
+                        } else {
+                            action = Overwrite(i);
+                        }
+                        break;
+                    }
+                }
+            }
+
+            match action {
+                Drop => (),
+                Overwrite(i) => combined[i] = p,
+                Insert => {
+                    // New packet.  Add it to combined.
+                    combined.push(p);
+
+                    // Because the caller might insert the same packet
+                    // multiple times, we need to also add it to
+                    // packet_map.
+                    let i = combined.len() - 1;
+                    match packet_map.entry(hash) {
+                        Entry::Occupied(mut oe) => {
+                            oe.get_mut().push(i)
+                        }
+                        Entry::Vacant(ve) => {
+                            ve.insert(vec![ i ]);
+                        }
+                    }
+                }
             }
         }
 
@@ -4031,7 +4084,7 @@ mod test {
     }
 
     #[test]
-    fn insert_packets() {
+    fn insert_packets_add_sig() {
         use crate::armor;
         use crate::packet::Tag;
 
@@ -4054,6 +4107,145 @@ mod test {
     }
 
     #[test]
+    fn insert_packets_update_sig() -> Result<()> {
+        use std::time::Duration;
+
+        use crate::packet::signature::subpacket::Subpacket;
+        use crate::packet::signature::subpacket::SubpacketValue;
+
+        let (cert, _) = CertBuilder::general_purpose(None, Some("Test"))
+            .generate()?;
+        let packets = cert.clone().into_packets().count();
+
+        // Merge a signature with different unhashed subpacket areas.
+        // Make sure only the last variant is merged.
+        let sig = cert.primary_key().self_signatures().next()
+            .expect("binding signature");
+
+        let a = Subpacket::new(
+            SubpacketValue::SignatureExpirationTime(
+                Duration::new(1, 0).try_into()?),
+            false)?;
+        let b = Subpacket::new(
+            SubpacketValue::SignatureExpirationTime(
+                Duration::new(2, 0).try_into()?),
+            false)?;
+
+        let mut sig_a = sig.clone();
+        sig_a.unhashed_area_mut().add(a)?;
+        let mut sig_b = sig.clone();
+        sig_b.unhashed_area_mut().add(b)?;
+
+        // Insert sig_a, make sure it (and it alone) appears.
+        let cert2 = cert.clone().insert_packets(sig_a.clone())?;
+        let mut sigs = cert2.primary_key().self_signatures();
+        assert_eq!(sigs.next(), Some(&sig_a));
+        assert!(sigs.next().is_none());
+        assert_eq!(cert2.clone().into_packets().count(), packets);
+
+        // Insert sig_b, make sure it (and it alone) appears.
+        let cert2 = cert.clone().insert_packets(sig_b.clone())?;
+        let mut sigs = cert2.primary_key().self_signatures();
+        assert_eq!(sigs.next(), Some(&sig_b));
+        assert!(sigs.next().is_none());
+        assert_eq!(cert2.clone().into_packets().count(), packets);
+
+        // Insert sig_a and sig_b.  Make sure sig_b (and it alone)
+        // appears.
+        let cert2 = cert.clone().insert_packets(
+            vec![ sig_a.clone().into(), sig_b.clone() ])?;
+        let mut sigs = cert2.primary_key().self_signatures();
+        assert_eq!(sigs.next(), Some(&sig_b));
+        assert!(sigs.next().is_none());
+        assert_eq!(cert2.clone().into_packets().count(), packets);
+
+        // Insert sig_b and sig_a.  Make sure sig_a (and it alone)
+        // appears.
+        let cert2 = cert.clone().insert_packets(
+            vec![ sig_b.clone().into(), sig_a.clone() ])?;
+        let mut sigs = cert2.primary_key().self_signatures();
+        assert_eq!(sigs.next(), Some(&sig_a));
+        assert!(sigs.next().is_none());
+        assert_eq!(cert2.clone().into_packets().count(), packets);
+
+        Ok(())
+    }
+
+    #[test]
+    fn insert_packets_add_userid() -> Result<()> {
+        let (cert, _) = CertBuilder::general_purpose(None, Some("a"))
+            .generate()?;
+        let packets = cert.clone().into_packets().count();
+
+        let uid_a = UserID::from("a");
+        let uid_b = UserID::from("b");
+
+        // Insert a, make sure it appears once.
+        let cert2 = cert.clone().insert_packets(uid_a.clone())?;
+        let mut uids = cert2.userids();
+        assert_eq!(uids.next().unwrap().userid(), &uid_a);
+        assert!(uids.next().is_none());
+        assert_eq!(cert2.clone().into_packets().count(), packets);
+
+        // Insert b, make sure it also appears.
+        let cert2 = cert.clone().insert_packets(uid_b.clone())?;
+        let mut uids: Vec<UserID>
+            = cert2.userids().map(|ua| ua.userid().clone()).collect();
+        uids.sort();
+        let mut uids = uids.iter();
+        assert_eq!(uids.next().unwrap(), &uid_a);
+        assert_eq!(uids.next().unwrap(), &uid_b);
+        assert!(uids.next().is_none());
+        assert_eq!(cert2.clone().into_packets().count(), packets + 1);
+
+        Ok(())
+    }
+
+    #[test]
+    fn insert_packets_update_key() -> Result<()> {
+        use crate::crypto::Password;
+
+        let (cert, _) = CertBuilder::new().generate()?;
+        let packets = cert.clone().into_packets().count();
+        assert_eq!(cert.keys().count(), 1);
+
+        let key = cert.keys().secret().next().unwrap().key();
+        assert!(key.has_secret());
+        let key_a = key.clone().encrypt_secret(&Password::from("a"))?
+            .role_into_primary();
+        let key_b = key.clone().encrypt_secret(&Password::from("b"))?
+            .role_into_primary();
+
+        // Insert variant a.
+        let cert2 = cert.clone().insert_packets(key_a.clone())?;
+        assert_eq!(cert2.primary_key().key().parts_as_secret().unwrap(),
+                   &key_a);
+        assert_eq!(cert2.clone().into_packets().count(), packets);
+
+        // Insert variant b.
+        let cert2 = cert.clone().insert_packets(key_b.clone())?;
+        assert_eq!(cert2.primary_key().key().parts_as_secret().unwrap(),
+                   &key_b);
+        assert_eq!(cert2.clone().into_packets().count(), packets);
+
+        // Insert variant a then b.  We should keep b.
+        let cert2 = cert.clone().insert_packets(
+            vec![ key_a.clone(), key_b.clone() ])?;
+        assert_eq!(cert2.primary_key().key().parts_as_secret().unwrap(),
+                   &key_b);
+        assert_eq!(cert2.clone().into_packets().count(), packets);
+
+        // Insert variant b then a.  We should keep a.
+        let cert2 = cert.clone().insert_packets(
+            vec![ key_b.clone(), key_a.clone() ])?;
+        assert_eq!(cert2.primary_key().key().parts_as_secret().unwrap(),
+                   &key_a);
+        assert_eq!(cert2.clone().into_packets().count(), packets);
+
+        Ok(())
+    }
+
+    #[test]
     fn set_validity_period() {
         let p = &P::new();
 
diff --git a/openpgp/src/packet/key.rs b/openpgp/src/packet/key.rs
index a008ffb9..8eedacac 100644
--- a/openpgp/src/packet/key.rs
+++ b/openpgp/src/packet/key.rs
@@ -89,6 +89,7 @@
 use std::fmt;
 use std::cmp::Ordering;
 use std::convert::TryInto;
+use std::hash::Hasher;
 use std::time;
 
 #[cfg(test)]
@@ -917,6 +918,23 @@ impl<P, R> Key4<P, R>
     {
         self.public_cmp(b) == Ordering::Equal
     }
+
+    /// Hashes everything but any secret key material into state.
+    ///
+    /// This is an alternate implementation of [`Hash`], which never
+    /// hashes the secret key material.
+    ///
+    ///   [`Hash`]: https://doc.rust-lang.org/stable/std/hash/trait.Hash.html
+    pub fn public_hash<H>(&self, state: &mut H)
+        where H: Hasher
+    {
+        use std::hash::Hash;
+
+        self.common.hash(state);
+        self.creation_time.hash(state);
+        self.pk_algo.hash(state);
+        Hash::hash(&self.mpis(), state);
+    }
 }
 
 impl<R> Key4<key::PublicParts, R>
diff --git a/openpgp/src/packet/mod.rs b/openpgp/src/packet/mod.rs
index ed68e42f..d82f4965 100644
--- a/openpgp/src/packet/mod.rs
+++ b/openpgp/src/packet/mod.rs
@@ -157,6 +157,7 @@
 //! [partial body encoding]: https://tools.ietf.org/html/rfc4880#section-4.2.2.4
 //! [`Key::public_cmp`]: enum.Key.html#method.public_cmp
 use std::fmt;
+use std::hash::Hasher;
 use std::ops::{Deref, DerefMut};
 use std::slice;
 use std::iter::IntoIterator;
@@ -370,6 +371,47 @@ impl Packet {
             _ => Some(self.tag()),
         }
     }
+
+    /// Hashes most everything into state.
+    ///
+    /// This is an alternate implementation of [`Hash`], which does
+    /// not hash:
+    ///
+    ///   - The unhashed subpacket area of Signature packets.
+    ///   - Secret key material.
+    ///
+    ///   [`Hash`]: https://doc.rust-lang.org/stable/std/hash/trait.Hash.html
+    ///
+    /// Unlike [`Signature::normalize`], this method ignores
+    /// authenticated packets in the unhashed subpacket area.
+    ///
+    ///   [`Signature::normalize`]: struct.Signature.html#method.normalize
+    pub fn normalized_hash<H>(&self, state: &mut H)
+        where H: Hasher
+    {
+        use std::hash::Hash;
+
+        match self {
+            Packet::Signature(sig) => sig.normalized_hash(state),
+            Packet::OnePassSig(x) => Hash::hash(&x, state),
+            Packet::PublicKey(k) => k.public_hash(state),
+            Packet::PublicSubkey(k) => k.public_hash(state),
+            Packet::SecretKey(k) => k.public_hash(state),
+            Packet::SecretSubkey(k) => k.public_hash(state),
+            Packet::Marker(x) => Hash::hash(&x, state),
+            Packet::Trust(x) => Hash::hash(&x, state),
+            Packet::UserID(x) => Hash::hash(&x, state),
+            Packet::UserAttribute(x) => Hash::hash(&x, state),
+            Packet::Literal(x) => Hash::hash(&x, state),
+            Packet::CompressedData(x) => Hash::hash(&x, state),
+            Packet::PKESK(x) => Hash::hash(&x, state),
+            Packet::SKESK(x) => Hash::hash(&x, state),
+            Packet::SEIP(x) => Hash::hash(&x, state),
+            Packet::MDC(x) => Hash::hash(&x, state),
+            Packet::AED(x) => Hash::hash(&x, state),
+            Packet::Unknown(x) => Hash::hash(&x, state),
+        }
+    }
 }
 
 // Allow transparent access of common fields.
diff --git a/openpgp/src/packet/signature.rs b/openpgp/src/packet/signature.rs
index 87d4c511..cd9ead07 100644
--- a/openpgp/src/packet/signature.rs
+++ b/openpgp/src/packet/signature.rs
@@ -116,6 +116,7 @@
 
 use std::cmp::Ordering;
 use std::fmt;
+use std::hash::Hasher;
 use std::ops::{Deref, DerefMut};
 use std::time::SystemTime;
 
@@ -1993,6 +1994,11 @@ impl crate::packet::Signature {
     /// which need to be validated.  Ignoring the unhashed subpackets
     /// means that we can deduplicate signatures using this predicate.
     ///
+    /// Unlike [`Signature::normalize`], this method ignores
+    /// authenticated packets in the unhashed subpacket area.
+    ///
+    ///   [`Signature::normalize`]: #method.normalize
+    ///
     /// # Examples
     ///
     /// ```
@@ -2047,6 +2053,11 @@ impl crate::packet::Signature {
     /// which need to be validated.  Ignoring the unhashed subpackets
     /// means that we can deduplicate signatures using this predicate.
     ///
+    /// Unlike [`Signature::normalize`], this method ignores
+    /// authenticated packets in the unhashed subpacket area.
+    ///
+    ///   [`Signature::normalize`]: #method.normalize
+    ///
     /// # Examples
     ///
     /// ```
@@ -2091,6 +2102,31 @@ impl crate::packet::Signature {
             .then_with(|| self.mpis().cmp(other.mpis()))
     }
 
+    /// Hashes everything but the unhashed subpacket area into state.
+    ///
+    /// This is an alternate implementation of [`Hash`], which does
+    /// not hash the unhashed subpacket area.
+    ///
+    ///   [`Hash`]: https://doc.rust-lang.org/stable/std/hash/trait.Hash.html
+    ///
+    /// Unlike [`Signature::normalize`], this method ignores
+    /// authenticated packets in the unhashed subpacket area.
+    ///
+    ///   [`Signature::normalize`]: #method.normalize
+    pub fn normalized_hash<H>(&self, state: &mut H)
+        where H: Hasher
+    {
+        use std::hash::Hash;
+
+        self.version.hash(state);
+        self.typ.hash(state);
+        self.pk_algo.hash(state);
+        self.hash_algo.hash(state);
+        self.hashed_area().hash(state);
+        self.digest_prefix().hash(state);
+        Hash::hash(&self.mpis(), state);
+    }
+
     /// Normalizes the signature.
     ///
     /// This function normalizes the *unhashed* signature subpackets.