5 files changed, 23 insertions, 4 deletions
diff --git a/src/Network.hs b/src/Network.hs
index c9a2d8b..7195129 100644
--- a/src/Network.hs
+++ b/src/Network.hs
@@ -4,7 +4,7 @@ module Network (
     getNextPeerChange,
     ServerOptions(..), serverIdentity, defaultServerOptions,
 
-    Peer, peerServer,
+    Peer, peerServer, peerStorage,
     PeerAddress(..), peerAddress,
     PeerIdentity(..), peerIdentity,
     PeerChannel(..),
@@ -101,7 +101,7 @@ data Peer = Peer
     , peerServer_ :: Server
     , peerIdentityVar :: TVar PeerIdentity
     , peerChannel :: TVar PeerChannel
-    , peerStorage :: Storage
+    , peerStorage_ :: Storage
     , peerInStorage :: PartialStorage
     , peerOutQueue :: TQueue (Bool, [TransportHeaderItem], TransportPacket)
     , peerSentPackets :: TVar [SentPacket]
@@ -120,6 +120,9 @@ data SentPacket = SentPacket
 peerServer :: Peer -> Server
 peerServer = peerServer_
 
+peerStorage :: Peer -> Storage
+peerStorage = peerStorage_
+
 instance Eq Peer where
     (==) = (==) `on` peerIdentityVar
 
diff --git a/src/Network.hs-boot b/src/Network.hs-boot
index 1ec6daa..f251e5e 100644
--- a/src/Network.hs-boot
+++ b/src/Network.hs-boot
@@ -1,4 +1,8 @@
 module Network where
 
+import Storage
+
 data Server
 data Peer
+
+peerStorage :: Peer -> Storage
diff --git a/src/Service.hs b/src/Service.hs
index 7d5c7e7..22c983b 100644
--- a/src/Service.hs
+++ b/src/Service.hs
@@ -111,6 +111,9 @@ data ServiceHandlerState s = ServiceHandlerState
 newtype ServiceHandler s a = ServiceHandler (ReaderT (ServiceInput s) (WriterT [ServiceReply s] (StateT (ServiceHandlerState s) (ExceptT String IO))) a)
     deriving (Functor, Applicative, Monad, MonadReader (ServiceInput s), MonadWriter [ServiceReply s], MonadState (ServiceHandlerState s), MonadError String, MonadIO)
 
+instance MonadStorage (ServiceHandler s) where
+    getStorage = asks $ peerStorage . svcPeer
+
 instance MonadHead LocalState (ServiceHandler s) where
     updateLocalHead f = do
         (ls, x) <- liftIO . f =<< gets svcLocal
diff --git a/src/State.hs b/src/State.hs
index 358d958..e112aca 100644
--- a/src/State.hs
+++ b/src/State.hs
@@ -2,7 +2,7 @@ module State (
     LocalState(..),
     SharedState, SharedType(..),
     SharedTypeID, mkSharedTypeID,
-    MonadHead(..),
+    MonadStorage(..), MonadHead(..),
 
     loadLocalStateHead,
     updateLocalState, updateLocalState_,
@@ -81,9 +81,15 @@ instance SharedType (Maybe ComposedIdentity) where
     sharedTypeID _ = mkSharedTypeID "0c6c1fe0-f2d7-4891-926b-c332449f7871"
 
 
-class MonadHead a m where
+class Monad m => MonadStorage m where
+    getStorage :: m Storage
+
+class MonadStorage m => MonadHead a m where
     updateLocalHead :: (Stored a -> IO (Stored a, b)) -> m b
 
+instance Monad m => MonadStorage (ReaderT (Head a) m) where
+    getStorage = asks $ refStorage . headRef
+
 instance (HeadType a, MonadIO m) => MonadHead a (ReaderT (Head a) m) where
     updateLocalHead f = do
         h <- ask
diff --git a/src/Test.hs b/src/Test.hs
index 9c5319b..8bd34ea 100644
--- a/src/Test.hs
+++ b/src/Test.hs
@@ -167,6 +167,9 @@ instance MonadFail CommandM where
 instance MonadRandom CommandM where
     getRandomBytes = liftIO . getRandomBytes
 
+instance MonadStorage CommandM where
+    getStorage = asks tiStorage
+
 instance MonadHead LocalState CommandM where
     updateLocalHead f = do
         Just h <- gets tsHead