Use proper address family when creating Socket from (PeerAddr IP)

2 files changed, 19 insertions, 15 deletions

diff --git a/src/Network/BitTorrent/Core/PeerAddr.hs b/src/Network/BitTorrent/Core/PeerAddr.hs
index e3e09fd6..63885144 100644
--- a/src/Network/BitTorrent/Core/PeerAddr.hs
+++ b/src/Network/BitTorrent/Core/PeerAddr.hs
@@ -23,6 +23,7 @@ module Network.BitTorrent.Core.PeerAddr
          PeerAddr(..)
        , defaultPorts
        , peerSockAddr
+       , peerSocket
 
          -- * Peer storage
        , PeerStore
@@ -296,14 +297,24 @@ defaultPorts =  [6881..6889]
 _resolvePeerAddr :: (IPAddress i) => PeerAddr HostName -> PeerAddr i
 _resolvePeerAddr = undefined
 
--- | Convert peer info from tracker or DHT announce query response to
--- socket address. Usually used to intiate connection between peers.
---
+_peerSockAddr :: PeerAddr IP -> (Family, SockAddr)
+_peerSockAddr PeerAddr {..} =
+    case peerHost of
+          IPv4 ipv4 ->
+              (AF_INET, SockAddrInet peerPort (toHostAddress  ipv4))
+          IPv6 ipv6 ->
+              (AF_INET6, SockAddrInet6 peerPort 0 (toHostAddress6 ipv6) 0)
+
 peerSockAddr :: PeerAddr IP -> SockAddr
-peerSockAddr PeerAddr {..} =
-  case peerHost of
-    IPv4 ipv4 -> SockAddrInet  peerPort   (toHostAddress  ipv4)
-    IPv6 ipv6 -> SockAddrInet6 peerPort 0 (toHostAddress6 ipv6) 0
+peerSockAddr = snd . _peerSockAddr
+
+-- | Create a socket connected to the address specified in a peerAddr
+peerSocket :: SocketType -> PeerAddr IP -> IO Socket
+peerSocket socketType pa = do
+    let (family, addr) = _peerSockAddr pa
+    sock <- socket family socketType defaultProtocol
+    connect sock addr
+    return sock
 
 {-----------------------------------------------------------------------
 --  Peer storage
diff --git a/src/Network/BitTorrent/Exchange/Wire.hs b/src/Network/BitTorrent/Exchange/Wire.hs
index a3b60b99..31da3f0c 100644
--- a/src/Network/BitTorrent/Exchange/Wire.hs
+++ b/src/Network/BitTorrent/Exchange/Wire.hs
@@ -518,13 +518,6 @@ initiateHandshake sock hs = do
   sendHandshake sock hs
   recvHandshake sock
 
--- | Tries to connect to peer using reasonable default parameters.
-connectToPeer :: PeerAddr IP -> IO Socket
-connectToPeer p = do
-  sock <- socket AF_INET Stream Network.Socket.defaultProtocol
-  connect sock (peerSockAddr p)
-  return sock
-
 {-----------------------------------------------------------------------
 --  Wire
 -----------------------------------------------------------------------}
@@ -648,7 +641,7 @@ reconnect = undefined
 --
 connectWire :: Handshake -> PeerAddr IP -> ExtendedCaps -> Wire () -> IO ()
 connectWire hs addr extCaps wire =
-  bracket (connectToPeer addr) close $ \ sock -> do
+  bracket (peerSocket Stream addr) close $ \ sock -> do
     hs' <- initiateHandshake sock hs
 
     Prelude.mapM_ (\(t,e) -> unless t $ throwIO $ ProtocolError e) [