1 files changed, 41 insertions, 0 deletions

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