1 files changed, 45 insertions, 13 deletions

diff --git a/src/Network/BitTorrent/Address.hs b/src/Network/BitTorrent/Address.hs
index a72a8a07..ed02f7ff 100644
--- a/src/Network/BitTorrent/Address.hs
+++ b/src/Network/BitTorrent/Address.hs
@@ -26,6 +26,8 @@ module Network.BitTorrent.Address
        ( -- * Address
          Address (..)
        , fromAddr
+       , PortNumber
+       , SockAddr
 
          -- ** IP
        , IPv4
@@ -81,16 +83,18 @@ module Network.BitTorrent.Address
 
          -- * Utils
        , libUserAgent
+       , sockAddrPort
+       , getBindAddress
        ) where
 
 import Control.Applicative
 import Control.Monad
+import Control.Exception (onException)
 import Data.BEncode as BE
-import Data.BEncode   as BS
 import Data.BEncode.BDict (BKey)
 import Data.Bits
-import Data.ByteString as BS
-import Data.ByteString.Internal as BS
+import qualified Data.ByteString as BS
+import qualified Data.ByteString.Internal as BS
 import Data.ByteString.Base16 as Base16
 import Data.ByteString.Char8 as BC
 import Data.ByteString.Char8 as BS8
@@ -167,19 +171,21 @@ setPort port (SockAddrInet6 _ f h s) = SockAddrInet6 port f h s
 setPort _    addr                    = addr
 {-# INLINE setPort #-}
 
-getPort :: SockAddr -> Maybe PortNumber
-getPort (SockAddrInet  p   _  ) = Just p
-getPort (SockAddrInet6 p _ _ _) = Just p
-getPort _                       = Nothing
-{-# INLINE getPort #-}
+-- | Obtains the port associated with a socket address
+-- if one is associated with it.
+sockAddrPort :: SockAddr -> Maybe PortNumber
+sockAddrPort (SockAddrInet  p   _  ) = Just p
+sockAddrPort (SockAddrInet6 p _ _ _) = Just p
+sockAddrPort _                       = Nothing
+{-# INLINE sockAddrPort #-}
 
 instance Address a => Address (NodeAddr a) where
   toSockAddr NodeAddr {..} = setPort nodePort $ toSockAddr nodeHost
-  fromSockAddr sa = NodeAddr <$> fromSockAddr sa <*> getPort sa
+  fromSockAddr sa = NodeAddr <$> fromSockAddr sa <*> sockAddrPort sa
 
 instance Address a => Address (PeerAddr a) where
   toSockAddr PeerAddr {..} = setPort peerPort $ toSockAddr peerHost
-  fromSockAddr sa = PeerAddr Nothing <$> fromSockAddr sa <*> getPort sa
+  fromSockAddr sa = PeerAddr Nothing <$> fromSockAddr sa <*> sockAddrPort sa
 
 {-----------------------------------------------------------------------
 --  Peer id
@@ -366,7 +372,7 @@ instance BEncode PortNumber where
   toBEncode   = toBEncode    .  fromEnum
   fromBEncode = fromBEncode >=> portNumber
     where
-      portNumber :: Integer -> BS.Result PortNumber
+      portNumber :: Integer -> BE.Result PortNumber
       portNumber n
         | 0 <= n && n <= fromIntegral (maxBound :: Word16)
         = pure $ fromIntegral n
@@ -414,7 +420,7 @@ ipToBEncode :: Show i => i -> BValue
 ipToBEncode ip = BString $ BS8.pack $ show ip
 {-# INLINE ipToBEncode #-}
 
-ipFromBEncode :: Read a => BValue -> BS.Result a
+ipFromBEncode :: Read a => BValue -> BE.Result a
 ipFromBEncode (BString (BS8.unpack -> ipStr))
   | Just ip <- readMaybe (ipStr) = pure ip
   |         otherwise            = decodingError $ "IP: " ++ ipStr
@@ -1166,7 +1172,7 @@ fingerprint pid = either (const def) id $ runGet getCI (getPeerId pid)
         'e' -> Fingerprint <$> getBitCometImpl <*> getBitCometVersion
         'F' -> Fingerprint <$> getBitCometImpl <*> getBitCometVersion
         c   -> do
-          c1 <- w2c <$> S.lookAhead getWord8
+          c1 <- BS.w2c <$> S.lookAhead getWord8
           if c1 == 'P'
             then do
                  _ <- getWord8
@@ -1252,3 +1258,29 @@ bep42 addr (NodeId r)
                     bs     -> bs
         where msk | BS.length ip == 4  =  ip4mask
                   | otherwise          =  ip6mask
+
+
+-- | Given a string specifying a port (numeric or service name)
+-- and a flag indicating whether you want to support IPv6, this
+-- function will return a SockAddr to bind to.  If the input
+-- is not understood as a port number, zero will be set in order
+-- to ask the system for an unused port.
+--
+-- TODO: Also interpret local ip address specifications in the input
+-- string.
+getBindAddress :: String -> Bool -> IO SockAddr
+getBindAddress listenPortString enabled6 = do
+    -- Bind addresses for localhost
+    xs <- getAddrInfo (Just (defaultHints {addrFlags = [AI_PASSIVE] }))
+                      Nothing
+                      (Just listenPortString)
+            `onException` return []
+    -- We prefer IPv6 because that can also handle connections from IPv4
+    -- clients...
+    let (x6s,x4s) = partition (\s -> addrFamily s == AF_INET6) xs
+        listenAddr =
+          case if enabled6 then x6s++x4s else x4s of
+            AddrInfo { addrAddress = listenAddr } : _ -> listenAddr
+            _ -> SockAddrInet6 (parsePort listenPortString) 0 iN6ADDR_ANY 0
+                    where parsePort s = fromMaybe 0 $ readMaybe s
+    return listenAddr