Implemented more of the Mainline DHT rewrite.

2 files changed, 472 insertions, 42 deletions

diff --git a/Mainline.hs b/Mainline.hs
index d24b3376..12d1540b 100644
--- a/Mainline.hs
+++ b/Mainline.hs
@@ -4,32 +4,51 @@
 {-# LANGUAGE DeriveTraversable          #-}
 {-# LANGUAGE FlexibleInstances          #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE PatternSynonyms            #-}
+{-# LANGUAGE StandaloneDeriving         #-}
 module Mainline where
 
+import Control.Applicative
 import Control.Arrow
 import Control.Concurrent.STM
 import Crypto.Random
-import Data.BEncode             as BE
-import Data.BEncode.BDict       as BE
+import Data.BEncode                       as BE
+import qualified Data.BEncode.BDict       as BE
+         ;import Data.BEncode.BDict       (BKey)
+import Data.Bits
+import Data.Bits.ByteString
 import Data.Bool
-import Data.ByteArray
-import Data.ByteString          (ByteString)
-import Data.ByteString          as B
-import Data.ByteString.Lazy     (toStrict)
+import qualified Data.ByteArray           as BA
+         ;import Data.ByteArray           (ByteArrayAccess)
+import qualified Data.ByteString          as B
+         ;import Data.ByteString          (ByteString)
+import Data.ByteString.Lazy               (toStrict)
 import Data.Data
+import Data.Default
+import Data.Hashable
 import Data.IP
+import Data.List
 import Data.Maybe
 import Data.Monoid
-import qualified Data.Serialize as S
+import qualified Data.Serialize           as S
+import Data.Set                           (Set)
+import Data.Torrent
 import Data.Typeable
 import Data.Word
-import Network.Address          (Address, fromSockAddr, sockAddrPort,
-                                 toSockAddr, withPort)
+import Network.Address                    (Address, fromSockAddr, setPort,
+                                           sockAddrPort, toSockAddr)
+import Network.BitTorrent.DHT.ContactInfo as Peers
+import Network.BitTorrent.DHT.Token       as Token
+import qualified Network.DHT.Routing      as R
+         ;import Network.DHT.Routing      (Info, Timestamp, getTimestamp)
 import Network.QueryResponse
 import Network.Socket
 
 newtype NodeId = NodeId ByteString
- deriving (Eq,Ord,Show,ByteArrayAccess, BEncode)
+ deriving (Eq,Ord,Show,ByteArrayAccess, BEncode, Bits)
+
+instance FiniteBits NodeId where
+    finiteBitSize _ = 160
 
 data NodeInfo = NodeInfo
   { nodeId   :: NodeId
@@ -37,8 +56,44 @@ data NodeInfo = NodeInfo
   , nodePort :: PortNumber
   }
 
+instance Hashable NodeInfo where
+  hashWithSalt s ni = hashWithSalt s (nodeIP ni , nodePort ni)
+  {-# INLINE hashWithSalt #-}
+
+
+
+{-
+
+-- | KRPC 'compact list' compatible encoding: contact information for
+-- nodes is encoded as a 26-byte string. Also known as "Compact node
+-- info" the 20-byte Node ID in network byte order has the compact
+-- IP-address/port info concatenated to the end.
+  get = NodeInfo <$> (NodeId <$> S.getBytes 20 ) <*> S.get <*> S.get
+-}
+
+getNodeInfo4 :: S.Get NodeInfo
+getNodeInfo4 = NodeInfo <$> (NodeId <$> S.getBytes 20)
+                        <*> (IPv4 <$> S.get)
+                        <*> S.get
+
+putNodeInfo4 :: NodeInfo -> S.Put
+putNodeInfo4 (NodeInfo (NodeId nid) (IPv4 ip) port)
+               = S.putByteString nid >> S.put ip >> S.put port
+putNodeInfo4 _ = return ()
+
+getNodeInfo6 :: S.Get NodeInfo
+getNodeInfo6 = NodeInfo <$> (NodeId <$> S.getBytes 20)
+                        <*> (IPv6 <$> S.get)
+                        <*> S.get
+
+putNodeInfo6 :: NodeInfo -> S.Put
+putNodeInfo6 (NodeInfo (NodeId nid) (IPv6 ip) port)
+               = S.putByteString nid >> S.put ip >> S.put port
+putNodeInfo6 _ = return ()
+
+
 nodeAddr :: NodeInfo -> SockAddr
-nodeAddr (NodeInfo _ ip port) = toSockAddr ip `withPort` port
+nodeAddr (NodeInfo _ ip port) = setPort port $ toSockAddr ip
 
 nodeInfo :: NodeId -> SockAddr -> Either String NodeInfo
 nodeInfo nid saddr
@@ -110,7 +165,7 @@ instance BE.BEncode (Message BValue) where
 
 encodeMessage (Q origin tid a meth ro)
     = case a of
-        BDict args -> encodeQuery tid meth (BDict $ genericArgs origin ro `union` args)
+        BDict args -> encodeQuery tid meth (BDict $ genericArgs origin ro `BE.union` args)
         _          -> encodeQuery tid meth a -- XXX: Not really a valid query.
 encodeMessage (R origin tid v ip)
     = case v of
@@ -148,6 +203,60 @@ encodePacket :: Message BValue -> NodeInfo -> (ByteString, SockAddr)
 encodePacket msg ni = ( toStrict $ BE.encode msg
                       , nodeAddr ni )
 
+classify :: Message BValue -> MessageClass String Method TransactionId
+classify (Q { msgID = tid, qryMethod = meth }) = IsQuery meth tid
+classify (R { msgID = tid                   }) = IsResponse tid
+
+encodePayload :: BEncode a => TransactionId -> NodeInfo -> NodeInfo -> a -> Message BValue
+encodePayload tid self dest b = R (nodeId self) tid (Right $ BE.toBEncode b) (Just $ nodeAddr dest)
+
+errorPayload :: TransactionId -> NodeInfo -> NodeInfo -> Error -> Message a
+errorPayload tid self dest e = R (nodeId self) tid (Left e) (Just $ nodeAddr dest)
+
+decodePayload :: BEncode a => Message BValue -> Either String a
+decodePayload msg = BE.fromBEncode $ qryPayload msg
+
+type Handler = MethodHandler String TransactionId NodeInfo (Message BValue)
+
+handler :: ( BEncode a
+           , BEncode b
+           ) =>
+           (NodeInfo -> a -> IO b) -> Maybe Handler
+handler f = Just $ MethodHandler decodePayload encodePayload f
+
+
+handlerE :: ( BEncode a
+            , BEncode b
+            ) =>
+            (NodeInfo -> a -> IO (Either Error b)) -> Maybe Handler
+handlerE f = Just $ MethodHandler decodePayload enc f
+ where
+    enc tid self dest (Left e)  = errorPayload  tid self dest e
+    enc tid self dest (Right b) = encodePayload tid self dest b
+
+type AnnounceSet = Set (InfoHash, PortNumber)
+
+data SwarmsDatabase = SwarmsDatabase
+    { contactInfo   :: !( TVar PeerStore     ) -- ^ Published by other nodes.
+    , sessionTokens :: !( TVar SessionTokens ) -- ^ Query session IDs.
+    , announceInfo  :: !( TVar AnnounceSet   ) -- ^ To publish by this node.
+    }
+
+newSwarmsDatabase :: IO SwarmsDatabase
+newSwarmsDatabase = do
+    toks <- nullSessionTokens
+    atomically
+        $ SwarmsDatabase <$> newTVar def
+                    <*> newTVar toks
+                    <*> newTVar def
+
+type RoutingInfo = Info NodeInfo NodeId
+
+data Routing = Routing
+    { routing4 :: !( TVar (Maybe RoutingInfo) )
+    , routing6 :: !( TVar (Maybe RoutingInfo) )
+    }
+
 newClient ::
       SockAddr -> IO (Client String Method TransactionId NodeInfo (Message BValue))
 newClient addr = do
@@ -156,13 +265,26 @@ newClient addr = do
     self <- atomically $ newTVar
              $ NodeInfo nid (fromMaybe (toEnum 0) $ fromSockAddr addr)
                             (fromMaybe 0 $ sockAddrPort addr)
-    -- drg <- getSystemDRG
-    let net = layerTransport parsePacket encodePacket udp
+    routing <- atomically $ Routing <$> newTVar Nothing <*> newTVar Nothing
+    swarms <- newSwarmsDatabase
+    let net = onInbound grok $ layerTransport parsePacket encodePacket udp
+        grok _ _ = do
+            -- TODO Update kademlia table.
+            -- TODO Update external ip address and update BEP-42 node id.
+            return ()
         dispatch tbl = DispatchMethods
             { classifyInbound = classify
             , lookupHandler = handlers
             , tableMethods = tbl
             }
+
+        handlers :: Method -> Maybe Handler
+        handlers ( Method "ping"          ) = handler pingH
+        handlers ( Method "find_node"     ) = handler $ findNodeH routing
+        handlers ( Method "get_peers"     ) = handler $ getPeersH routing swarms
+        handlers ( Method "announce_peer" ) = handlerE $ announceH swarms
+        handlers ( Method meth            ) = Just $ defaultHandler meth
+
         mapT = transactionMethods mapMethods gen
         gen :: Word16 -> (TransactionId, Word16)
         gen cnt = (TransactionId $ S.encode cnt, cnt+1)
@@ -176,29 +298,287 @@ newClient addr = do
             , clientResponseId    = return
             }
 
-classify :: Message BValue -> MessageClass String Method TransactionId
-classify (Q { msgID = tid, qryMethod = meth }) = IsQuery meth tid
-classify (R { msgID = tid                   }) = IsResponse tid
-
-encodePayload tid self dest b = R (nodeId self) tid (Right $ BE.toBEncode b) (Just $ nodeAddr dest)
-
-errorPayload tid self dest e = R (nodeId self) tid (Left e) (Just $ nodeAddr dest)
-
-decodePayload :: BEncode a => Message BValue -> Either String a
-decodePayload msg = BE.fromBEncode $ qryPayload msg
-
-handler f = Just $ MethodHandler decodePayload encodePayload f
-
-handlers :: Method -> Maybe (MethodHandler String TransactionId NodeInfo (Message BValue))
-handlers (Method "ping"     ) = error "handler pingH"
-handlers (Method "find_node") = error "find_node"
-handlers (Method "get_peers") = error "get_peers"
-handlers (Method meth       ) = Just $ MethodHandler decodePayload errorPayload (defaultH meth)
-
-data Ping = Ping
+defaultHandler :: ByteString -> Handler
+defaultHandler meth = MethodHandler decodePayload errorPayload returnError
+ where
+    returnError :: NodeInfo -> BValue -> IO Error
+    returnError _ _ = return $ Error MethodUnknown ("Unknown method " <> meth)
+
+
+data Ping = Ping deriving Show
+
+-- Pong is the same as Ping.
+type Pong = Ping
+pattern Pong = Ping
+
+instance BEncode Ping where
+  toBEncode Ping = toDict endDict
+  fromBEncode _  = pure Ping
+
+data WantIP = Want_IP4 | Want_IP6 | Want_Both
+ deriving (Eq, Enum, Ord, Show)
+
+wantList :: WantIP -> [ByteString]
+wantList Want_IP4  = ["ip4"]
+wantList Want_IP6  = ["ip6"]
+wantList Want_Both = ["ip4","ip6"]
+
+instance BEncode WantIP where
+    toBEncode w = toBEncode $ wantList w
+    fromBEncode bval = do
+        wants <- fromBEncode bval
+        let _ = wants :: [ByteString]
+        case (elem "ip4" wants, elem "ip6" wants) of
+             (True,True)  -> Right Want_Both
+             (True,False) -> Right Want_IP4
+             (False,True) -> Right Want_IP6
+             _            -> Left "Unrecognized IP type."
+
+data FindNode = FindNode NodeId (Maybe WantIP)
+
+instance BEncode FindNode where
+  toBEncode (FindNode nid iptyp) = toDict   $ target_key .=! nid
+                                              .: want_key .=? iptyp
+                                              .: endDict
+  fromBEncode                    = fromDict $ FindNode  <$>! target_key
+                                                        <*>? want_key
+
+data NodeFound = NodeFound
+    { nodes4 :: [NodeInfo]
+    , nodes6 :: [NodeInfo]
+    }
+
+instance BEncode NodeFound where
+  toBEncode (NodeFound ns ns6) = toDict $
+       nodes_key .=?
+        (if Prelude.null ns then Nothing
+                            else Just (S.runPut (mapM_ putNodeInfo4 ns)))
+    .: nodes6_key .=?
+        (if Prelude.null ns6 then Nothing
+                             else Just (S.runPut (mapM_ putNodeInfo6 ns6)))
+    .: endDict
 
-pingH :: NodeInfo -> Ping -> IO Ping
-pingH = error "pingH"
+  fromBEncode bval = NodeFound <$> ns4 <*> ns6
+   where
+    ns4 = fromDict (binary getNodeInfo4 nodes_key) bval
+    ns6 = fromDict (binary getNodeInfo6 nodes6_key) bval
+
+binary :: S.Get a -> BKey -> BE.Get [a]
+binary get k = field (req k) >>= either (fail . format) return .
+    S.runGet (many get)
+  where
+    format str = "fail to deserialize " ++ show k ++ " field: " ++ str
+
+pingH :: NodeInfo -> Ping -> IO Pong
+pingH _ Ping = return Pong
+
+-- | True if the argument is an IPv4-mapped address with prefix ::FFFF:0:0/96
+-- as defined in RFC 4291.
+is4mapped :: IPv6 -> Bool
+is4mapped ip
+    | [0,0,0,0,0,0,0xffff,_] <- fromIPv6 ip
+                = True
+    | otherwise = False
+
+prefer4or6 :: NodeInfo -> Maybe WantIP -> WantIP
+prefer4or6 addr iptyp = fromMaybe unspecified iptyp
+ where
+    unspecified = case nodeIP addr of
+        IPv4 _               -> Want_IP4
+        IPv6 a | is4mapped a -> Want_IP4
+               | otherwise   -> Want_IP6
+
+findNodeH :: Routing -> NodeInfo -> FindNode -> IO NodeFound
+findNodeH routing addr (FindNode node iptyp) = do
+    let preferred = prefer4or6 addr iptyp
+    ks  <- bool (return []) (go $ routing4 routing) (preferred /= Want_IP6)
+    ks6 <- bool (return []) (go $ routing6 routing) (preferred /= Want_IP4)
+    return $ NodeFound ks ks6
+ where
+    go var = do
+        let myid   = error "TODO myid"   :: NodeId
+            k      = error "TODO k"      :: Int
+            nobkts = error "TODO nobkts" :: Int
+        nfo <- atomically $ readTVar var
+        let tbl = maybe (R.nullTable myid nobkts) R.myBuckets nfo
+        return $ R.kclosest nodeId k node tbl
+
+data GetPeers = GetPeers InfoHash (Maybe WantIP)
+
+instance BEncode GetPeers where
+  toBEncode (GetPeers ih iptyp)
+              = toDict   $ info_hash_key .=! ih
+                        .: want_key .=? iptyp
+                        .: endDict
+  fromBEncode = fromDict $ GetPeers <$>! info_hash_key <*>? want_key
+
+
+type PeerList = Either [NodeInfo] [PeerAddr]
+
+data GotPeers = GotPeers
+  { -- | If the queried node has no peers for the infohash, returned
+    -- the K nodes in the queried nodes routing table closest to the
+    -- infohash supplied in the query.
+    peers        :: PeerList
+
+    -- | The token value is a required argument for a future
+    -- announce_peer query.
+  , grantedToken :: Token
+  } -- deriving (Show, Eq, Typeable)
+
+nodeIsIPv6 :: NodeInfo -> Bool
+nodeIsIPv6 (NodeInfo _ (IPv6 _) _) = True
+nodeIsIPv6 _                       = False
+
+instance BEncode GotPeers where
+  toBEncode GotPeers {..} = toDict $
+    case peers of
+      Left  ns
+        | let (ns6,ns4) = partition nodeIsIPv6 ns
+        ->
+           nodes_key .=? (if null ns4 then Nothing
+                                      else Just $ S.runPut (mapM_ putNodeInfo4 ns4))
+        .: nodes6_key .=? (if null ns6 then Nothing
+                                       else Just $ S.runPut (mapM_ putNodeInfo4 ns6))
+        .: token_key .=! grantedToken
+        .: endDict
+      Right ps ->
+           token_key .=! grantedToken
+        .: peers_key .=! map S.encode ps -- TODO: Spec says we shouldn't mix ip4/ip6 here.
+                                         -- (We could filter in MethodHandler.)
+        .: endDict
 
-defaultH :: ByteString -> NodeInfo -> BValue -> IO Error
-defaultH meth _ _ = return $ Error MethodUnknown ("Unknown method " <> meth)
+  fromBEncode = fromDict $ do
+     mns4 <- optional (binary getNodeInfo4 nodes_key)        -- "nodes"
+     mns6 <- optional (binary getNodeInfo6 nodes6_key)       -- "nodes6"
+     let mns = ((++) <$> mns4 <*> mns6)
+               <|> mns4
+               <|> mns6
+     tok <- field    (req    token_key)                      -- "token"
+     mps <- optional (field (req peers_key) >>= decodePeers) -- "values"
+     case (Right <$> mps) <|> (Left <$> mns) of
+       Nothing -> fail "get_peers: neihter peers nor nodes key is valid"
+       Just xs -> pure $ GotPeers xs tok
+     where
+       decodePeers = either fail pure . mapM S.decode
+
+getPeersH :: Routing -> SwarmsDatabase -> NodeInfo -> GetPeers -> IO GotPeers
+getPeersH routing (SwarmsDatabase peers toks _) naddr (GetPeers ih iptyp) = do
+  ps <- do
+      tm <- getTimestamp
+      ps <- atomically $ do
+        (ps,store') <- Peers.freshPeers ih tm <$> readTVar peers
+        writeTVar peers store'
+        return ps
+      if null ps
+        then Left <$> error "TODO: getClosest ih"
+        else return (Right ps)
+  tok <- grantToken toks naddr
+  return $ GotPeers ps tok
+
+-- | Announce that the peer, controlling the querying node, is
+-- downloading a torrent on a port.
+data Announce = Announce
+  { -- | If set, the 'port' field should be ignored and the source
+    -- port of the UDP packet should be used as the peer's port
+    -- instead. This is useful for peers behind a NAT that may not
+    -- know their external port, and supporting uTP, they accept
+    -- incoming connections on the same port as the DHT port.
+    impliedPort :: Bool
+
+    -- | infohash of the torrent;
+  , topic    :: InfoHash
+
+    -- | some clients announce the friendly name of the torrent here.
+  , announcedName :: Maybe ByteString
+
+    -- | the port /this/ peer is listening;
+  , port     :: PortNumber
+
+    -- TODO: optional boolean "seed" key
+
+    -- | received in response to a previous get_peers query.
+  , sessionToken :: Token
+
+  } deriving (Show, Eq, Typeable)
+
+peer_ip_key   = "ip"
+peer_id_key   = "peer id"
+peer_port_key = "port"
+msg_type_key   = "msg_type"
+piece_key      = "piece"
+total_size_key = "total_size"
+node_id_key :: BKey
+node_id_key = "id"
+read_only_key :: BKey
+read_only_key = "ro"
+want_key :: BKey
+want_key = "want"
+target_key :: BKey
+target_key = "target"
+nodes_key :: BKey
+nodes_key = "nodes"
+nodes6_key :: BKey
+nodes6_key = "nodes6"
+info_hash_key :: BKey
+info_hash_key = "info_hash"
+peers_key :: BKey
+peers_key = "values"
+token_key :: BKey
+token_key = "token"
+name_key :: BKey
+name_key = "name"
+port_key :: BKey
+port_key = "port"
+implied_port_key :: BKey
+implied_port_key = "implied_port"
+
+instance BEncode Announce where
+  toBEncode Announce {..} = toDict $
+       implied_port_key .=? flagField impliedPort
+    .: info_hash_key    .=! topic
+    .: name_key         .=? announcedName
+    .: port_key         .=! port
+    .: token_key        .=! sessionToken
+    .: endDict
+    where
+      flagField flag = if flag then Just (1 :: Int) else Nothing
+
+  fromBEncode = fromDict $ do
+    Announce <$> (boolField <$> optional (field (req implied_port_key)))
+             <*>! info_hash_key
+             <*>? name_key
+             <*>! port_key
+             <*>! token_key
+    where
+      boolField = maybe False (/= (0 :: Int))
+
+
+
+-- | The queried node must verify that the token was previously sent
+-- to the same IP address as the querying node. Then the queried node
+-- should store the IP address of the querying node and the supplied
+-- port number under the infohash in its store of peer contact
+-- information.
+data Announced = Announced
+  deriving (Show, Eq, Typeable)
+
+instance BEncode Announced where
+  toBEncode   _ = toBEncode Ping
+  fromBEncode _ = pure  Announced
+
+announceH :: SwarmsDatabase -> NodeInfo -> Announce -> IO (Either Error Announced)
+announceH (SwarmsDatabase peers toks _) naddr announcement = do
+  checkToken toks naddr (sessionToken announcement)
+    >>= bool (Left  <$> return (Error ProtocolError "invalid parameter: token"))
+             (Right <$> go)
+ where
+    go = do
+      let annPort  = if impliedPort announcement
+                        then nodePort naddr
+                        else port announcement
+          peerAddr = PeerAddr Nothing (nodeIP naddr) annPort
+      atomically
+            $ modifyTVar' peers
+            $ insertPeer (topic announcement) (announcedName announcement) peerAddr
+      return Announced
diff --git a/src/Network/BitTorrent/DHT/Token.hs b/src/Network/BitTorrent/DHT/Token.hs
index e3a6b1f6..08079b75 100644
--- a/src/Network/BitTorrent/DHT/Token.hs
+++ b/src/Network/BitTorrent/DHT/Token.hs
@@ -24,6 +24,10 @@ module Network.BitTorrent.DHT.Token
 
          -- * Session tokens
        , TokenMap
+       , SessionTokens
+       , nullSessionTokens
+       , checkToken
+       , grantToken
 
          -- ** Construction
        , Network.BitTorrent.DHT.Token.tokens
@@ -52,7 +56,7 @@ import Data.Hashable
 import Data.String
 import Data.Time
 import System.Random
-
+import Control.Concurrent.STM
 import Network.Address
 
 -- TODO use ShortByteString
@@ -77,7 +81,7 @@ 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 => NodeAddr a -> Secret -> Token
+makeToken :: Hashable a => a -> Secret -> Token
 makeToken n s = Token $ toBS $ hashWithSalt s n
   where
     toBS = toStrict . toLazyByteString . int64BE . fromIntegral
@@ -104,7 +108,7 @@ tokens seed = (`evalState` mkStdGen seed) $
 -- 'update's.
 --
 -- Typically used to handle find_peers query.
-lookup :: Hashable a => NodeAddr a -> TokenMap -> Token
+lookup :: Hashable a => a -> TokenMap -> Token
 lookup addr TokenMap {..} = makeToken addr curSecret
 
 -- | Check if token is valid.
@@ -112,7 +116,7 @@ lookup addr TokenMap {..} = makeToken addr curSecret
 -- 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 => NodeAddr a -> Token -> TokenMap -> Bool
+member :: Hashable a => a -> Token -> TokenMap -> Bool
 member addr token TokenMap {..} = token `L.elem` valid
   where valid = makeToken addr <$> [curSecret, prevSecret]
 
@@ -130,3 +134,49 @@ update TokenMap {..} = TokenMap
     }
   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)
+