From 0c5c03357144de4acb872dc4d8c6ba4b6f6ae76e Mon Sep 17 00:00:00 2001
From: Joe Crayne <joe@jerkface.net>
Date: Tue, 16 Jul 2019 17:49:16 -0400
Subject: Faster subkey verificaiton.

---
 lib/KeyDB.hs      | 29 +++++++++++++++++++++++++----
 lib/Transforms.hs | 38 ++++++++++++++++++++++++++++++--------
 2 files changed, 55 insertions(+), 12 deletions(-)

diff --git a/lib/KeyDB.hs b/lib/KeyDB.hs
index 0bc0fb3..c92f614 100644
--- a/lib/KeyDB.hs
+++ b/lib/KeyDB.hs
@@ -5,11 +5,13 @@ module KeyDB
     , SubKey(..)
     , KeyData(..)
     , KeyDB
+    , KeyGrip(..)
     , emptyKeyDB
     , keyData
     , kkData
     , lookupKeyData
     , lookupByGrip
+    , associatedKeys
     , fingerprintGrip
     , smallprGrip
     , transmute
@@ -26,6 +28,7 @@ module KeyDB
     , flattenKeys
     , flattenFiltered
     , UidString(..)
+    , buildGripMap
     ) where
 
 import Control.Monad
@@ -79,6 +82,17 @@ data KeyDB = KeyDB
     , byGrip   :: IMap KeyGrip [KeyKey]
     } deriving Show
 
+
+-- | TODO: This is an optimization to legacy (pre-KeyDB) code.  Ultimately it
+-- should be unneccessary.
+buildGripMap :: [Packet] -> IMap KeyGrip [Packet]
+buildGripMap ps = foldr go I.empty ps
+ where
+    go pkt m =  I.alter (\case Just ks -> Just (pkt:ks)
+                               Nothing -> Just [pkt])
+                        (fingerprintGrip . fingerprint $ pkt)
+                        m
+
 emptyKeyDB :: KeyDB
 emptyKeyDB = KeyDB { byKeyKey = Map.empty, byGrip = I.empty }
 
@@ -110,16 +124,23 @@ transmute perform update db = do
                            -- Note: We currently leave deleted-keys in the byGrip map.
                            , concatMap snd $ Map.elems bkk )
 
+mpGrip :: MappedPacket -> KeyGrip
+mpGrip mp = fingerprintGrip $ fingerprint $ packet mp
+
+associatedKeys :: KeyData -> [MappedPacket]
+associatedKeys kd = keyMappedPacket kd : [ k | SubKey k _ <- Map.elems (keySubKeys kd) ]
+
 alterKeyDB :: (Maybe KeyData -> Maybe KeyData) -> KeyKey -> KeyDB -> KeyDB
 alterKeyDB update kk db = db
     { byKeyKey = Map.alter update kk (byKeyKey db)
     , byGrip = case Map.lookup kk (byKeyKey db) of
         Just _  -> byGrip db
         Nothing -> case update Nothing of
-            Just kd -> I.alter (\case Nothing  -> Just [kk]
-                                      Just kks -> Just $ mergeL [kk] kks)
-                               (fingerprintGrip $ fingerprint $ packet $ keyMappedPacket kd)
-                               (byGrip db)
+            Just kd -> let go g m = I.alter (\case Nothing  -> Just [kk]
+                                                   Just kks -> Just $ mergeL [kk] kks)
+                                            g
+                                            m
+                        in foldr go (byGrip db) $ map mpGrip $ associatedKeys kd
             Nothing -> byGrip db
     }
 
diff --git a/lib/Transforms.hs b/lib/Transforms.hs
index e7097ba..f3cd5e3 100644
--- a/lib/Transforms.hs
+++ b/lib/Transforms.hs
@@ -16,6 +16,7 @@ import Data.Ord
 import Data.OpenPGP
 import Data.OpenPGP.Util
 import Data.Word
+import qualified IntMapClass as I
 import KeyDB
 import KeyRing.Types
 import FunctorToMaybe
@@ -236,12 +237,19 @@ has_tag tag p = isSignaturePacket p
 
 
 
-verifyBindings :: [Packet] -> [Packet] -> ([SignatureOver], [SignatureOver])
-verifyBindings keys nonkeys = (top ++ filter isSubkeySignature embedded,othersigs)
+verifyBindings :: I.IMap KeyGrip [Packet] -> [Packet] -> ([SignatureOver], [SignatureOver])
+verifyBindings gmap nonkeys = (top ++ filter isSubkeySignature embedded,othersigs)
  where
     verified = do
-        sig <- signatures (Message nonkeys)
-        let v = verify (Message keys) sig
+        sigs <- signatures (Message nonkeys)
+        sig <- signatures_over sigs
+        let grip = issuerGrip sig
+            gks = concat [ ks | g <- maybeToList grip
+                              , ks <- maybeToList $ I.lookup g gmap ]
+            kmsg = Message
+                    $ if null gks then maybe (concat $ I.elems gmap) (const []) grip
+                                  else gks
+            v = verify kmsg (sigs { signatures_over = [sig] })
         guard (not . null $ signatures_over v)
         return v
     (top,othersigs) = partition isSubkeySignature verified
@@ -292,13 +300,15 @@ getBindings ::
 getBindings pkts = (sigs,bindings)
  where
     (sigs,concat->bindings) = unzip $ do
-        keys <- disjoint_fp (filter isKey pkts)
-        let (bs,sigs) = verifyBindings keys pkts
+        keys <- take 1 $ disjoint_fp (filter isKey pkts)
+        let gmap = buildGripMap keys
+            (bs,sigs) = verifyBindings gmap pkts
         return . ((keys,sigs),) $ do
             b <- bs -- trace ("sigs = "++show (map (map signature_issuer . signatures_over) sigs)) bs
             i <- map signature_issuer (signatures_over b)
             i <- maybeToList i
-            who <- maybeToList $ find_key (show . fingerprint) (Message keys) i
+            g <- maybeToList $ smallprGrip i
+            who <- take 1 $ concat $ maybeToList $ I.lookup g gmap
             let (code,claimants) =
                     case () of
                         _ | who == topkey b -> (1,[])
@@ -627,7 +637,19 @@ makeSig doDecrypt top fname subkey_p tags mbsig = do
 
 -- TODO: Use fingerprint to narrow candidates.
 candidateSignerKeys :: KeyDB -> Packet -> [Packet]
-candidateSignerKeys db sig = map keyPacket $ keyData db
+candidateSignerKeys db sig =
+    case issuerGrip sig of
+        Just g -> concatMap (map packet . associatedKeys) $ lookupByGrip g db
+        _      -> map keyPacket $ keyData db
+
+issuerGrip :: Packet -> Maybe KeyGrip
+issuerGrip sig = do
+    IssuerPacket hexfp <- find isIssuer (hashed_subpackets sig ++ unhashed_subpackets sig)
+    smallprGrip hexfp
+
+isIssuer :: SignatureSubpacket -> Bool
+isIssuer (IssuerPacket _) = True
+isIssuer _                = False
 
 performManipulations ::
                 (PacketDecrypter)
-- 
cgit v1.2.3