1 files changed, 16 insertions, 9 deletions
diff --git a/src/Storage/Key.hs b/src/Storage/Key.hs
index 3ed4a66..8e6d04c 100644
--- a/src/Storage/Key.hs
+++ b/src/Storage/Key.hs
@@ -3,9 +3,13 @@ module Storage.Key (
     storeKey, loadKey,
 ) where
 
+import Control.Concurrent.MVar
+import Control.Monad
+
 import Data.ByteArray
 import qualified Data.ByteString.Char8 as BC
 import qualified Data.ByteString.Lazy as BL
+import qualified Data.Map as M
 
 import System.FilePath
 import System.IO.Error
@@ -20,17 +24,20 @@ class Storable pub => KeyPair sec pub | sec -> pub, pub -> sec where
     keyFromData :: ScrubbedBytes -> Stored pub -> Maybe sec
 
 
-keyStorage :: Storage -> FilePath
-keyStorage (Storage base) = base </> "keys"
-
-keyFilePath :: KeyPair sec pub => Stored pub -> FilePath
-keyFilePath pkey = keyStorage (storedStorage pkey) </> (BC.unpack $ showRef $ storedRef pkey)
+keyFilePath :: KeyPair sec pub => FilePath -> Stored pub -> FilePath
+keyFilePath sdir pkey = sdir </> "keys" </> (BC.unpack $ showRef $ storedRef pkey)
 
 storeKey :: KeyPair sec pub => sec -> IO ()
-storeKey key = writeFileOnce (keyFilePath $ keyGetPublic key) (BL.fromStrict $ convert $ keyGetData key)
+storeKey key = do
+    let spub = keyGetPublic key
+    case stBacking $ storedStorage spub of
+         StorageDir dir -> writeFileOnce (keyFilePath dir spub) (BL.fromStrict $ convert $ keyGetData key)
+         StorageMemory { memKeys = kstore } -> modifyMVar_ kstore $ return . M.insert (refDigest $ storedRef spub) (keyGetData key)
 
 loadKey :: KeyPair sec pub => Stored pub -> IO (Maybe sec)
 loadKey spub = do
-    tryIOError (BC.readFile (keyFilePath spub)) >>= \case
-        Right kdata -> return $ keyFromData (convert kdata) spub
-        Left _ -> return Nothing
+    case stBacking $ storedStorage spub of
+         StorageDir dir -> tryIOError (BC.readFile (keyFilePath dir spub)) >>= \case
+             Right kdata -> return $ keyFromData (convert kdata) spub
+             Left _ -> return Nothing
+         StorageMemory { memKeys = kstore } -> (flip keyFromData spub <=< M.lookup (refDigest $ storedRef spub)) <$> readMVar kstore