7 files changed, 22 insertions, 26 deletions

diff --git a/src/Data/Torrent/Bitfield.hs b/src/Data/Torrent/Bitfield.hs
index 17278a2c..176ba2c0 100644
--- a/src/Data/Torrent/Bitfield.hs
+++ b/src/Data/Torrent/Bitfield.hs
@@ -55,10 +55,8 @@ module Data.Torrent.Bitfield
        , fromBitmap, toBitmap
        , toList
 
-#if  defined (TESTING)
          -- * Debug
        , mkBitfield
-#endif
        ) where
 
 import Control.Monad
diff --git a/src/Network/BitTorrent/Core.hs b/src/Network/BitTorrent/Core.hs
index b4329598..990a5975 100644
--- a/src/Network/BitTorrent/Core.hs
+++ b/src/Network/BitTorrent/Core.hs
@@ -8,5 +8,6 @@
 --   Re-export every @Network.BitTorrent.Core.*@ module.
 --
 module Network.BitTorrent.Core (module Core) where
-import Network.BitTorrent.Core.PeerId   as Core
-import Network.BitTorrent.Core.PeerAddr as Core
\ No newline at end of file
+import Network.BitTorrent.Core.Fingerprint as Core
+import Network.BitTorrent.Core.PeerId      as Core
+import Network.BitTorrent.Core.PeerAddr    as Core
\ No newline at end of file
diff --git a/src/Network/BitTorrent/Exchange/Message.hs b/src/Network/BitTorrent/Exchange/Message.hs
index 85ad76d6..b879e212 100644
--- a/src/Network/BitTorrent/Exchange/Message.hs
+++ b/src/Network/BitTorrent/Exchange/Message.hs
@@ -81,18 +81,15 @@ import Data.Ord
 import Data.Serialize as S
 import Data.Text as T
 import Data.Typeable
-import Data.Tuple
 import Data.Word
 import Network
 import Network.Socket hiding (KeepAlive)
-import Network.Socket.ByteString as BS
 import Text.PrettyPrint as PP
 import Text.PrettyPrint.Class
 
 import Data.Torrent.Bitfield
 import Data.Torrent.InfoHash
-import Network.BitTorrent.Core.PeerId
-import Network.BitTorrent.Core.PeerAddr ()
+import Network.BitTorrent.Core
 import Network.BitTorrent.Exchange.Block
 
 {-----------------------------------------------------------------------
diff --git a/src/Network/BitTorrent/Exchange/Session.hs b/src/Network/BitTorrent/Exchange/Session.hs
index d2a9aaaf..c24e2b0b 100644
--- a/src/Network/BitTorrent/Exchange/Session.hs
+++ b/src/Network/BitTorrent/Exchange/Session.hs
@@ -15,7 +15,7 @@ import Text.PrettyPrint
 
 import Data.Torrent.Bitfield
 import Data.Torrent.InfoHash
-import Network.BitTorrent.Core.PeerAddr
+import Network.BitTorrent.Core
 import Network.BitTorrent.Exchange.Message
 import Network.BitTorrent.Exchange.Status
 
diff --git a/src/Network/BitTorrent/Exchange/Wire.hs b/src/Network/BitTorrent/Exchange/Wire.hs
index dd77a915..68c9b355 100644
--- a/src/Network/BitTorrent/Exchange/Wire.hs
+++ b/src/Network/BitTorrent/Exchange/Wire.hs
@@ -38,8 +38,7 @@ import Text.PrettyPrint as PP hiding (($$), (<>))
 import Text.PrettyPrint.Class
 
 import Data.Torrent.InfoHash
-import Network.BitTorrent.Core.PeerId
-import Network.BitTorrent.Core.PeerAddr
+import Network.BitTorrent.Core
 import Network.BitTorrent.Exchange.Message
 
 
@@ -147,15 +146,21 @@ connectToPeer p = do
 
 type Wire = ConduitM Message Message (ReaderT Connection IO)
 
-validate :: Wire ()
-validate = do
-  mmsg <- await
-  case mmsg of
-    Nothing  -> return ()
-    Just msg -> do
-      valid <- lift $ asks (`isAllowed` msg)
-      if valid then yield msg else error "TODO"
-
+validate :: ChannelSide -> Wire ()
+validate side = await >>= maybe (return ()) yieldCheck
+  where
+    yieldCheck msg = do
+      caps <- lift $ asks connCaps
+      case requires msg of
+        Nothing  -> return ()
+        Just ext
+          | allowed caps ext -> yield msg
+          |     otherwise    -> monadThrow $ ProtocolError $ InvalidMessage side ext
+
+validate' action = do
+  validate RemotePeer
+  action
+  validate ThisPeer
 
 runWire :: Wire () -> Socket -> Connection -> IO ()
 runWire action sock = runReaderT $
@@ -173,8 +178,6 @@ sendMessage msg = do
 recvMessage :: Wire Message
 recvMessage = undefined
 
-
-
 extendedHandshake :: Wire ()
 extendedHandshake = undefined
 
diff --git a/src/Network/BitTorrent/Tracker/Message.hs b/src/Network/BitTorrent/Tracker/Message.hs
index e8e4cf0e..212181b9 100644
--- a/src/Network/BitTorrent/Tracker/Message.hs
+++ b/src/Network/BitTorrent/Tracker/Message.hs
@@ -88,8 +88,7 @@ import Text.Read (readMaybe)
 
 import Data.Torrent.InfoHash
 import Data.Torrent.Progress
-import Network.BitTorrent.Core.PeerId
-import Network.BitTorrent.Core.PeerAddr
+import Network.BitTorrent.Core
 
 
 {-----------------------------------------------------------------------
diff --git a/src/Network/BitTorrent/Tracker/Wai.hs b/src/Network/BitTorrent/Tracker/Wai.hs
index 51fd8153..770816b4 100644
--- a/src/Network/BitTorrent/Tracker/Wai.hs
+++ b/src/Network/BitTorrent/Tracker/Wai.hs
@@ -12,8 +12,6 @@ import Network.HTTP.Types
 import Network.Wai
 
 import Data.Torrent.Progress
-import Network.BitTorrent.Core.PeerId
-import Network.BitTorrent.Core.PeerAddr
 import Network.BitTorrent.Tracker.Message