2 files changed, 60 insertions, 98 deletions

diff --git a/src/Network/BitTorrent/Exchange/Session.hs b/src/Network/BitTorrent/Exchange/Session.hs
index 0d4f3d02..416e00fd 100644
--- a/src/Network/BitTorrent/Exchange/Session.hs
+++ b/src/Network/BitTorrent/Exchange/Session.hs
@@ -110,7 +110,7 @@ newSession logFun addr rootPath dict = do
 closeSession :: Session -> IO ()
 closeSession = undefined
 
-instance MonadIO m => MonadLogger (Connected Session m) where
+instance MonadLogger (Connected Session) where
   monadLoggerLog loc src lvl msg = do
     conn <- ask
     ses  <- asks connSession
@@ -143,9 +143,9 @@ insert addr ses @ Session {..} = do
       let hs = Handshake def caps infohash tpeerId
       chan <- dupChan broadcast
       connectWire ses hs addr ecaps chan $ do
-        conn <- getConnection
+        conn <- ask
 --      liftIO $ modifyMVar_ connections $ pure . M.insert addr conn
-        resizeBitfield (totalPieces storage)
+        lift $ resizeBitfield (totalPieces storage)
         logEvent "Connection established"
         exchange
 --      liftIO $ modifyMVar_ connections $ pure . M.delete addr
@@ -162,12 +162,12 @@ deleteAll = undefined
 
 withStatusUpdates :: StatusUpdates a -> Wire Session a
 withStatusUpdates m = do
-  Session {..} <- getSession
+  Session {..} <- asks connSession
   liftIO $ runStatusUpdates status m
 
 getThisBitfield :: Wire Session Bitfield
 getThisBitfield = do
-  ses <- getSession
+  ses <- asks connSession
   liftIO $ SS.getBitfield (status ses)
 
 readBlock :: BlockIx -> Storage -> IO (Block BL.ByteString)
@@ -181,8 +181,8 @@ readBlock bix @ BlockIx {..} s = do
 
 sendBroadcast :: PeerMessage msg => msg -> Wire Session ()
 sendBroadcast msg = do
-  Session {..} <- getSession
-  ecaps <- getExtCaps
+  Session {..} <- asks connSession
+  ecaps <- use connExtCaps
   liftIO $ writeChan broadcast (envelop ecaps msg)
 
 {-----------------------------------------------------------------------
@@ -191,9 +191,9 @@ sendBroadcast msg = do
 
 fillRequestQueue :: Wire Session ()
 fillRequestQueue = do
-  maxN <- getAdvertisedQueueLength
-  rbf  <- getRemoteBitfield
-  addr <- connRemoteAddr <$> getConnection
+  maxN <- lift $ getAdvertisedQueueLength
+  rbf  <- use connBitfield
+  addr <- asks connRemoteAddr
   blks <- withStatusUpdates $ do
     n <- getRequestQueueLength addr
     scheduleBlocks addr rbf (maxN - n)
@@ -201,13 +201,13 @@ fillRequestQueue = do
 
 tryFillRequestQueue :: Wire Session ()
 tryFillRequestQueue = do
-  allowed <- canDownload <$> getStatus
+  allowed <- canDownload <$> use connStatus
   when allowed $ do
     fillRequestQueue
 
 interesting :: Wire Session ()
 interesting = do
-  addr <- connRemoteAddr <$> getConnection
+  addr <- asks connRemoteAddr
   logMessage  (Status (Interested True))
   sendMessage (Interested True)
   logMessage  (Status (Choking    False))
@@ -220,17 +220,17 @@ interesting = do
 
 handleStatus :: StatusUpdate -> Wire Session ()
 handleStatus s = do
-  updateConnStatus RemotePeer s
+  connStatus %= over remoteStatus (updateStatus s)
   case s of
     Interested _     -> return ()
     Choking    True  -> do
-      addr <- connRemoteAddr <$> getConnection
+      addr <- asks connRemoteAddr
       withStatusUpdates (resetPending addr)
     Choking    False -> tryFillRequestQueue
 
 handleAvailable :: Available -> Wire Session ()
 handleAvailable msg = do
-  updateRemoteBitfield $ case msg of
+  connBitfield %= case msg of
     Have     ix -> BF.insert ix
     Bitfield bf -> const     bf
 
@@ -245,15 +245,15 @@ handleAvailable msg = do
 
 handleTransfer :: Transfer -> Wire Session ()
 handleTransfer (Request bix) = do
-  Session {..} <- getSession
+  Session {..} <- asks connSession
   bitfield <- getThisBitfield
-  upload   <- canUpload <$> getStatus
+  upload   <- canUpload <$> use connStatus
   when (upload && ixPiece bix `BF.member` bitfield) $ do
     blk <- liftIO $ readBlock bix storage
     sendMessage (Piece blk)
 
 handleTransfer (Piece   blk) = do
-  Session {..} <- getSession
+  Session {..} <- asks connSession
   isSuccess <- withStatusUpdates (pushBlock blk storage)
   case isSuccess of
     Nothing -> liftIO $ throwIO $ userError "block is not requested"
diff --git a/src/Network/BitTorrent/Exchange/Wire.hs b/src/Network/BitTorrent/Exchange/Wire.hs
index 64fa3295..e88b3ae5 100644
--- a/src/Network/BitTorrent/Exchange/Wire.hs
+++ b/src/Network/BitTorrent/Exchange/Wire.hs
@@ -18,41 +18,46 @@ module Network.BitTorrent.Exchange.Wire
          Connected
        , Wire
 
-         -- ** Connection
+         -- * Connection
        , Connection
+
+         -- ** Identity
        , connRemoteAddr
-       , connProtocol
-       , connCaps
        , connTopic
        , connRemotePeerId
        , connThisPeerId
+
+         -- ** Capabilities
+       , connProtocol
+       , connCaps
+       , connExtCaps
+       , connRemoteEhs
+
+         -- ** State
+       , connStatus
+       , connBitfield
+
+         -- ** Env
        , connOptions
        , connSession
+       , connStats
 
-         -- ** Setup
+         -- * Setup
        , runWire
        , connectWire
        , acceptWire
        , resizeBitfield
 
-         -- ** Messaging
+         -- * Messaging
        , recvMessage
        , sendMessage
        , filterQueue
        , getAdvertisedQueueLength
 
-         -- ** Query
-       , getConnection
-       , getSession
-       , getStatus
-       , getRemoteBitfield
-       , updateConnStatus
-       , updateRemoteBitfield
-       , getExtCaps
-       , getStats
+         -- * Query
        , getMetadata
 
-         -- ** Exceptions
+         -- * Exceptions
        , ChannelSide   (..)
        , ProtocolError (..)
        , WireFailure   (..)
@@ -60,15 +65,15 @@ module Network.BitTorrent.Exchange.Wire
        , isWireFailure
        , disconnectPeer
 
-         -- ** Stats
+         -- * Stats
        , ByteStats       (..)
        , FlowStats       (..)
        , ConnectionStats (..)
 
-         -- ** Flood detection
+         -- * Flood detection
        , FloodDetector   (..)
 
-         -- ** Options
+         -- * Options
        , Options         (..)
        ) where
 
@@ -449,6 +454,10 @@ data ConnectionState = ConnectionState {
     -- | If @not (allowed ExtExtended connCaps)@ then this set is always
     -- empty. Otherwise it has the BEP10 extension protocol mandated mapping of
     -- 'MessageId' to the message type for the remote peer.
+    --
+    --  Note that this value can change in current session if either
+    --  this or remote peer will initiate rehandshaking.
+    --
     _connExtCaps      :: !ExtendedCaps
 
     -- | Current extended handshake information from the remote peer
@@ -456,6 +465,9 @@ data ConnectionState = ConnectionState {
 
     -- | Various stats about messages sent and received. Stats can be
     -- used to protect /this/ peer against flood attacks.
+    --
+    -- Note that this value will change with the next sent or received
+    -- message.
   , _connStats        :: !ConnectionStats
 
   , _connStatus       :: !ConnectionStatus
@@ -565,70 +577,40 @@ initiateHandshake sock hs = do
 -----------------------------------------------------------------------}
 
 -- | do not expose this so we can change it without breaking api
-newtype Connected s m a = Connected { runConnected :: (ReaderT (Connection s) m a) }
+newtype Connected s a = Connected { runConnected :: (ReaderT (Connection s) IO a) }
     deriving (Functor, Applicative, Monad
              , MonadIO, MonadReader (Connection s), MonadThrow
              )
 
-instance MonadIO m => MonadState ConnectionState (Connected s m) where
+instance MonadState ConnectionState (Connected s) where
     get   = Connected (asks connState) >>= liftIO . readIORef
     put x = Connected (asks connState) >>= liftIO . flip writeIORef x
 
-instance MonadTrans (Connected s) where
-    lift = Connected . lift
-
 -- | A duplex channel connected to a remote peer which keep tracks
 -- connection parameters.
-type Wire s a = ConduitM Message Message (Connected s IO) a
+type Wire s a = ConduitM Message Message (Connected s) a
 
 {-----------------------------------------------------------------------
 --  Query
 -----------------------------------------------------------------------}
 
-setExtCaps :: ExtendedCaps -> Wire s ()
-setExtCaps x = lift $ connExtCaps .= x
-
--- | Get current extended capabilities. Note that this value can
--- change in current session if either this or remote peer will
--- initiate rehandshaking.
-getExtCaps :: Wire s ExtendedCaps
-getExtCaps = lift $ use connExtCaps
-
-setRemoteEhs :: ExtendedHandshake -> Wire s ()
-setRemoteEhs x = lift $ connRemoteEhs .= x
-
-getRemoteEhs :: Wire s ExtendedHandshake
-getRemoteEhs = lift $ use connRemoteEhs
-
--- | Get current stats. Note that this value will change with the next
--- sent or received message.
-getStats :: Wire s ConnectionStats
-getStats = lift $ use connStats
-
--- | See the 'Connection' section for more info.
-getConnection :: Wire s (Connection s)
-getConnection = lift ask
-
-getSession :: Wire s s
-getSession = lift (asks connSession)
-
 -- TODO configurable
 defQueueLength :: Int
 defQueueLength = 1
 
-getAdvertisedQueueLength :: Wire s Int
+getAdvertisedQueueLength :: Connected s Int
 getAdvertisedQueueLength = do
-  ExtendedHandshake {..} <- getRemoteEhs
+  ExtendedHandshake {..} <- use connRemoteEhs
   return $ fromMaybe defQueueLength ehsQueueLength
 
 {-----------------------------------------------------------------------
 --  Wrapper
 -----------------------------------------------------------------------}
 
-putStats :: ChannelSide -> Message -> Connected s IO ()
+putStats :: ChannelSide -> Message -> Connected s ()
 putStats side msg = connStats %= addStats side (stats msg)
 
-validate :: ChannelSide -> Message -> Connected s IO ()
+validate :: ChannelSide -> Message -> Connected s ()
 validate side msg = do
   caps <- asks connCaps
   case requires msg of
@@ -696,8 +678,8 @@ extendedHandshake caps = do
   msg <- recvMessage
   case msg of
     Extended (EHandshake remoteEhs@(ExtendedHandshake {..})) -> do
-      setExtCaps $ ehsCaps <> caps
-      setRemoteEhs remoteEhs
+      connExtCaps   .= (ehsCaps <> caps)
+      connRemoteEhs .= remoteEhs
     _ -> protocolError HandshakeRefused
 
 rehandshake :: ExtendedCaps -> Wire s ()
@@ -776,29 +758,9 @@ acceptWire sock peerAddr wire = do
   bracket (return sock) close $ \ _ -> do
     error "acceptWire: not implemented"
 
-{-----------------------------------------------------------------------
---  Connection Status
------------------------------------------------------------------------}
-
-getStatus :: Wire s ConnectionStatus
-getStatus = lift $ use connStatus
-
-updateConnStatus :: ChannelSide -> StatusUpdate -> Wire s ()
-updateConnStatus side u = lift $ do
-    connStatus %= (over (statusSide side) (updateStatus u))
-  where
-    statusSide ThisPeer   = clientStatus
-    statusSide RemotePeer = remoteStatus
-
-getRemoteBitfield :: Wire s Bitfield
-getRemoteBitfield = lift $ use connBitfield
-
-updateRemoteBitfield :: (Bitfield -> Bitfield) -> Wire s ()
-updateRemoteBitfield f = lift $ connBitfield %= f
-
 -- | Used when size of bitfield becomes known.
-resizeBitfield :: Int -> Wire s ()
-resizeBitfield n = updateRemoteBitfield (adjustSize n)
+resizeBitfield :: Int -> Connected s ()
+resizeBitfield n = connBitfield %= adjustSize n
 
 {-----------------------------------------------------------------------
 --  Metadata exchange
@@ -833,7 +795,7 @@ fetchMetadata = loop 0
 getMetadata :: Wire s InfoDict
 getMetadata = do
   chunks <- fetchMetadata
-  Connection {..} <- getConnection
+  Connection {..} <- ask
   case BE.decode (BS.concat chunks) of
      Right (infodict @ InfoDict {..})
        | connTopic == idInfoHash -> return infodict