It builds!

1 files changed, 75 insertions, 1 deletions

diff --git a/src/Network/DHT/Mainline.hs b/src/Network/DHT/Mainline.hs
index 29d4231d..aefd7742 100644
--- a/src/Network/DHT/Mainline.hs
+++ b/src/Network/DHT/Mainline.hs
@@ -88,6 +88,10 @@ module Network.DHT.Mainline
        , Announce  (..)
        , Announced (..)
 #endif
+       , DHTData(..)
+       , SessionTokens(..)
+       , grantToken
+       , checkToken
        ) where
 
 import Control.Applicative
@@ -118,13 +122,19 @@ import Network.DatagramServer.Mainline
 import Data.Maybe
 
 import Data.Torrent (InfoHash)
-import Network.BitTorrent.DHT.Token
+import Network.BitTorrent.DHT.Token as T
+import Network.BitTorrent.DHT.ContactInfo
 #ifdef VERSION_bencoding
 import Network.DatagramServer ()
 #endif
 import Network.DatagramServer.Types hiding (Query,Response)
 import Network.DHT.Types
 import Network.DHT.Routing
+import Data.Time
+import Control.Concurrent.STM
+import System.Random
+import Data.Hashable
+
 
 {-----------------------------------------------------------------------
 -- envelopes
@@ -472,6 +482,59 @@ bep42 addr (NodeId r)
         where msk | BS.length ip == 4  =  ip4mask
                   | otherwise          =  ip6mask
 
+{-----------------------------------------------------------------------
+-- Tokens policy
+-----------------------------------------------------------------------}
+
+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 a => TVar SessionTokens -> NodeAddr a -> IO Token
+grantToken sessionTokens addr = do
+  tryUpdateSecret sessionTokens
+  toks <- readTVarIO sessionTokens
+  return $ T.lookup addr $ tokenMap toks
+
+-- | Throws 'HandlerError' if the token is invalid or already
+-- expired. See 'TokenMap' for details.
+checkToken :: Hashable a => TVar SessionTokens -> NodeAddr a -> Token -> IO Bool
+checkToken sessionTokens addr questionableToken = do
+  tryUpdateSecret sessionTokens
+  toks <- readTVarIO sessionTokens
+  return $ T.member addr questionableToken (tokenMap toks)
+
+
+--------------------------
+
+
 instance Kademlia KMessageOf where
     data Ping KMessageOf = Ping
       deriving (Show, Eq, Typeable)
@@ -479,10 +542,17 @@ instance Kademlia KMessageOf where
       deriving (Show, Eq, Typeable)
     newtype NodeFound KMessageOf ip = NodeFound [NodeInfo KMessageOf ip ()]
       deriving (Show, Eq, Typeable)
+    data DHTData KMessageOf ip = TorrentData
+            { contactInfo   :: !(TVar    (PeerStore ip     )) -- ^ published by other nodes;
+            , sessionTokens :: !(TVar     SessionTokens     ) -- ^ query session IDs.
+            }
+
     pingMessage _ = Ping
     pongMessage _ = Ping
     findNodeMessage _ k = FindNode (toNodeId k)
+    findWho (FindNode nid) = nid
     foundNodes (NodeFound ns) = ns
+    foundNodesMessage ns = NodeFound ns
 
     dhtAdjustID _ fallback ip0 arrival
             = fromMaybe fallback $ do
@@ -494,3 +564,7 @@ instance Kademlia KMessageOf where
 
     namePing      _ = "ping"
     nameFindNodes _ = "find-nodes"
+
+    initializeDHTData = TorrentData
+                <$> newTVarIO def
+                <*> (newTVarIO =<< nullSessionTokens)