-- |
--   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)