4 files changed, 15 insertions, 4 deletions

diff --git a/bench/Main.hs b/bench/Main.hs
index 269e60f5..6e8a0ce3 100644
--- a/bench/Main.hs
+++ b/bench/Main.hs
@@ -7,10 +7,15 @@ import Criterion.Main
 import Data.ByteString (ByteString)
 import qualified Data.ByteString as B
 import Data.Serialize
+import Network
+
 import Network.BitTorrent as BT
 import Data.Bitfield as BT
 
 
+instance NFData PortNumber where
+  rnf = rnf . (fromIntegral :: PortNumber -> Int)
+
 instance NFData BlockIx where
   rnf (BlockIx a b c) = a `deepseq` b `deepseq` rnf c
 
diff --git a/network-bittorrent.cabal b/network-bittorrent.cabal
index 339dc376..0a89603a 100644
--- a/network-bittorrent.cabal
+++ b/network-bittorrent.cabal
@@ -127,6 +127,7 @@ benchmark benchmarks
   build-depends:       base == 4.*
                      , bytestring >= 0.10.2.0
                      , cereal
+                     , network
 
                      , criterion
                      , deepseq
diff --git a/src/Network/BitTorrent/PeerWire/Message.hs b/src/Network/BitTorrent/PeerWire/Message.hs
index 1c520221..5544dca7 100644
--- a/src/Network/BitTorrent/PeerWire/Message.hs
+++ b/src/Network/BitTorrent/PeerWire/Message.hs
@@ -6,13 +6,13 @@ module Network.BitTorrent.PeerWire.Message
 
 import Control.Applicative
 import qualified Data.ByteString as B
-
 import Data.Serialize
+import Network
 
 import Network.BitTorrent.PeerWire.Block
 import Data.Bitfield
 
-import Data.Array
+
 
 -- | Messages used in communication between peers.
 --
@@ -48,7 +48,7 @@ data Message = KeepAlive
                -- used during "End Game".
              | Cancel   BlockIx
 
-             | Port     Int
+             | Port     PortNumber
 
                -- | BEP 6: Then peer have all pieces it might send the
                --   'HaveAll' message instead of 'Bitfield'
diff --git a/tests/Encoding.hs b/tests/Encoding.hs
index fe0e1d7e..d43045bc 100644
--- a/tests/Encoding.hs
+++ b/tests/Encoding.hs
@@ -13,6 +13,8 @@ import Test.Framework.Providers.QuickCheck2 (testProperty)
 import Test.QuickCheck
 
 import Network.URI
+import Network
+
 
 import Data.Bitfield
 import Data.Torrent
@@ -33,6 +35,9 @@ instance Arbitrary Block where
 
 deriving instance Arbitrary Bitfield
 
+instance Arbitrary PortNumber where
+  arbitrary = fromIntegral <$> (arbitrary :: Gen Word16)
+
 instance Arbitrary Message where
   arbitrary = oneof
     [ pure KeepAlive
@@ -45,7 +50,7 @@ instance Arbitrary Message where
     , Request <$> arbitrary
     , Piece <$> arbitrary
     , Cancel <$> arbitrary
-    , Port <$> choose (0, fromIntegral (maxBound :: Word16))
+    , Port <$> arbitrary
     ]
 
 instance Arbitrary PeerID where