Replace storedStorage usage with MonadHead

13 files changed, 111 insertions, 99 deletions

diff --git a/src/Attach.hs b/src/Attach.hs
index 67828aa..48d18d8 100644
--- a/src/Attach.hs
+++ b/src/Attach.hs
@@ -42,9 +42,8 @@ instance PairingResult AttachIdentity where
 
     pairingVerifyResult (AttachIdentity sdata keys) = do
         curid <- lsIdentity . fromStored <$> svcGetLocal
-        secret <- maybe (throwError "failed to load own secret key") return =<<
-            liftIO (loadKey $ iddKeyIdentity $ fromStored $ signedData $ fromStored curid)
-        sdata' <- liftIO $ wrappedStore (storedStorage sdata) =<< signAdd secret (fromStored sdata)
+        secret <- loadKey $ iddKeyIdentity $ fromStored $ signedData $ fromStored curid
+        sdata' <- mstore =<< signAdd secret (fromStored sdata)
         return $ do
             guard $ iddKeyIdentity (fromStored $ signedData $ fromStored sdata) ==
                 iddKeyIdentity (fromStored $ signedData $ fromStored curid)
@@ -52,26 +51,25 @@ instance PairingResult AttachIdentity where
             guard $ iddPrev (fromStored $ signedData $ fromStored $ idData identity) == [curid]
             return (identity, keys)
 
-    pairingFinalizeRequest (identity, keys) = updateLocalHead_ $ \slocal -> liftIO $ do
+    pairingFinalizeRequest (identity, keys) = updateLocalHead_ $ \slocal -> do
         let owner = finalOwner identity
-            st = storedStorage slocal
+        st <- getStorage
         pkeys <- mapM (copyStored st) [ idKeyIdentity owner, idKeyMessage owner ]
-        mapM_ storeKey $ catMaybes [ keyFromData sec pub | sec <- keys, pub <- pkeys ]
+        liftIO $ mapM_ storeKey $ catMaybes [ keyFromData sec pub | sec <- keys, pub <- pkeys ]
 
         shared <- makeSharedStateUpdate st (Just owner) (lsShared $ fromStored slocal)
-        wrappedStore st (fromStored slocal)
+        mstore (fromStored slocal)
             { lsIdentity = idData identity
             , lsShared = [ shared ]
             }
 
     pairingFinalizeResponse = do
-        st <- storedStorage <$> svcGetLocal
         owner <- mergeSharedIdentity
         pid <- asks svcPeerIdentity
-        secret <- maybe (throwError "failed to load secret key") return =<< liftIO (loadKey $ idKeyIdentity owner)
-        identity <- liftIO $ wrappedStore st =<< sign secret =<< wrappedStore st (emptyIdentityData $ idKeyIdentity pid)
+        secret <- loadKey $ idKeyIdentity owner
+        identity <- mstore =<< sign secret =<< mstore (emptyIdentityData $ idKeyIdentity pid)
             { iddPrev = [idData pid], iddOwner = Just (idData owner) }
-        skeys <- liftIO $ map keyGetData . catMaybes <$> mapM loadKey [ idKeyIdentity owner, idKeyMessage owner ]
+        skeys <- map keyGetData . catMaybes <$> mapM loadKeyMb [ idKeyIdentity owner, idKeyMessage owner ]
         return $ AttachIdentity identity skeys
 
     defaultPairingAttributes _ = PairingAttributes
diff --git a/src/Channel.hs b/src/Channel.hs
index a1773bd..167c1ba 100644
--- a/src/Channel.hs
+++ b/src/Channel.hs
@@ -77,13 +77,13 @@ instance Storable ChannelAcceptData where
 keySize :: Int
 keySize = 32
 
-createChannelRequest :: (MonadIO m) => Storage -> UnifiedIdentity -> UnifiedIdentity -> m (Stored ChannelRequest)
-createChannelRequest st self peer = liftIO $ do
-    (_, xpublic) <- generateKeys st
-    Just skey <- loadKey $ idKeyMessage self
-    wrappedStore st =<< sign skey =<< wrappedStore st ChannelRequest { crPeers = sort [idData self, idData peer], crKey = xpublic }
+createChannelRequest :: (MonadStorage m, MonadIO m, MonadError String m) => UnifiedIdentity -> UnifiedIdentity -> m (Stored ChannelRequest)
+createChannelRequest self peer = do
+    (_, xpublic) <- liftIO . generateKeys =<< getStorage
+    skey <- loadKey $ idKeyMessage self
+    mstore =<< sign skey =<< mstore ChannelRequest { crPeers = sort [idData self, idData peer], crKey = xpublic }
 
-acceptChannelRequest :: (MonadIO m, MonadError String m) => UnifiedIdentity -> UnifiedIdentity -> Stored ChannelRequest -> m (Stored ChannelAccept, Channel)
+acceptChannelRequest :: (MonadStorage m, MonadIO m, MonadError String m) => UnifiedIdentity -> UnifiedIdentity -> Stored ChannelRequest -> m (Stored ChannelAccept, Channel)
 acceptChannelRequest self peer req = do
     case sequence $ map validateIdentity $ crPeers $ fromStored $ signedData $ fromStored req of
         Nothing -> throwError $ "invalid peers in channel request"
@@ -95,11 +95,10 @@ acceptChannelRequest self peer req = do
     when (idKeyMessage peer `notElem` (map (sigKey . fromStored) $ signedSignature $ fromStored req)) $
         throwError $ "channel requent not signed by peer"
 
-    let st = storedStorage req
+    (xsecret, xpublic) <- liftIO . generateKeys =<< getStorage
+    skey <- loadKey $ idKeyMessage self
+    acc <- mstore =<< sign skey =<< mstore ChannelAccept { caRequest = req, caKey = xpublic }
     liftIO $ do
-        (xsecret, xpublic) <- generateKeys st
-        Just skey <- loadKey $ idKeyMessage self
-        acc <- wrappedStore st =<< sign skey =<< wrappedStore st ChannelAccept { caRequest = req, caKey = xpublic }
         let chPeers = crPeers $ fromStored $ signedData $ fromStored req
             chKey = BA.take keySize $ dhSecret xsecret $
                 fromStored $ crKey $ fromStored $ signedData $ fromStored req
@@ -125,9 +124,7 @@ acceptedChannel self peer acc = do
     when (idKeyMessage self `notElem` (map (sigKey . fromStored) $ signedSignature $ fromStored req)) $
         throwError $ "original channel request not signed by us"
 
-    xsecret <- liftIO (loadKey $ crKey $ fromStored $ signedData $ fromStored req) >>= \case
-        Just key -> return key
-        Nothing  -> throwError $ "secret key not found"
+    xsecret <- loadKey $ crKey $ fromStored $ signedData $ fromStored req
     let chPeers = crPeers $ fromStored $ signedData $ fromStored req
         chKey = BA.take keySize $ dhSecret xsecret $
             fromStored $ caKey $ fromStored $ signedData $ fromStored acc
diff --git a/src/Identity.hs b/src/Identity.hs
index 834e5ee..9653077 100644
--- a/src/Identity.hs
+++ b/src/Identity.hs
@@ -21,7 +21,8 @@ module Identity (
 import Control.Arrow
 import Control.Monad
 import Control.Monad.Except
-import qualified Control.Monad.Identity as I
+import Control.Monad.Identity qualified as I
+import Control.Monad.Reader
 
 import Data.Either
 import Data.Foldable
@@ -122,13 +123,17 @@ createIdentity st name owner = do
 
     let signOwner idd
             | Just o <- owner = do
-                Just ownerSecret <- loadKey (iddKeyIdentity $ fromStored $ signedData $ fromStored $ idData o)
+                Just ownerSecret <- loadKeyMb (iddKeyIdentity $ fromStored $ signedData $ fromStored $ idData o)
                 signAdd ownerSecret idd
             | otherwise = return idd
 
-    Just identity <- return . validateIdentity =<< wrappedStore st =<< signOwner =<< sign secret =<<
-        wrappedStore st (emptyIdentityData public)
-        { iddName = name, iddOwner = idData <$> owner, iddKeyMessage = Just publicMsg }
+    Just identity <- flip runReaderT st $ do
+        return . validateIdentity =<< mstore =<< signOwner =<< sign secret =<<
+            mstore (emptyIdentityData public)
+                { iddName = name
+                , iddOwner = idData <$> owner
+                , iddKeyMessage = Just publicMsg
+                }
     return identity
 
 validateIdentity :: Stored (Signed IdentityData) -> Maybe UnifiedIdentity
@@ -192,7 +197,7 @@ lookupProperty sel topHeads = findResult filteredLayers
           findResult [xs] = Just $ snd $ minimumBy (comparing fst) xs
           findResult (_:rest) = findResult rest
 
-mergeIdentity :: Foldable m => Identity m -> IO UnifiedIdentity
+mergeIdentity :: (Foldable f, MonadStorage m, MonadError String m, MonadIO m) => Identity f -> m UnifiedIdentity
 mergeIdentity idt | Just idt' <- toUnifiedIdentity idt = return idt'
 mergeIdentity idt = do
     (owner, ownerData) <- case idOwner_ idt of
@@ -201,11 +206,9 @@ mergeIdentity idt = do
                     | otherwise -> do owner <- mergeIdentity cowner
                                       return (Just owner, Just $ idData owner)
 
-    (sid:_) <- return $ toList $ idDataF idt
-    let st = storedStorage sid
-        public = idKeyIdentity idt
-    Just secret <- loadKey public
-    sdata <- wrappedStore st =<< sign secret =<< wrappedStore st (emptyIdentityData public)
+    let public = idKeyIdentity idt
+    secret <- loadKey public
+    sdata <- mstore =<< sign secret =<< mstore (emptyIdentityData public)
         { iddPrev = toList $ idDataF idt, iddOwner = ownerData }
     return $ idt { idData_ = I.Identity sdata, idOwner_ = toComposedIdentity <$> owner }
 
diff --git a/src/Main.hs b/src/Main.hs
index b3f503d..cbefeb2 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -113,7 +113,8 @@ main = do
                 Nothing -> error "ref does not exist"
                 Just refs
                     | Just idt <- validateIdentityF $ map wrappedLoad refs -> do
-                        BC.putStrLn . showRefDigest . refDigest . storedRef . idData =<< interactiveIdentityUpdate idt
+                        BC.putStrLn . showRefDigest . refDigest . storedRef . idData =<<
+                            (either fail return <=< runExceptT $ runReaderT (interactiveIdentityUpdate idt) st)
                     | otherwise -> error "invalid identity"
 
         ["test"] -> runTestTool st
@@ -413,7 +414,7 @@ cmdDiscoveryInit = void $ do
 cmdDiscovery :: Command
 cmdDiscovery = void $ do
     Just peer <- gets csIcePeer
-    st <- gets (storedStorage . headStoredObject . csHead)
+    st <- getStorage
     sref <- asks ciLine
     eprint <- asks ciPrint
     liftIO $ readRef st (BC.pack sref) >>= \case
diff --git a/src/Message.hs b/src/Message.hs
index 41a88b0..ac67f07 100644
--- a/src/Message.hs
+++ b/src/Message.hs
@@ -73,8 +73,8 @@ instance Service DirectMessage where
         let msg = fromStored smsg
         powner <- asks $ finalOwner . svcPeerIdentity
         erb <- svcGetLocal
-        let st = storedStorage erb
-            DirectMessageThreads prev _ = lookupSharedValue $ lsShared $ fromStored erb
+        st <- getStorage
+        let DirectMessageThreads prev _ = lookupSharedValue $ lsShared $ fromStored erb
             sent = findMsgProperty powner msSent prev
             received = findMsgProperty powner msReceived prev
             received' = filterAncestors $ smsg : received
@@ -153,21 +153,20 @@ findMsgProperty pid sel mss = concat $ flip findProperty mss $ \x -> do
 sendDirectMessage :: (Foldable f, Applicative f, MonadHead LocalState m, MonadError String m)
                   => Identity f -> Text -> m (Stored DirectMessage)
 sendDirectMessage pid text = updateLocalHead $ \ls -> do
-    let st = storedStorage ls
-        self = localIdentity $ fromStored ls
+    let self = localIdentity $ fromStored ls
         powner = finalOwner pid
-    flip updateSharedState ls $ \(DirectMessageThreads prev _) -> liftIO $ do
+    flip updateSharedState ls $ \(DirectMessageThreads prev _) -> do
         let sent = findMsgProperty powner msSent prev
             received = findMsgProperty powner msReceived prev
 
-        time <- getZonedTime
-        smsg <- wrappedStore st DirectMessage
+        time <- liftIO getZonedTime
+        smsg <- mstore DirectMessage
             { msgFrom = toComposedIdentity $ finalOwner self
             , msgPrev = filterAncestors $ sent ++ received
             , msgTime = time
             , msgText = text
             }
-        next <- wrappedStore st $ MessageState
+        next <- mstore MessageState
             { msPrev = prev
             , msPeer = powner
             , msSent = [smsg]
diff --git a/src/Network.hs b/src/Network.hs
index 3614de0..96f8527 100644
--- a/src/Network.hs
+++ b/src/Network.hs
@@ -22,6 +22,7 @@ import Control.Concurrent.STM
 import Control.Exception
 import Control.Monad
 import Control.Monad.Except
+import Control.Monad.Reader
 import Control.Monad.State
 
 import qualified Data.ByteString.Char8 as BC
@@ -527,7 +528,7 @@ withPeerIdentity peer act = liftIO $ atomically $ readTVar (peerIdentityVar peer
 
 setupChannel :: UnifiedIdentity -> Peer -> UnifiedIdentity -> WaitingRefCallback
 setupChannel identity peer upid = do
-    req <- createChannelRequest (peerStorage peer) identity upid
+    req <- flip runReaderT (peerStorage peer) $ createChannelRequest identity upid
     let reqref = refDigest $ storedRef req
     let hitems =
             [ TrChannelRequest reqref
@@ -544,7 +545,7 @@ setupChannel identity peer upid = do
 handleChannelRequest :: Peer -> UnifiedIdentity -> Ref -> WaitingRefCallback
 handleChannelRequest peer identity req = do
     withPeerIdentity peer $ \upid -> do
-        (acc, ch) <- acceptChannelRequest identity upid (wrappedLoad req)
+        (acc, ch) <- flip runReaderT (peerStorage peer) $ acceptChannelRequest identity upid (wrappedLoad req)
         liftIO $ atomically $ do
             getPeerChannel peer >>= \case
                 ChannelPeerRequest wr | wrDigest wr == refDigest req -> do
diff --git a/src/PubKey.hs b/src/PubKey.hs
index 483a94b..f69d739 100644
--- a/src/PubKey.hs
+++ b/src/PubKey.hs
@@ -1,6 +1,6 @@
 module PubKey (
     PublicKey, SecretKey,
-    KeyPair(generateKeys), loadKey,
+    KeyPair(generateKeys), loadKey, loadKeyMb,
     Signature(sigKey), Signed, signedData, signedSignature,
     sign, signAdd, isSignedBy,
 
@@ -97,14 +97,14 @@ instance Storable a => Storable (Signed a) where
                 throwError "signature verification failed"
         return $ Signed sdata sigs
 
-sign :: SecretKey -> Stored a -> IO (Signed a)
+sign :: MonadStorage m => SecretKey -> Stored a -> m (Signed a)
 sign secret val = signAdd secret $ Signed val []
 
-signAdd :: SecretKey -> Signed a -> IO (Signed a)
+signAdd :: MonadStorage m => SecretKey -> Signed a -> m (Signed a)
 signAdd (SecretKey secret spublic) (Signed val sigs) = do
     let PublicKey public = fromStored spublic
         sig = ED.sign secret public $ storedRef val
-    ssig <- wrappedStore (storedStorage val) $ Signature spublic sig
+    ssig <- mstore $ Signature spublic sig
     return $ Signed val (ssig : sigs)
 
 isSignedBy :: Signed a -> Stored PublicKey -> Bool
diff --git a/src/State.hs b/src/State.hs
index 12f9db9..b575ffa 100644
--- a/src/State.hs
+++ b/src/State.hs
@@ -2,7 +2,8 @@ module State (
     LocalState(..),
     SharedState, SharedType(..),
     SharedTypeID, mkSharedTypeID,
-    MonadStorage(..), MonadHead(..),
+
+    MonadHead(..),
     updateLocalHead_,
 
     loadLocalStateHead,
@@ -83,19 +84,12 @@ instance SharedType (Maybe ComposedIdentity) where
     sharedTypeID _ = mkSharedTypeID "0c6c1fe0-f2d7-4891-926b-c332449f7871"
 
 
-class Monad m => MonadStorage m where
-    getStorage :: m Storage
-
 class (MonadIO m, MonadStorage m) => MonadHead a m where
     updateLocalHead :: (Stored a -> m (Stored a, b)) -> m b
 
 updateLocalHead_ :: MonadHead a m => (Stored a -> m (Stored a)) -> m ()
 updateLocalHead_ f = updateLocalHead (fmap (,()) . f)
 
-
-instance Monad m => MonadStorage (ReaderT (Head a) m) where
-    getStorage = asks $ headStorage
-
 instance (HeadType a, MonadIO m) => MonadHead a (ReaderT (Head a) m) where
     updateLocalHead f = do
         h <- ask
@@ -146,7 +140,7 @@ updateSharedState :: forall a b m. (SharedType a, MonadHead LocalState m) => (a
 updateSharedState f = \ls -> do
     let shared = lsShared $ fromStored ls
         val = lookupSharedValue shared
-        st = storedStorage ls
+    st <- getStorage
     (val', x) <- f val
     (,x) <$> if toComponents val' == toComponents val
                 then return ls
@@ -170,36 +164,36 @@ makeSharedStateUpdate st val prev = liftIO $ wrappedStore st SharedState
 mergeSharedIdentity :: (MonadHead LocalState m, MonadError String m) => m UnifiedIdentity
 mergeSharedIdentity = updateLocalHead $ updateSharedState $ \case
     Just cidentity -> do
-        identity <- liftIO $ mergeIdentity cidentity
+        identity <- mergeIdentity cidentity
         return (Just $ toComposedIdentity identity, identity)
     Nothing -> throwError "no existing shared identity"
 
 updateSharedIdentity :: (MonadHead LocalState m, MonadError String m) => m ()
 updateSharedIdentity = updateLocalHead_ $ updateSharedState_ $ \case
     Just identity -> do
-        Just . toComposedIdentity <$> liftIO (interactiveIdentityUpdate identity)
+        Just . toComposedIdentity <$> interactiveIdentityUpdate identity
     Nothing -> throwError "no existing shared identity"
 
-interactiveIdentityUpdate :: Foldable m => Identity m -> IO UnifiedIdentity
+interactiveIdentityUpdate :: (Foldable f, MonadStorage m, MonadIO m, MonadError String m) => Identity f -> m UnifiedIdentity
 interactiveIdentityUpdate identity = do
-    let st = storedStorage $ head $ toList $ idDataF $ identity
-        public = idKeyIdentity identity
-
-    T.putStr $ T.concat $ concat
-        [ [ T.pack "Name" ]
-        , case idName identity of
-               Just name -> [T.pack " [", name, T.pack "]"]
-               Nothing -> []
-        , [ T.pack ": " ]
-        ]
-    hFlush stdout
-    name <- T.getLine
+    let public = idKeyIdentity identity
+
+    name <- liftIO $ do
+        T.putStr $ T.concat $ concat
+            [ [ T.pack "Name" ]
+            , case idName identity of
+                   Just name -> [T.pack " [", name, T.pack "]"]
+                   Nothing -> []
+            , [ T.pack ": " ]
+            ]
+        hFlush stdout
+        T.getLine
 
     if  | T.null name -> mergeIdentity identity
         | otherwise -> do
-            Just secret <- loadKey public
-            maybe (error "created invalid identity") return . validateIdentity =<<
-                wrappedStore st =<< sign secret =<< wrappedStore st (emptyIdentityData public)
+            secret <- loadKey public
+            maybe (throwError "created invalid identity") return . validateIdentity =<<
+                mstore =<< sign secret =<< mstore (emptyIdentityData public)
                 { iddPrev = toList $ idDataF identity
                 , iddName = Just name
                 }
diff --git a/src/Storage.hs b/src/Storage.hs
index d5d14e3..69e8ab6 100644
--- a/src/Storage.hs
+++ b/src/Storage.hs
@@ -25,6 +25,8 @@ module Storage (
     WatchedHead,
     watchHead, watchHeadWith, unwatchHead,
 
+    MonadStorage(..),
+
     Storable(..), ZeroStorable(..),
     StorableText(..), StorableDate(..), StorableUUID(..),
 
@@ -41,7 +43,7 @@ module Storage (
     loadZRef,
 
     Stored,
-    fromStored, storedRef, storedStorage,
+    fromStored, storedRef,
     wrappedStore, wrappedLoad,
     copyStored,
 
@@ -525,6 +527,22 @@ unwatchHead (WatchedHead st wid _) = do
         StorageMemory { memWatchers = mvar } -> modifyMVar_ mvar $ return . delWatcher
 
 
+class Monad m => MonadStorage m where
+    getStorage :: m Storage
+    mstore :: Storable a => a -> m (Stored a)
+
+    default mstore :: MonadIO m => Storable a => a -> m (Stored a)
+    mstore x = do
+        st <- getStorage
+        wrappedStore st x
+
+instance MonadIO m => MonadStorage (ReaderT Storage m) where
+    getStorage = ask
+
+instance MonadIO m => MonadStorage (ReaderT (Head a) m) where
+    getStorage = asks $ headStorage
+
+
 class Storable a where
     store' :: a -> Store
     load' :: Load a
@@ -862,9 +880,6 @@ fromStored (Stored _ x) = x
 storedRef :: Stored a -> Ref
 storedRef (Stored ref _) = ref
 
-storedStorage :: Stored a -> Storage
-storedStorage (Stored (Ref st _) _) = st
-
 wrappedStore :: MonadIO m => Storable a => Storage -> a -> m (Stored a)
 wrappedStore st x = do ref <- liftIO $ store st x
                        return $ Stored ref x
diff --git a/src/Storage/Internal.hs b/src/Storage/Internal.hs
index 402d924..b68d0f7 100644
--- a/src/Storage/Internal.hs
+++ b/src/Storage/Internal.hs
@@ -175,6 +175,9 @@ instance Eq (Stored' c a) where
 instance Ord (Stored' c a) where
     compare (Stored r1 _) (Stored r2 _) = compare (refDigest r1) (refDigest r2)
 
+storedStorage :: Stored' c a -> Storage' c
+storedStorage (Stored (Ref st _) _) = st
+
 
 type Complete = Identity
 type Partial = Either RefDigest
diff --git a/src/Storage/Key.hs b/src/Storage/Key.hs
index 28fc989..7d36da3 100644
--- a/src/Storage/Key.hs
+++ b/src/Storage/Key.hs
@@ -1,10 +1,11 @@
 module Storage.Key (
     KeyPair(..),
-    storeKey, loadKey,
+    storeKey, loadKey, loadKeyMb,
 ) where
 
 import Control.Concurrent.MVar
 import Control.Monad
+import Control.Monad.Except
 
 import Data.ByteArray
 import qualified Data.ByteString.Char8 as BC
@@ -34,8 +35,11 @@ storeKey key = do
          StorageDir { dirPath = 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
+loadKey :: (KeyPair sec pub, MonadIO m, MonadError String m) => Stored pub -> m sec
+loadKey = maybe (throwError "secret key not found") return <=< loadKeyMb
+
+loadKeyMb :: (KeyPair sec pub, MonadIO m) => Stored pub -> m (Maybe sec)
+loadKeyMb spub = liftIO $ do
     case stBacking $ storedStorage spub of
          StorageDir { dirPath = dir } -> tryIOError (BC.readFile (keyFilePath dir spub)) >>= \case
              Right kdata -> return $ keyFromData (convert kdata) spub
diff --git a/src/Sync.hs b/src/Sync.hs
index b1c0ab0..dd801b5 100644
--- a/src/Sync.hs
+++ b/src/Sync.hs
@@ -27,7 +27,7 @@ instance Service SyncService where
                 let current = sort $ lsShared $ fromStored ls
                     updated = filterAncestors (added : current)
                 if current /= updated
-                   then wrappedStore (storedStorage ls) (fromStored ls) { lsShared = updated }
+                   then mstore (fromStored ls) { lsShared = updated }
                    else return ls
 
     serviceNewPeer = notifyPeer . lsShared . fromStored =<< svcGetLocal
@@ -43,4 +43,4 @@ notifyPeer shared = do
     self <- svcSelf
     when (finalOwner pid `sameIdentity` finalOwner self) $ do
         forM_ shared $ \sh ->
-            replyStoredRef =<< (wrappedStore (storedStorage sh) . SyncPacket) sh
+            replyStoredRef =<< (mstore . SyncPacket) sh
diff --git a/src/Test.hs b/src/Test.hs
index a506345..3f59239 100644
--- a/src/Test.hs
+++ b/src/Test.hs
@@ -426,31 +426,28 @@ cmdWatchSharedIdentity = do
 cmdUpdateLocalIdentity :: Command
 cmdUpdateLocalIdentity = do
     [name] <- asks tiParams
-    updateLocalHead_ $ \ls -> liftIO $ do
+    updateLocalHead_ $ \ls -> do
         Just identity <- return $ validateIdentity $ lsIdentity $ fromStored ls
-        let st = storedStorage ls
-            public = idKeyIdentity identity
+        let public = idKeyIdentity identity
 
-        Just secret <- loadKey public
+        secret <- loadKey public
         nidata <- maybe (error "created invalid identity") (return . idData) . validateIdentity =<<
-            wrappedStore st =<< sign secret =<< wrappedStore st (emptyIdentityData public)
+            mstore =<< sign secret =<< mstore (emptyIdentityData public)
             { iddPrev = toList $ idDataF identity
             , iddName = Just name
             }
-        wrappedStore st $ (fromStored ls) { lsIdentity = nidata }
+        mstore (fromStored ls) { lsIdentity = nidata }
 
 cmdUpdateSharedIdentity :: Command
 cmdUpdateSharedIdentity = do
     [name] <- asks tiParams
     updateLocalHead_ $ updateSharedState_ $ \case
         Nothing -> throwError "no existing shared identity"
-        Just identity -> liftIO $ do
-            let st = storedStorage $ head $ idDataF identity
-                public = idKeyIdentity identity
-
-            Just secret <- loadKey public
+        Just identity -> do
+            let public = idKeyIdentity identity
+            secret <- loadKey public
             maybe (error "created invalid identity") (return . Just . toComposedIdentity) . validateIdentity =<<
-                wrappedStore st =<< sign secret =<< wrappedStore st (emptyIdentityData public)
+                mstore =<< sign secret =<< mstore (emptyIdentityData public)
                 { iddPrev = toList $ idDataF identity
                 , iddName = Just name
                 }