4 files changed, 104 insertions, 4 deletions

diff --git a/src/Network/BitTorrent/Core/Node.hs b/src/Network/BitTorrent/Core/Node.hs
index f6ab7d82..f163a895 100644
--- a/src/Network/BitTorrent/Core/Node.hs
+++ b/src/Network/BitTorrent/Core/Node.hs
@@ -39,6 +39,7 @@ import Data.Bits
 import Data.ByteString as BS
 import Data.BEncode as BE
 import Data.Default
+import Data.Hashable
 import Data.IP
 import Data.List as L
 import Data.Ord
@@ -126,6 +127,10 @@ instance BEncode a => BEncode (NodeAddr a) where
   fromBEncode b = uncurry NodeAddr <$> fromBEncode b
   {-# INLINE fromBEncode #-}
 
+instance Hashable a => Hashable (NodeAddr a) where
+  hashWithSalt s NodeAddr {..} = hashWithSalt s (nodeHost, nodePort)
+  {-# INLINE hashWithSalt #-}
+
 -- | Example:
 --
 --   @nodePort \"127.0.0.1:6881\" == 6881@
@@ -160,7 +165,7 @@ instance Eq a => Ord (NodeInfo a) where
 instance Serialize a => Serialize (NodeInfo a) where
   get = NodeInfo <$> get <*> get
   put NodeInfo {..} = put nodeId >> put nodeAddr
-
+{-
 type CompactInfo = ByteString
 
 data NodeList a = CompactNodeList [NodeInfo a]
@@ -176,3 +181,4 @@ encodeCompact = S.runPut . mapM_ put
 
 --encodePeerList :: [PeerAddr] -> [BEncode]
 --encodePeerList = undefined
+-}
\ No newline at end of file
diff --git a/src/Network/BitTorrent/Core/PeerAddr.hs b/src/Network/BitTorrent/Core/PeerAddr.hs
index 261df8d0..9f5c7c5d 100644
--- a/src/Network/BitTorrent/Core/PeerAddr.hs
+++ b/src/Network/BitTorrent/Core/PeerAddr.hs
@@ -38,6 +38,7 @@ import Data.BEncode.BDict (BKey)
 import Data.ByteString.Char8 as BS8
 import Data.Char
 import Data.Default
+import Data.Hashable
 import Data.HashMap.Strict as HM
 import Data.IP
 import Data.List      as L
@@ -81,6 +82,10 @@ instance Serialize PortNumber where
   put = putWord16be . fromIntegral
   {-# INLINE put #-}
 
+instance Hashable PortNumber where
+  hashWithSalt s = hashWithSalt s . fromEnum
+  {-# INLINE hashWithSalt #-}
+
 {-----------------------------------------------------------------------
 --  IP addr
 -----------------------------------------------------------------------}
diff --git a/src/Network/BitTorrent/DHT/Message.hs b/src/Network/BitTorrent/DHT/Message.hs
index 49629755..83413752 100644
--- a/src/Network/BitTorrent/DHT/Message.hs
+++ b/src/Network/BitTorrent/DHT/Message.hs
@@ -30,7 +30,6 @@ module Network.BitTorrent.DHT.Message
 import Control.Applicative
 import Data.BEncode as BE
 import Data.BEncode.BDict
-import Data.ByteString as BS
 import Data.List as L
 import Data.Monoid
 import Data.Serialize as S
@@ -40,14 +39,13 @@ import Network.KRPC
 
 import Data.Torrent.InfoHash
 import Network.BitTorrent.Core
+import Network.BitTorrent.DHT.Token
 import Network.KRPC ()
 
 {-----------------------------------------------------------------------
 -- envelopes
 -----------------------------------------------------------------------}
 
-type Token = ByteString
-
 node_id_key :: BKey
 node_id_key = "id"
 
diff --git a/src/Network/BitTorrent/DHT/Token.hs b/src/Network/BitTorrent/DHT/Token.hs
new file mode 100644
index 00000000..6705f932
--- /dev/null
+++ b/src/Network/BitTorrent/DHT/Token.hs
@@ -0,0 +1,91 @@
+-- | The return value for a query for peers includes an opaque value
+-- known as the \"token.\" For a node to announce that its controlling
+-- peer is downloading a torrent, it must present the token received
+-- from the same queried node in a recent query for peers. When a node
+-- attempts to "announce" a torrent, the queried node checks the token
+-- against the querying node's IP address. This is to prevent
+-- malicious hosts from signing up other hosts for torrents. Since the
+-- token is merely returned by the querying node to the same node it
+-- received the token from, the implementation is not defined. Tokens
+-- must be accepted for a reasonable amount of time after they have
+-- been distributed.
+--
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+module Network.BitTorrent.DHT.Token
+       ( Token
+       , TokenMap
+       , Network.BitTorrent.DHT.Token.tokens
+       , Network.BitTorrent.DHT.Token.lookup
+       , Network.BitTorrent.DHT.Token.member
+       , Network.BitTorrent.DHT.Token.updateInterval
+       , Network.BitTorrent.DHT.Token.update
+       ) where
+
+import Control.Applicative
+import Control.Monad.State
+import Data.BEncode (BEncode)
+import Data.ByteString as BS
+import Data.ByteString.Lazy as BL
+import Data.ByteString.Lazy.Builder as BS
+import Data.Default
+import Data.List as L
+import Data.Hashable
+import Data.String
+import Data.Time
+import System.Random
+
+import Network.BitTorrent.Core
+
+
+type Secret = Int
+
+-- TODO use ShortByteString
+newtype Token = Token BS.ByteString
+  deriving (Show, Eq, BEncode, IsString)
+
+-- | Meaningless token, for testsing purposes only.
+instance Default Token where
+  def = Token "0xdeadbeef"
+
+-- The BitTorrent implementation uses the SHA1 hash of the IP address
+-- concatenated onto a secret, we use hashable instead.
+makeToken :: Hashable a => NodeAddr a -> Secret -> Token
+makeToken n s = Token $ toBS $ hashWithSalt s n
+  where
+    toBS = toStrict . toLazyByteString . int64BE . fromIntegral
+
+-- | Constant space token map based on secret.
+data TokenMap = TokenMap
+  { prevSecret :: {-# UNPACK #-} !Secret
+  , curSecret  :: {-# UNPACK #-} !Secret
+  , generator  :: {-# UNPACK #-} !StdGen
+  } deriving Show
+
+tokens :: Int -> TokenMap
+tokens seed = (`evalState` mkStdGen seed) $
+  TokenMap <$> state next
+           <*> state next
+           <*> get
+
+lookup :: Hashable a => NodeAddr a -> TokenMap -> Token
+lookup addr TokenMap {..} = makeToken addr curSecret
+
+-- | If token is not set 'Network.KRPC.ProtocolError' should be sent
+-- back.
+member :: Hashable a => NodeAddr a -> Token -> TokenMap -> Bool
+member addr token TokenMap {..} = token `L.elem` valid
+  where valid = makeToken addr <$> [curSecret, prevSecret]
+
+-- Secret changes every five minutes and tokens up to ten minutes old
+-- are accepted.
+updateInterval :: NominalDiffTime
+updateInterval = 5 * 60
+
+update :: TokenMap -> TokenMap
+update TokenMap {..} = TokenMap
+    { prevSecret = curSecret
+    , curSecret  = newSecret
+    , generator  = newGen
+    }
+  where
+    (newSecret, newGen) = next generator
\ No newline at end of file