2 files changed, 34 insertions, 20 deletions

diff --git a/src/Network/BitTorrent/Core/PeerAddr.hs b/src/Network/BitTorrent/Core/PeerAddr.hs
index 63ae04b9..b9d4878e 100644
--- a/src/Network/BitTorrent/Core/PeerAddr.hs
+++ b/src/Network/BitTorrent/Core/PeerAddr.hs
@@ -133,6 +133,20 @@ instance BEncode IPv6 where
   fromBEncode = ipFromBEncode
   {-# INLINE fromBEncode #-}
 
+instance Serialize IP where
+    put (IPv4 ip) = put ip
+    put (IPv6 ip) = put ip
+
+    -- | When 'get'ing an IP it must be 'isolate'd to the appropriate number of
+    -- bytes since we have no other way of telling which address type we are
+    -- trying to parse
+    get = do
+      n <- remaining
+      case n of
+        4  -> IPv4 <$> get
+        16 -> IPv6 <$> get
+        _ -> fail "Wrong number of bytes remaining to parse IP"
+
 instance Serialize IPv4 where
     put = putWord32host    .  toHostAddress
     get = fromHostAddress <$> getWord32host
diff --git a/src/Network/BitTorrent/Exchange/Wire.hs b/src/Network/BitTorrent/Exchange/Wire.hs
index 109f6551..c0658961 100644
--- a/src/Network/BitTorrent/Exchange/Wire.hs
+++ b/src/Network/BitTorrent/Exchange/Wire.hs
@@ -432,8 +432,9 @@ data Connection = Connection
     -- future.
     connProtocol     :: !ProtocolName
 
-    -- | A set of enabled extensions. This value used to check if a
-    -- message is allowed to be sent or received.
+    -- | Set of enabled core extensions, i.e. the pre BEP10 extension
+    -- mecahnism. This value is used to check if a message is allowed to be sent
+    -- or received.
   , connCaps         :: !Caps
 
     -- | /Both/ peers handshaked with this infohash. A connection can
@@ -450,9 +451,9 @@ data Connection = Connection
     -- |
   , connOptions      :: !Options
 
-    -- | If @not (allowed ExtExtended connCaps)@ then this set is
-    -- always empty. Otherwise it has extension protocol 'MessageId'
-    -- map.
+    -- | 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.
   , connExtCaps      :: !(IORef ExtendedCaps)
 
     -- | Current extended handshake information from the remote peer
@@ -632,9 +633,9 @@ rehandshake caps = undefined
 reconnect :: Wire ()
 reconnect = undefined
 
--- | Initiate 'Wire' connection and handshake with a peer. This
--- function will also do extension protocol handshake if 'ExtExtended'
--- is enabled on both sides.
+-- | Initiate 'Wire' connection and handshake with a peer. This function will
+-- also do the BEP10 extension protocol handshake if 'ExtExtended' is enabled on
+-- both sides.
 --
 -- This function can throw 'WireFailure' exception.
 --
@@ -643,20 +644,19 @@ connectWire hs addr extCaps wire =
   bracket (connectToPeer addr) close $ \ sock -> do
     hs' <- initiateHandshake sock hs
 
-    unless (def           == hsProtocol hs') $ do
-      throwIO $ ProtocolError $ InvalidProtocol (hsProtocol hs')
-
-    unless (hsProtocol hs == hsProtocol hs') $ do
-      throwIO $ ProtocolError $ UnexpectedProtocol (hsProtocol hs')
-
-    unless (hsInfoHash hs == hsInfoHash hs') $ do
-      throwIO $ ProtocolError $ UnexpectedTopic (hsInfoHash hs')
-
-    unless (hsPeerId hs' == fromMaybe (hsPeerId hs') (peerId addr)) $ do
-      throwIO $ ProtocolError $ UnexpectedPeerId (hsPeerId hs')
+    Prelude.mapM_ (\(t,e) -> unless t $ throwIO $ ProtocolError e) [
+      (def == hsProtocol hs'
+      , InvalidProtocol $ hsProtocol hs'),
+      (hsProtocol hs == hsProtocol hs'
+      , UnexpectedProtocol $ hsProtocol hs'),
+      (hsInfoHash hs == hsInfoHash hs'
+      , UnexpectedTopic $ hsInfoHash hs'),
+      (hsPeerId hs' == fromMaybe (hsPeerId hs') (peerId addr)
+      , UnexpectedPeerId $ hsPeerId hs')
+      ]
 
     let caps = hsReserved hs <> hsReserved hs'
-    let wire' = if ExtExtended `allowed` caps
+        wire' = if ExtExtended `allowed` caps
                 then extendedHandshake extCaps >> wire
                 else wire