Factor out some new libraries

	word64-map:
		Data.Word64Map
	network-addr:
		Network.Address
	tox-crypto:
		Crypto.Tox
	lifted-concurrent:
		Control.Concurrent.Lifted.Instrument
		Control.Concurrent.Async.Lifted.Instrument
	psq-wrap:
		Data.Wrapper.PSQInt
		Data.Wrapper.PSQ
	minmax-psq:
		Data.MinMaxPSQ
        tasks:
                Control.Concurrent.Tasks
	kad:
		Network.Kademlia
		Network.Kademlia.Bootstrap
		Network.Kademlia.Routing
		Network.Kademlia.CommonAPI
		Network.Kademlia.Persistence
		Network.Kademlia.Search

3 files changed, 468 insertions, 0 deletions

diff --git a/dht/src/Network/BitTorrent/DHT/ContactInfo.hs b/dht/src/Network/BitTorrent/DHT/ContactInfo.hs
new file mode 100644
index 00000000..ec7e6658
--- /dev/null
+++ b/dht/src/Network/BitTorrent/DHT/ContactInfo.hs
@@ -0,0 +1,254 @@
+{-# LANGUAGE BangPatterns #-}
+module Network.BitTorrent.DHT.ContactInfo
+       ( PeerStore
+       , PeerAddr(..)
+       , Network.BitTorrent.DHT.ContactInfo.lookup
+       , Network.BitTorrent.DHT.ContactInfo.freshPeers
+       , Network.BitTorrent.DHT.ContactInfo.insertPeer
+       , deleteOlderThan
+       , knownSwarms
+       ) where
+
+import Control.Applicative
+import Data.Default
+import Data.List as L
+import Data.Maybe
+import Data.HashMap.Strict as HM
+import Data.Serialize
+import Data.Semigroup
+import Data.Wrapper.PSQ as PSQ
+import Data.Time.Clock.POSIX
+import Data.ByteString (ByteString)
+import Data.Word
+
+import Data.Torrent
+import Network.Address
+
+-- {-
+-- import Data.HashMap.Strict as HM
+--
+-- import Data.Torrent.InfoHash
+-- import Network.Address
+--
+-- -- increase prefix when table is too large
+-- -- decrease prefix when table is too small
+-- -- filter outdated peers
+--
+-- {-----------------------------------------------------------------------
+-- --  PeerSet
+-- -----------------------------------------------------------------------}
+--
+-- type PeerSet a = [(PeerAddr, NodeInfo a, Timestamp)]
+--
+-- -- compare PSQueue vs Ordered list
+--
+-- takeNewest :: PeerSet a -> [PeerAddr]
+-- takeNewest = undefined
+--
+-- dropOld :: Timestamp -> PeerSet a -> PeerSet a
+-- dropOld = undefined
+--
+-- insert :: PeerAddr -> Timestamp -> PeerSet a -> PeerSet a
+-- insert = undefined
+--
+-- type Mask = Int
+-- type Size = Int
+-- type Timestamp = Int
+--
+-- {-----------------------------------------------------------------------
+-- --  InfoHashMap
+-- -----------------------------------------------------------------------}
+--
+-- -- compare handwritten prefix tree versus IntMap
+--
+-- data Tree a
+--   = Nil
+--   | Tip !InfoHash !(PeerSet a)
+--   | Bin !InfoHash !Mask !Size !Timestamp (Tree a) (Tree a)
+--
+-- insertTree :: InfoHash -> a -> Tree a -> Tree a
+-- insertTree = undefined
+--
+-- type Prio = Int
+--
+-- --shrink :: ContactInfo ip -> Int
+-- shrink  Nil      = Nil
+-- shrink (Tip _ _) = undefined
+-- shrink (Bin _ _) = undefined
+--
+-- {-----------------------------------------------------------------------
+-- -- InfoHashMap
+-- -----------------------------------------------------------------------}
+--
+-- -- compare new design versus HashMap
+--
+-- data IntMap k p a
+-- type ContactInfo = Map InfoHash Timestamp (Set (PeerAddr IP) Timestamp)
+--
+-- data ContactInfo ip = PeerStore
+--   { maxSize    :: Int
+--   , prefixSize :: Int
+--   , thisNodeId :: NodeId
+--
+--   , count      :: Int  -- ^ Cached size of the 'peerSet'
+--   , peerSet    :: HashMap InfoHash [PeerAddr ip]
+--   }
+--
+-- size :: ContactInfo ip -> Int
+-- size = undefined
+--
+-- prefixSize :: ContactInfo ip -> Int
+-- prefixSize = undefined
+--
+-- lookup :: InfoHash -> ContactInfo ip -> [PeerAddr ip]
+-- lookup = undefined
+--
+-- insert :: InfoHash -> PeerAddr ip -> ContactInfo ip -> ContactInfo ip
+-- insert = undefined
+--
+-- -- | Limit in size.
+-- prune :: NodeId -> Int -> ContactInfo ip -> ContactInfo ip
+-- prune pref targetSize  Nil      = Nil
+-- prune pref targetSize (Tip _ _) = undefined
+--
+-- -- | Remove expired entries.
+-- splitGT :: Timestamp -> ContactInfo ip -> ContactInfo ip
+-- splitGT = undefined
+-- -}
+
+-- | Storage used to keep track a set of known peers in client,
+-- tracker or DHT sessions.
+newtype PeerStore = PeerStore (HashMap InfoHash SwarmData)
+
+type Timestamp = POSIXTime
+
+data SwarmData = SwarmData
+    { peers :: !(PSQ PeerAddr Timestamp)
+    , name :: !(Maybe ByteString)
+    }
+
+-- | This wrapper will serialize an ip address with a '4' or '6' prefix byte
+-- to indicate whether it is IPv4 or IPv6.
+--
+-- Note: it does not serialize port numbers.
+newtype SerializeAddress a = SerializeAddress { unserializeAddress :: a }
+
+instance Address a => Serialize (SerializeAddress a) where
+    get = SerializeAddress <$> do
+        c <- get
+        case (c::Word8) of
+            0x34 -> do ip4 <- get
+                       return $ fromJust $ fromAddr (ip4::IPv4)
+            0x36 -> do ip6 <- get
+                       return $ fromJust $ fromAddr (ip6::IPv6)
+            _ -> return $ error "cannot deserialize non-IP SerializeAddress"
+    put (SerializeAddress a)
+        | Just ip4 <- fromAddr a
+                    = put (0x34::Word8) >> put (ip4::IPv4)
+        | Just ip6 <- fromAddr a
+                    = put (0x36::Word8) >> put (ip6::IPv6)
+        | otherwise = return $ error "cannot serialize non-IP SerializeAddress"
+
+
+instance Serialize SwarmData where
+    get = flip SwarmData <$> get
+                         <*> ( PSQ.fromList . L.map parseAddr <$> get )
+     where
+        parseAddr (pid,addr,port) = PeerAddr { peerId = pid
+                                             , peerHost = unserializeAddress addr
+                                             , peerPort = port
+                                             }
+                                      :-> 0
+
+    put SwarmData{..} = do
+        put name
+        put $ L.map (\(addr :-> _) -> (peerId addr, SerializeAddress addr, peerPort addr))
+                -- XXX: should we serialize the timestamp?
+            $ PSQ.toList peers
+
+knownSwarms :: PeerStore -> [ (InfoHash, Int, Maybe ByteString) ]
+knownSwarms (PeerStore m) = L.map (\(ih,SwarmData q n) -> (ih,PSQ.size q,n)) $ HM.toList m
+
+swarmSingleton :: PeerAddr -> SwarmData
+swarmSingleton a = SwarmData
+    { peers = PSQ.singleton a 0
+    , name = Nothing }
+
+swarmInsert :: SwarmData -> SwarmData -> SwarmData
+swarmInsert new old = SwarmData
+    { peers = L.foldl' (\q (a :-> t) -> PSQ.insertWith newerTimeStamp a t q) (peers old) (PSQ.toList $ peers new)
+    , name = name new <|> name old -- TODO: decodeUtf8' check
+    }
+ where
+    newerTimeStamp newtime oldtime = if newtime > oldtime then newtime else oldtime
+
+isSwarmOccupied :: SwarmData -> Bool
+isSwarmOccupied SwarmData{..} = not $ PSQ.null peers
+
+-- | Empty store.
+instance Default (PeerStore) where
+  def = PeerStore HM.empty
+  {-# INLINE def #-}
+
+instance Semigroup PeerStore where
+  PeerStore a <> PeerStore b =
+    PeerStore (HM.unionWith swarmInsert a b)
+  {-# INLINE (<>) #-}
+
+-- | Monoid under union operation.
+instance Monoid PeerStore where
+  mempty  = def
+  {-# INLINE mempty #-}
+
+  mappend (PeerStore a) (PeerStore b) =
+    PeerStore (HM.unionWith swarmInsert a b)
+  {-# INLINE mappend #-}
+
+-- | Can be used to store peers between invocations of the client
+-- software.
+instance Serialize PeerStore where
+  get = PeerStore . HM.fromList <$> get
+  put (PeerStore m) = put (L.filter (isSwarmOccupied . snd) $ HM.toList m)
+
+-- | Returns all peers associated with a given info hash.
+lookup :: InfoHash -> PeerStore -> [PeerAddr]
+lookup ih (PeerStore m) = maybe [] (PSQ.keys . peers) $ HM.lookup ih m
+
+batchSize :: Int
+batchSize = 64
+
+-- | Used in 'get_peers' DHT queries.
+freshPeers :: InfoHash -> Timestamp -> PeerStore -> ([PeerAddr], PeerStore)
+freshPeers ih tm (PeerStore m) = fromMaybe ([],PeerStore m) $ do
+    swarm <- HM.lookup ih m
+    let ps0 = take batchSize $ unfoldr (incomp minView) (peers swarm)
+        peers' = case reverse ps0 of
+                    (_,psq):_ -> psq
+                    _         -> peers swarm
+        ps = L.map (key . fst) ps0
+        m' = HM.insert ih swarm { peers = L.foldl' (\q p -> PSQ.insert p tm q) peers' ps } m
+    return $! m' `seq` (ps,PeerStore m')
+
+incomp :: (x -> Maybe (r,x)) -> x -> Maybe ((r,x),x)
+incomp !f !x = do
+   (result,x') <- f x
+   pure $! ( (result,x'), x' )
+
+-- | Used in 'announce_peer' DHT queries.
+insertPeer :: InfoHash -> Maybe ByteString -> PeerAddr -> PeerStore -> PeerStore
+insertPeer !ih !name !a !(PeerStore m) = seq a' $ PeerStore (HM.insertWith swarmInsert ih a' m)
+ where
+    a' = SwarmData { peers = PSQ.singleton a 0
+                   , name = name }
+
+deleteOlderThan :: POSIXTime -> PeerStore -> PeerStore
+deleteOlderThan cutoff (PeerStore m) = PeerStore $ HM.mapMaybe gc m
+    where
+        gc :: SwarmData -> Maybe SwarmData
+        gc swarms = fmap (\ps -> swarms { peers = ps }) $ gcPSQ (peers swarms)
+
+        gcPSQ :: PSQKey a => PSQ a Timestamp -> Maybe (PSQ a Timestamp)
+        gcPSQ ps = case minView ps of
+            Nothing                            -> Nothing
+            Just (_ :-> tm, ps') | tm < cutoff -> gcPSQ ps'
+            Just _                             -> Just ps
diff --git a/dht/src/Network/BitTorrent/DHT/Readme.md b/dht/src/Network/BitTorrent/DHT/Readme.md
new file mode 100644
index 00000000..e2352f10
--- /dev/null
+++ b/dht/src/Network/BitTorrent/DHT/Readme.md
@@ -0,0 +1,13 @@
+References
+==========
+
+Some good references excluding BEPs:
+
+* [Kademlia wiki page][kademlia-wiki]
+* [Kademlia: A Peer-to-peer Information System Based on the XOR Metric][kademlia-paper]
+* [BitTorrent Mainline DHT Measurement][mldht]
+* Profiling a Million User DHT. (paper)
+
+[kademlia-wiki]: http://en.wikipedia.org/wiki/Kademlia
+[kademlia-paper]: http://pdos.csail.mit.edu/~petar/papers/maymounkov-kademlia-lncs.pdf
+[mldht]: http://www.cs.helsinki.fi/u/jakangas/MLDHT/
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)
+