6 files changed, 62 insertions, 12 deletions

diff --git a/dht-client.cabal b/dht-client.cabal
index 347924ee..de6eb33f 100644
--- a/dht-client.cabal
+++ b/dht-client.cabal
@@ -91,6 +91,7 @@ library
                     Control.TriadCommittee
                     Crypto.Tox
                     Text.XXD
+                    Roster
 
   build-depends:      base
                     , containers
@@ -196,6 +197,10 @@ executable dhtd
                      , stm
                      , cereal
                      , bencoding
+                     , unordered-containers
+                     , vector
+                     , text
+
   if flag(thread-debug)
     build-depends: time
     cpp-options:   -DTHREAD_DEBUG
diff --git a/src/Network/QueryResponse.hs b/src/Network/QueryResponse.hs
index 9563fa7c..fca6d5cc 100644
--- a/src/Network/QueryResponse.hs
+++ b/src/Network/QueryResponse.hs
@@ -27,7 +27,9 @@ import qualified Data.IntMap.Strict as IntMap
          ;import Data.IntMap.Strict (IntMap)
 import qualified Data.Map.Strict    as Map
          ;import Data.Map.Strict    (Map)
-import qualified Data.Word64Map as W64Map
+import qualified Data.Word64Map     as W64Map
+         ;import Data.Word64Map     (Word64Map)
+import Data.Word
 import Data.Maybe
 import Data.Typeable
 import Network.Socket
@@ -333,7 +335,7 @@ intMapMethods :: TableMethods IntMap Int
 intMapMethods = TableMethods IntMap.insert IntMap.delete IntMap.lookup
 
 -- | Methods for using 'Data.Word64Map'.
-w64MapMethods :: TableMethods IntMap Int
+w64MapMethods :: TableMethods Word64Map Word64
 w64MapMethods = TableMethods W64Map.insert W64Map.delete W64Map.lookup
 
 -- | Methods for using 'Data.Map'
diff --git a/src/Network/Tox.hs b/src/Network/Tox.hs
index c587578d..e9220fcb 100644
--- a/src/Network/Tox.hs
+++ b/src/Network/Tox.hs
@@ -87,6 +87,7 @@ import GHC.TypeLits
 
 import Crypto.Tox
 import Data.Word64Map                        (fitsInInt)
+import qualified Data.Word64Map              (empty)
 import Network.Tox.Crypto.Transport          (NetCrypto)
 import qualified Network.Tox.DHT.Handlers    as DHT
 import qualified Network.Tox.DHT.Transport   as DHT
@@ -189,8 +190,8 @@ newClient drg net classify selfAddr handlers modifytbl modifynet = do
             intmap_var <- atomically $ newTVar (drg, mempty)
             return $ Right (intmapT,intmap_var)
          else do
-            let word64mapT = transactionMethods (contramap w64key w64MapMethods) gen
-            map_var <- atomically $ newTVar (drg, mempty)
+            let word64mapT = transactionMethods (contramap w64Key w64MapMethods) gen
+            map_var <- atomically $ newTVar (drg, Data.Word64Map.empty)
             return $ Left (word64mapT,map_var)
     let dispatch tbl var handlers = DispatchMethods
                 { classifyInbound = classify
@@ -236,12 +237,12 @@ getContactInfo Tox{toxCryptoKeys,toxRouting} = join $ atomically $ do
         ns = filter (DHT.isGlobal . nodeIP) [n4,n6]
               ++ concat (zipWith (\a b -> [a,b]) n4s n6s)
     return $ do
-    timestamp <- round . (* 1000000) <$> getPOSIXTime
-    return DHT.DHTPublicKey
-        { dhtpkNonce = timestamp
-        , dhtpk      = id2key self
-        , dhtpkNodes = DHT.SendNodes $ take 4 ns
-        }
+        timestamp <- round . (* 1000000) <$> getPOSIXTime
+        return DHT.DHTPublicKey
+            { dhtpkNonce = timestamp
+            , dhtpk      = id2key self
+            , dhtpkNodes = DHT.SendNodes $ take 4 ns
+            }
 
 isLocalHost :: SockAddr -> Bool
 isLocalHost (SockAddrInet _ host32) = (fromBE32 host32 == 0x7f000001)
diff --git a/src/Network/Tox/DHT/Handlers.hs b/src/Network/Tox/DHT/Handlers.hs
index 9cdf0d06..840e2e6b 100644
--- a/src/Network/Tox/DHT/Handlers.hs
+++ b/src/Network/Tox/DHT/Handlers.hs
@@ -1,6 +1,7 @@
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE PatternSynonyms            #-}
 {-# LANGUAGE TupleSections              #-}
+{-# LANGUAGE CPP                        #-}
 module Network.Tox.DHT.Handlers where
 
 import Network.Tox.DHT.Transport as DHTTransport
@@ -31,6 +32,7 @@ import Data.Maybe
 import Data.Bits
 import Data.Serialize (Serialize)
 import Data.Word
+import Data.List
 import System.IO
 
 data TransactionId = TransactionId
@@ -175,6 +177,7 @@ getNodesH routing addr (GetNodes nid) = do
     ks6 <- go append6 $ routing6 routing
     let (ns1,ns2) = case preferred of Want_IP6 -> (ks6,ks)
                                       Want_IP4 -> (ks,ks6)
+                                      Want_Both -> error $ "BUG:unreachable at " ++ __FILE__ ++ ":" ++ show __LINE__
     return $ SendNodes
            $ if null ns2 then ns1
                          else take 4 (take 3 ns1 ++ ns2)
@@ -240,6 +243,42 @@ ping client addr = do
     hPutStrLn stderr $ show addr ++ " -pong-> " ++ show reply
     maybe (return False) (\Pong -> return True) $ join reply
 
+cookieRequest :: TVar [(SockAddr,(Int,NodeInfo))] -> PublicKey -> Client -> NodeInfo -> IO (Maybe Cookie)
+cookieRequest tvar myDhtKey client addr = do
+    let sockAddr = nodeAddr addr
+    let incAddr sockMap
+            = case partition ((==sockAddr) . fst) sockMap of
+                    ([],xs)                              -> insert (sockAddr, (1  ,addr)) xs
+                    ([(_,(c,addr'))],xs) | addr' == addr -> insert (sockAddr, (c+1,addr)) xs
+                    anythingElse -> error $ "unreachable at " ++ __FILE__ ++ show __LINE__ ++ show anythingElse ++ show (sockAddr,addr)
+        decAddr sockMap
+            = case partition ((==sockAddr) . fst) sockMap of
+                    ([],xs)                              -> xs -- unreachable?
+                    ([(_,(1,addr'))],xs) | addr' == addr -> xs
+                    ([(_,(c,addr'))],xs) | addr' == addr -> insert (sockAddr,(c-1,addr)) xs
+                    anythingElse -> error $ "unreachable at " ++ __FILE__ ++ show __LINE__ ++ show anythingElse ++ show (sockAddr,addr)
+    sockMap <- atomically $ do
+            mp <- incAddr <$> readTVar tvar
+            writeTVar tvar mp
+            return mp
+    let cookieSerializer
+            = MethodSerializer
+                   { methodTimeout  = \tid addr -> do
+                        modifyTVar tvar decAddr
+                        return (addr, 5000000)
+                   , method    = CookieRequestType
+                   , wrapQuery = \tid src dst cr -> DHTCookieRequest $ wrapAsymm tid src dst (, cr)
+                   , unwrapResponse = fmap snd . unCookie
+                   }
+        cookieRequest = CookieRequest myDhtKey
+    hPutStrLn stderr $ show addr ++ " <-- cookieRequest"
+    reply <- QR.sendQuery client cookieSerializer cookieRequest addr
+    hPutStrLn stderr $ show addr ++ " -cookieResponse-> " ++ show reply
+    return $ join reply
+
+unCookie (DHTCookie n24 fcookie) = Just fcookie
+unCookie _ = Nothing
+
 unsendNodes :: Message -> Maybe (Asymm (Nonce8,SendNodes))
 unsendNodes (DHTSendNodes asymm) = Just asymm
 unsendNodes _                    = Nothing
@@ -263,6 +302,7 @@ updateRouting client routing orouter naddr msg = do
     case prefer4or6 naddr Nothing of
         Want_IP4 -> updateTable client naddr orouter (routing4 routing) (committee4 routing) (sched4 routing)
         Want_IP6 -> updateTable client naddr orouter (routing6 routing) (committee6 routing) (sched6 routing)
+        Want_Both -> error $ "BUG:unreachable at " ++ __FILE__ ++ ":" ++ show __LINE__
 
 updateTable :: Client -> NodeInfo -> OnionRouter -> TVar (R.BucketList NodeInfo) -> TriadCommittee NodeId SockAddr -> TVar (Int.PSQ POSIXTime) -> IO ()
 updateTable client naddr orouter tbl committee sched = do
@@ -327,6 +367,7 @@ handlers :: TransportCrypto -> Routing -> PacketKind -> Maybe Handler
 handlers _      routing PingType          = Just $ MethodHandler (isPing snd) mkPong pingH
 handlers _      routing GetNodesType      = Just $ MethodHandler (isGetNodes snd) mkSendNodes $ getNodesH routing
 handlers crypto _       CookieRequestType = Just $ MethodHandler (isCookieRequest snd) mkCookie $ cookieRequestH crypto
+handlers _      _       _ = error "TODO handlers"
 
 nodeSearch :: Client -> Search NodeId (IP,PortNumber) () NodeInfo NodeInfo
 nodeSearch client  = Search
diff --git a/src/Network/Tox/DHT/Transport.hs b/src/Network/Tox/DHT/Transport.hs
index 5ebe8b15..79643fad 100644
--- a/src/Network/Tox/DHT/Transport.hs
+++ b/src/Network/Tox/DHT/Transport.hs
@@ -21,6 +21,7 @@ module Network.Tox.DHT.Transport
     , NoSpam(..)
     , verifyChecksum
     , CookieRequest(..)
+    , CookieResponse(..)
     , Cookie(..)
     , CookieData(..)
     , DHTRequest
diff --git a/src/Network/Tox/Onion/Transport.hs b/src/Network/Tox/Onion/Transport.hs
index 85cf095d..baadbbe8 100644
--- a/src/Network/Tox/Onion/Transport.hs
+++ b/src/Network/Tox/Onion/Transport.hs
@@ -681,8 +681,8 @@ decrypt :: TransportCrypto -> OnionMessage Encrypted -> OnionDestination r -> IO
 decrypt crypto msg addr = do
     (skey,pkey) <- selectKey crypto msg addr
     return $ do
-    msg <- sequenceMessage $ transcode (\n -> decryptMessage (skey,pkey) n . left (senderkey addr)) msg
-    Right (msg, addr)
+        msg <- sequenceMessage $ transcode (\n -> decryptMessage (skey,pkey) n . left (senderkey addr)) msg
+        Right (msg, addr)
 
 senderkey :: OnionDestination r -> t -> (PublicKey, t)
 senderkey addr e = (onionKey addr, e)