3 files changed, 18 insertions, 18 deletions

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