Simplify "mktorrent get" command

2 files changed, 11 insertions, 2 deletions

diff --git a/src/Network/BitTorrent/Exchange.hs b/src/Network/BitTorrent/Exchange.hs
index ce71e286..143bf090 100644
--- a/src/Network/BitTorrent/Exchange.hs
+++ b/src/Network/BitTorrent/Exchange.hs
@@ -27,6 +27,7 @@ module Network.BitTorrent.Exchange
 
          -- * Connections
        , connect
+       , connectSink
        ) where
 
 import Network.BitTorrent.Exchange.Manager
diff --git a/src/Network/BitTorrent/Exchange/Session.hs b/src/Network/BitTorrent/Exchange/Session.hs
index 30b7ed0e..2bd275bd 100644
--- a/src/Network/BitTorrent/Exchange/Session.hs
+++ b/src/Network/BitTorrent/Exchange/Session.hs
@@ -16,6 +16,7 @@ module Network.BitTorrent.Exchange.Session
 
          -- * Connection Set
        , connect
+       , connectSink
        , establish
 
          -- * Query
@@ -29,12 +30,13 @@ import Control.Concurrent.Chan.Split as CS
 import Control.Concurrent.STM
 import Control.Exception hiding (Handler)
 import Control.Lens
+import Control.Monad as M
 import Control.Monad.Logger
 import Control.Monad.Reader
 import Data.ByteString as BS
 import Data.ByteString.Lazy as BL
-import Data.Conduit
-import Data.Conduit.List as CL (iterM)
+import Data.Conduit as C
+import Data.Conduit.List as C
 import Data.Map as M
 import Data.Monoid
 import Data.Set  as S
@@ -333,6 +335,12 @@ establish :: PendingConnection -> Session -> IO ()
 establish conn = runConnection (acceptWire conn) (closePending conn)
                  (pendingPeer conn)
 
+-- | Conduit version of 'connect'.
+connectSink :: MonadIO m => Session -> Sink [PeerAddr IPv4] m ()
+connectSink s = C.mapM_ (liftIO . connectBatch)
+  where
+    connectBatch = M.mapM_ (\ addr -> connect (IPv4 <$> addr) s)
+
 -- | Why do we need this message?
 type BroadcastMessage = ExtendedCaps -> Message