Make PeerAddr polimorphic in its address type

2 files changed, 8 insertions, 7 deletions

diff --git a/src/Network/BitTorrent/Exchange/Assembler.hs b/src/Network/BitTorrent/Exchange/Assembler.hs
index 5dc7c5ca..aa009f49 100644
--- a/src/Network/BitTorrent/Exchange/Assembler.hs
+++ b/src/Network/BitTorrent/Exchange/Assembler.hs
@@ -65,6 +65,7 @@ import Data.IntMap.Strict as IM
 import Data.List as L
 import Data.Map as M
 import Data.Maybe
+import Data.IP
 
 import Data.Torrent.Piece
 import Network.BitTorrent.Core
@@ -79,7 +80,7 @@ type PieceMap = IntMap
 
 data Assembler = Assembler
   { -- | A set of blocks that have been 'Request'ed but not yet acked.
-    _inflight :: Map PeerAddr (PieceMap [BlockRange])
+    _inflight :: Map (PeerAddr IP) (PieceMap [BlockRange])
 
     -- | A set of blocks that but not yet assembled.
   , _pending  :: PieceMap Bucket
@@ -114,7 +115,7 @@ allowPiece pix a @ Assembler {..}  = over pending (IM.insert pix bkt) a
   where
     bkt = B.empty (piPieceLength info)
 
-allowedSet :: PeerAddr -> Assembler -> [BlockIx]
+allowedSet :: (PeerAddr IP) -> Assembler -> [BlockIx]
 allowedSet = undefined
 
 --inflight :: PeerAddr -> BlockIx -> Assembler -> Assembler
@@ -123,7 +124,7 @@ allowedSet = undefined
 --  You should check if a returned by peer block is actually have
 -- been requested and in-flight. This is needed to avoid "I send
 -- random corrupted block" attacks.
-insert :: PeerAddr -> Block a -> Assembler -> Assembler
+insert :: (PeerAddr IP) -> Block a -> Assembler -> Assembler
 insert = undefined
 
 {-
@@ -156,4 +157,4 @@ inserta :: Block a
         -> (PieceMap [Block a], Maybe (Piece a))
 inserta = undefined
 
--}
\ No newline at end of file
+-}
diff --git a/src/Network/BitTorrent/Exchange/Wire.hs b/src/Network/BitTorrent/Exchange/Wire.hs
index ae9babb3..27b4be12 100644
--- a/src/Network/BitTorrent/Exchange/Wire.hs
+++ b/src/Network/BitTorrent/Exchange/Wire.hs
@@ -505,7 +505,7 @@ initiateHandshake sock hs = do
   recvHandshake sock
 
 -- | Tries to connect to peer using reasonable default parameters.
-connectToPeer :: PeerAddr -> IO Socket
+connectToPeer :: (IPAddress i) => PeerAddr i -> IO Socket
 connectToPeer p = do
   sock <- socket AF_INET Stream Network.Socket.defaultProtocol
   connect sock (peerSockAddr p)
@@ -628,7 +628,7 @@ reconnect = undefined
 --
 -- This function can throw 'WireFailure' exception.
 --
-connectWire :: Handshake -> PeerAddr -> ExtendedCaps -> Wire () -> IO ()
+connectWire :: (IPAddress i) => Handshake -> PeerAddr i -> ExtendedCaps -> Wire () -> IO ()
 connectWire hs addr extCaps wire =
   bracket (connectToPeer addr) close $ \ sock -> do
     hs' <- initiateHandshake sock hs
@@ -673,7 +673,7 @@ connectWire hs addr extCaps wire =
 --
 --   This function can throw 'WireFailure' exception.
 --
-acceptWire :: Socket -> PeerAddr -> Wire () -> IO ()
+acceptWire :: (IPAddress i) => Socket -> PeerAddr i -> Wire () -> IO ()
 acceptWire sock peerAddr wire = do
   bracket (return sock) close $ \ _ -> do
     error "acceptWire: not implemented"