memoize encryption as well as decryption

1 files changed, 17 insertions, 9 deletions

diff --git a/lib/PacketTranscoder.hs b/lib/PacketTranscoder.hs
index 6d1d9b8..f4b4cce 100644
--- a/lib/PacketTranscoder.hs
+++ b/lib/PacketTranscoder.hs
@@ -3,6 +3,7 @@
 {-# LANGUAGE PatternGuards #-}
 module PacketTranscoder where
 
+import GHC.Stack
 import Control.Monad
 import Data.IORef
 import Data.List
@@ -127,7 +128,7 @@ makeMemoizingDecrypter operation ctx (workingkey,keys) = do
              $ listToMaybe $ filter (\sp -> isNothing (passSpecRingFile sp)
                                             && isNothing (passSpecKeySpec sp))
                                         $ passspecs
-        unkeysRef <- newIORef (Map.empty :: Map.Map KeyKey Packet)
+        unkeysRef <- newIORef (Map.empty :: Map.Map (KeyKey,SymmetricAlgorithm,S2K) Packet)
         return $ doDecrypt unkeysRef ({- pws `Map.union` -} pws2) defpw (not $ null agentspec)
     else let PassphraseMemoizer f = head chains
          in return f
@@ -138,7 +139,7 @@ makeMemoizingDecrypter operation ctx (workingkey,keys) = do
     (agentspec,passspecs) = partition isAgent passspecs0
                             where isAgent PassphraseAgent = True
                                   isAgent _               = False
-    doDecrypt :: IORef (Map.Map KeyKey Packet)
+    doDecrypt :: IORef (Map.Map (KeyKey,SymmetricAlgorithm,S2K) Packet)
                 -> Map.Map FilePath (IO S.ByteString)
                 -> Maybe (IO S.ByteString)
                 -> Bool
@@ -155,6 +156,9 @@ makeMemoizingDecrypter operation ctx (workingkey,keys) = do
                                     q <- fmap packet $ Map.lookup (keykey working) keys
                                     return (mp0, Query (packet mp0) (queryUID q) (Just working)))
 
+            dest_s2k' | dest_alg==Unencrypted = S2K 100 ""
+                      | otherwise             = dest_s2k
+
             wk = packet mp0
             kk = keykey wk
             fs = Map.keys $ locations mp
@@ -184,13 +188,16 @@ makeMemoizingDecrypter operation ctx (workingkey,keys) = do
                 case symmetric_algorithm wkun of
 
                     Unencrypted -> do
-                        writeIORef unkeysRef (Map.insert kk wkun unkeys)
+                        writeIORef unkeysRef (Map.insert (kk,Unencrypted,S2K 100 "") wkun unkeys)
                         ek <- case dest_alg of
                                 Unencrypted -> return $ Just wkun
-                                _           -> encryptSecretKey pw dest_s2k dest_alg wkun
+                                _           -> encryptSecretKey pw dest_s2k' dest_alg wkun
+
                         case ek of
                             Nothing   -> retryOrFail
-                            Just wken -> return $ KikiSuccess wken
+                            Just wken -> do
+                                modifyIORef unkeysRef (Map.insert (kk,dest_alg,dest_s2k') wken)
+                                return $ KikiSuccess wken
 
                     _ -> retryOrFail
 
@@ -225,13 +232,14 @@ makeMemoizingDecrypter operation ctx (workingkey,keys) = do
                     quit s
                     return ( maybe "" S8.pack mbpw, guard (ask /= AskNew) >> Just clear )
 
-        -- putStrLn $ concat [show (symmetric_algorithm wk,s2k wk)," --> ",show (dest_alg,dest_s2k)]
-        if symmetric_algorithm wk == dest_alg
-            && ( symmetric_algorithm wk == Unencrypted || s2k wk == dest_s2k )
+        calls <- currentCallStack
+        putStrLn $ concat [fingerprint wk," ", show (symmetric_algorithm wk,s2k wk)," --> ",show (dest_alg,dest_s2k)]
+        mapM_ putStrLn calls
+        if symmetric_algorithm wk == dest_alg && s2k wk == dest_s2k'
             then return (KikiSuccess wk)
             else maybe (decryptIt getpws)
                        (return . KikiSuccess)
-                    $ Map.lookup kk unkeys
+                    $ Map.lookup (kk,dest_alg,dest_s2k') unkeys
 
 keyQueries :: Maybe String -> Map InputFile (StreamInfo,Message) -> (Maybe MappedPacket, Map KeyKey (OriginMapped Query))
 keyQueries grip ringPackets = (mwk, fmap makeQuery keys)