1 files changed, 201 insertions, 0 deletions

diff --git a/dht/src/Network/BitTorrent/DHT/Token.hs b/dht/src/Network/BitTorrent/DHT/Token.hs
new file mode 100644
index 00000000..171cc8be
--- /dev/null
+++ b/dht/src/Network/BitTorrent/DHT/Token.hs
@@ -0,0 +1,201 @@
+-- |
+--   Copyright   :  (c) Sam Truzjan 2013
+--   License     :  BSD3
+--   Maintainer  :  pxqr.sta@gmail.com
+--   Stability   :  experimental
+--   Portability :  portable
+--
+-- 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, CPP #-}
+module Network.BitTorrent.DHT.Token
+       ( -- * Token
+         Token
+       , maxInterval
+       , toPaddedByteString
+       , fromPaddedByteString
+
+         -- * Session tokens
+       , TokenMap
+       , SessionTokens
+       , nullSessionTokens
+       , checkToken
+       , grantToken
+
+         -- ** Construction
+       , Network.BitTorrent.DHT.Token.tokens
+
+         -- ** Query
+       , Network.BitTorrent.DHT.Token.lookup
+       , Network.BitTorrent.DHT.Token.member
+
+         -- ** Modification
+       , Network.BitTorrent.DHT.Token.defaultUpdateInterval
+       , Network.BitTorrent.DHT.Token.update
+       ) where
+
+import Control.Arrow
+import Control.Monad.State
+#ifdef VERSION_bencoding
+import Data.BEncode (BEncode)
+#endif
+import Data.ByteString as BS
+import Data.ByteString.Char8 as B8
+import Data.ByteString.Lazy as BL
+import Data.ByteString.Lazy.Builder as BS
+import qualified Data.ByteString.Base16 as Base16
+import Data.Default
+import Data.List as L
+import Data.Hashable
+import Data.String
+import Data.Time
+import System.Random
+import Control.Concurrent.STM
+
+-- TODO use ShortByteString
+
+-- | An opaque value.
+newtype Token = Token BS.ByteString
+  deriving ( Eq, IsString
+#ifdef VERSION_bencoding
+           , BEncode
+#endif
+           )
+
+instance Show Token where
+    show (Token bs) = B8.unpack $ Base16.encode bs
+
+instance Read Token where
+    readsPrec i s = pure $ (Token *** B8.unpack) $ Base16.decode (B8.pack s)
+
+-- | Meaningless token, for testing purposes only.
+instance Default Token where
+  def = makeToken (0::Int) 0
+
+-- | Prepend token with 0x20 bytes to fill the available width.
+--
+-- If n > 8, then this will also guarantee a nonzero token, which is useful for
+-- Tox ping-id values for announce responses.
+toPaddedByteString :: Int -> Token -> BS.ByteString
+toPaddedByteString n (Token bs) = BS.append (BS.replicate (n - BS.length bs) 0x20) bs
+
+fromPaddedByteString :: Int -> BS.ByteString -> Token
+fromPaddedByteString n bs = Token $ BS.drop (n - len) bs
+ where
+    len = BS.length tok where Token tok = def
+
+-- | The secret value used as salt.
+type Secret = Int
+
+-- The BitTorrent implementation uses the SHA1 hash of the IP address
+-- concatenated onto a secret, we use hashable instead.
+makeToken :: Hashable a => a -> Secret -> Token
+makeToken n s = Token $ toBS $ hashWithSalt s n
+  where
+    toBS = toStrict . toLazyByteString . int64BE . fromIntegral
+{-# INLINE makeToken #-}
+
+-- | Constant space 'Node' to 'Token' map based on the secret value.
+data TokenMap = TokenMap
+  { prevSecret :: {-# UNPACK #-} !Secret
+  , curSecret  :: {-# UNPACK #-} !Secret
+  , generator  :: {-# UNPACK #-} !StdGen
+  } deriving Show
+
+-- | A new token map based on the specified seed value. Returned token
+-- map should be periodicatically 'update'd.
+--
+-- Normally, the seed value should vary between invocations of the
+-- client software.
+tokens :: Int -> TokenMap
+tokens seed = (`evalState` mkStdGen seed) $
+  TokenMap <$> state next
+           <*> state next
+           <*> get
+
+-- | Get token for the given node. A token becomes invalid after 2
+-- 'update's.
+--
+-- Typically used to handle find_peers query.
+lookup :: Hashable a => a -> TokenMap -> Token
+lookup addr TokenMap {..} = makeToken addr curSecret
+
+-- | Check if token is valid.
+--
+-- Typically used to handle 'Network.DHT.Mainline.Announce'
+-- query. If token is invalid the 'Network.KRPC.ProtocolError' should
+-- be sent back to the malicious node.
+member :: Hashable a => 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.
+defaultUpdateInterval :: NominalDiffTime
+defaultUpdateInterval = 5 * 60
+
+-- | Update current tokens.
+update :: TokenMap -> TokenMap
+update TokenMap {..} = TokenMap
+    { prevSecret = curSecret
+    , curSecret  = newSecret
+    , generator  = newGen
+    }
+  where
+    (newSecret, newGen) = next generator
+
+data SessionTokens = SessionTokens
+  { tokenMap    :: !TokenMap
+  , lastUpdate  :: !UTCTime
+  , maxInterval :: !NominalDiffTime
+  }
+
+nullSessionTokens :: IO SessionTokens
+nullSessionTokens = SessionTokens
+  <$> (tokens <$> randomIO)
+  <*> getCurrentTime
+  <*> pure defaultUpdateInterval
+
+-- TODO invalidate *twice* if needed
+invalidateTokens :: UTCTime -> SessionTokens -> SessionTokens
+invalidateTokens curTime ts @ SessionTokens {..}
+  | curTime `diffUTCTime` lastUpdate > maxInterval = SessionTokens
+    { tokenMap    = update tokenMap
+    , lastUpdate  = curTime
+    , maxInterval = maxInterval
+    }
+  |                  otherwise                     = ts
+
+{-----------------------------------------------------------------------
+-- Tokens
+-----------------------------------------------------------------------}
+
+tryUpdateSecret :: TVar SessionTokens -> IO ()
+tryUpdateSecret toks = do
+  curTime <- getCurrentTime
+  atomically $ modifyTVar' toks (invalidateTokens curTime)
+
+grantToken :: Hashable addr => TVar SessionTokens -> addr -> IO Token
+grantToken sessionTokens addr = do
+  tryUpdateSecret sessionTokens
+  toks <- readTVarIO sessionTokens
+  return $ Network.BitTorrent.DHT.Token.lookup addr $ tokenMap toks
+
+-- | Throws 'HandlerError' if the token is invalid or already
+-- expired. See 'TokenMap' for details.
+checkToken :: Hashable addr => TVar SessionTokens -> addr -> Token -> IO Bool
+checkToken sessionTokens addr questionableToken = do
+  tryUpdateSecret sessionTokens
+  toks <- readTVarIO sessionTokens
+  return $ member addr questionableToken (tokenMap toks)
+