3 files changed, 43 insertions, 7 deletions

diff --git a/src/Network/Tox.hs b/src/Network/Tox.hs
index 29591a23..af8114f4 100644
--- a/src/Network/Tox.hs
+++ b/src/Network/Tox.hs
@@ -284,7 +284,7 @@ newTox keydb addr = do
     orouter <- newOnionRouter
     (dhtcrypt,onioncrypt,dtacrypt,cryptonet) <- toxTransport crypto orouter lookupClose udp
     let dhtnet0 = layerTransport (DHT.decrypt crypto) (DHT.encrypt crypto) dhtcrypt
-    dhtclient <- newClient drg dhtnet0 DHT.classify (myAddr routing) (DHT.handlers routing) id
+    dhtclient <- newClient drg dhtnet0 DHT.classify (myAddr routing) (DHT.handlers crypto routing) id
                     $ \client net -> onInbound (DHT.updateRouting client routing orouter) net
 
     orouter <- forkRouteBuilder orouter $ \nid ni -> maybe [] (\(_,ns,_)->ns) <$> DHT.getNodes dhtclient nid ni
diff --git a/src/Network/Tox/DHT/Handlers.hs b/src/Network/Tox/DHT/Handlers.hs
index 4f891316..9cdf0d06 100644
--- a/src/Network/Tox/DHT/Handlers.hs
+++ b/src/Network/Tox/DHT/Handlers.hs
@@ -93,7 +93,7 @@ classify msg = mapMessage (\nonce24 (nonce8,_) -> go msg (TransactionId nonce8 n
     go (DHTGetNodes      {}) = IsQuery GetNodesType
     go (DHTPong          {}) = IsResponse
     go (DHTSendNodes     {}) = IsResponse
-    go (DHTCookieRequest {}) = IsQuery (PacketKind 0x18)
+    go (DHTCookieRequest {}) = IsQuery CookieRequestType
     go (DHTCookie        {}) = IsResponse
     go (DHTDHTRequest    {}) = IsQuery DHTRequestType
 
@@ -183,6 +183,23 @@ getNodesH routing addr (GetNodes nid) = do
 
     k = 4
 
+cookieRequestH :: TransportCrypto -> NodeInfo -> CookieRequest -> IO Cookie
+cookieRequestH crypto ni (CookieRequest remoteUserKey) = do
+    (n24,sym,us) <- atomically $ do
+        n24 <- transportNewNonce crypto
+        sym <- transportSymmetric crypto
+        us <- readTVar $ userKeys crypto
+        return (n24,sym,us)
+    timestamp <- round . (* 1000000) <$> getPOSIXTime
+    let dta = encodePlain $ CookieData
+            { cookieTime  = timestamp
+            , longTermKey = remoteUserKey
+            , dhtKey      = transportPublic crypto
+            }
+        edta = encryptSymmetric sym n24 dta
+    return $ Cookie n24 edta
+
+
 type Message = DHTMessage ((,) Nonce8)
 
 type Client = QR.Client String PacketKind TransactionId NodeInfo Message
@@ -298,9 +315,18 @@ isGetNodes _      _               = Left "Bad GetNodes"
 mkSendNodes :: TransactionId -> NodeInfo -> NodeInfo -> SendNodes -> DHTMessage ((,) Nonce8)
 mkSendNodes tid src dst sendnodes = DHTSendNodes $ wrapAsymm tid src dst (, sendnodes)
 
-handlers :: Routing -> PacketKind -> Maybe Handler
-handlers routing PingType     = Just $ MethodHandler (isPing snd) mkPong pingH
-handlers routing GetNodesType = Just $ MethodHandler (isGetNodes snd) mkSendNodes $ getNodesH routing
+isCookieRequest :: (f CookieRequest -> CookieRequest) -> DHTMessage f -> Either String CookieRequest
+isCookieRequest unpack (DHTCookieRequest a) = Right $ unpack $ asymmData a
+isCookieRequest _      _                    = Left "Bad cookie request"
+
+mkCookie :: TransactionId -> NodeInfo -> NodeInfo -> Cookie -> DHTMessage ((,) Nonce8)
+mkCookie (TransactionId n8 n24) src dst cookie = DHTCookie n24 (n8,cookie)
+
+
+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
 
 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 bebb8ae7..5ebe8b15 100644
--- a/src/Network/Tox/DHT/Transport.hs
+++ b/src/Network/Tox/DHT/Transport.hs
@@ -20,8 +20,9 @@ module Network.Tox.DHT.Transport
     , FriendRequest(..)
     , NoSpam(..)
     , verifyChecksum
-    , CookieRequest
-    , Cookie
+    , CookieRequest(..)
+    , Cookie(..)
+    , CookieData(..)
     , DHTRequest
     , mapMessage
     , encrypt
@@ -378,6 +379,15 @@ data CookieData = CookieData   --   16 (mac)
 instance Sized CookieData where
     size = ConstSize 72
 
+instance Serialize CookieData where
+    get = CookieData <$> get
+                     <*> (id2key <$> get)
+                     <*> (id2key <$> get)
+    put (CookieData tm userkey dhtkey) = do
+        put tm
+        put (key2id userkey)
+        put (key2id dhtkey)
+
 instance Sized CookieRequest where
     size = ConstSize 64 -- 32 byte key + 32 byte padding