1 files changed, 182 insertions, 0 deletions

diff --git a/DHTHandlers.hs b/DHTHandlers.hs
new file mode 100644
index 00000000..d98a5e60
--- /dev/null
+++ b/DHTHandlers.hs
@@ -0,0 +1,182 @@
+{-# LANGUAGE PatternSynonyms            #-}
+module DHTHandlers where
+
+import DHTTransport
+import Network.QueryResponse as QR hiding (Client)
+import qualified Network.QueryResponse as QR (Client)
+import ToxCrypto
+import ToxMessage  as Tox (PacketKind(..), pattern PingType, pattern GetNodesType, pattern DHTRequestType)
+
+import qualified Data.Wrapper.PSQInt as Int
+import Kademlia
+import Network.Address               (WantIP (..), ipFamily, testIdBit)
+import qualified Network.DHT.Routing as R
+import TriadCommittee
+
+import Control.Monad
+import Control.Concurrent.STM
+import Data.Time.Clock.POSIX  (POSIXTime, getPOSIXTime)
+import Network.Socket
+import Data.IP
+import Data.Maybe
+import Data.Bits
+
+data TransactionId = TransactionId
+ { transactionKey :: Nonce8  -- ^ Used to lookup pending query.
+ , cryptoNonce    :: Nonce24 -- ^ Used during the encryption layer.
+ }
+
+classify :: DHTMessage ((,) Nonce8) -> MessageClass String PacketKind TransactionId
+classify msg = mapMessage (\nonce24 (nonce8,_) -> go msg (TransactionId nonce8 nonce24)) msg
+ where
+    go (DHTPing          {}) = IsQuery PingType
+    go (DHTGetNodes      {}) = IsQuery GetNodesType
+    go (DHTPong          {}) = IsResponse
+    go (DHTSendNodes     {}) = IsResponse
+    go (DHTCookieRequest {}) = IsQuery (PacketKind 0x18)
+    go (DHTCookie        {}) = IsResponse
+    go (DHTDHTRequest    {}) = IsQuery DHTRequestType
+
+data Routing = Routing
+    { tentativeId :: NodeInfo
+    , sched4      :: !( TVar (Int.PSQ POSIXTime) )
+    , routing4    :: !( TVar (R.BucketList NodeInfo) )
+    , committee4  :: TriadCommittee NodeId SockAddr
+    , sched6      :: !( TVar (Int.PSQ POSIXTime) )
+    , routing6    :: !( TVar (R.BucketList NodeInfo) )
+    , committee6  :: TriadCommittee NodeId SockAddr
+    }
+
+
+-- TODO: This should cover more cases
+isLocal (IPv6 ip6) = (ip6 == toEnum 0)
+isLocal (IPv4 ip4) = (ip4 == toEnum 0)
+
+isGlobal = not . isLocal
+
+prefer4or6 :: NodeInfo -> Maybe WantIP -> WantIP
+prefer4or6 addr iptyp = fromMaybe (ipFamily $ nodeIP addr) iptyp
+
+toxSpace :: R.KademliaSpace NodeId NodeInfo
+toxSpace = R.KademliaSpace
+    { R.kademliaLocation = nodeId
+    , R.kademliaTestBit  = testIdBit
+    , R.kademliaXor      = xor
+    }
+
+
+pingH :: NodeInfo -> Ping -> IO Pong
+pingH _ Ping = return Pong
+
+getNodesH :: Routing -> NodeInfo -> GetNodes -> IO SendNodes
+getNodesH routing addr (GetNodes nid) = do
+    let preferred = prefer4or6 addr Nothing
+
+    (append4,append6) <- atomically $ do
+        ni4 <- R.thisNode <$> readTVar (routing4 routing)
+        ni6 <- R.thisNode <$> readTVar (routing6 routing)
+        return $ case ipFamily (nodeIP addr) of
+            Want_IP4 | isGlobal (nodeIP ni6) -> (id, (++ [ni6]))
+            Want_IP6 | isGlobal (nodeIP ni4) -> ((++ [ni4]), id)
+            _                                -> (id, id)
+    ks  <- go append4 $ routing4 routing
+    ks6 <- go append6 $ routing6 routing
+    let (ns1,ns2) = case preferred of Want_IP6 -> (ks6,ks)
+                                      Want_IP4 -> (ks,ks6)
+    return $ SendNodes
+           $ if null ns2 then ns1
+                         else take 4 (take 3 ns1 ++ ns2)
+ where
+    go f var = f . R.kclosest toxSpace k nid <$> atomically (readTVar var)
+
+    k = 4
+
+type Message = DHTMessage ((,) Nonce8)
+
+type Client = QR.Client String PacketKind TransactionId NodeInfo Message
+
+
+wrapAssym :: TransactionId -> NodeInfo -> NodeInfo -> dta -> Assym (Nonce8,dta)
+wrapAssym (TransactionId n8 n24) src dst dta = Assym
+    { senderKey = let NodeId pubkey = nodeId src in pubkey
+    , assymNonce = n24
+    , assymData = (n8, dta)
+    }
+
+serializer :: PacketKind
+                -> (Assym (Nonce8,ping) -> Message)
+                -> (Message -> Maybe (Assym (Nonce8,pong)))
+                -> MethodSerializer TransactionId NodeInfo Message PacketKind ping (Maybe pong)
+serializer pktkind mkping mkpong = MethodSerializer
+        { methodTimeout  = 5
+        , method         = pktkind
+        -- wrapQuery :: tid -> addr -> addr -> qry -> x
+        , wrapQuery      = \tid src dst ping -> mkping $ wrapAssym tid src dst ping
+        -- unwrapResponse :: x -> b
+        , unwrapResponse = fmap (snd . assymData) . mkpong
+        }
+
+
+unpong :: Message -> Maybe (Assym (Nonce8,Pong))
+unpong (DHTPong assym) = Just assym
+unpong _               = Nothing
+
+ping :: Client -> NodeInfo -> IO Bool
+ping client addr = do
+    reply <- QR.sendQuery client (serializer PingType DHTPing unpong) Ping addr
+    maybe (return False) (\Pong -> return True) $ join reply
+
+unsendNodes :: Message -> Maybe (Assym (Nonce8,SendNodes))
+unsendNodes (DHTSendNodes assym) = Just assym
+unsendNodes _                    = Nothing
+
+unwrapNodes :: SendNodes -> ( [NodeInfo], [NodeInfo], () )
+unwrapNodes (SendNodes ns) = (ns,ns,())
+
+getNodes :: Client -> NodeId -> NodeInfo -> IO (Maybe ([NodeInfo],[NodeInfo],()))
+getNodes client nid addr = do
+    reply <- QR.sendQuery client (serializer GetNodesType DHTGetNodes unsendNodes) (GetNodes nid) addr
+    return $ fmap unwrapNodes $ join reply
+
+updateRouting :: Client -> Routing -> NodeInfo -> Message -> IO ()
+updateRouting client routing naddr msg = do
+    -- hPutStrLn stderr $ "updateRouting "++show typ
+    -- TODO: check msg type
+    case prefer4or6 naddr Nothing of
+        Want_IP4 -> updateTable client naddr (routing4 routing) (committee4 routing) (sched4 routing)
+        Want_IP6 -> updateTable client naddr (routing6 routing) (committee6 routing) (sched6 routing)
+
+updateTable :: Client -> NodeInfo -> TVar (R.BucketList NodeInfo) -> TriadCommittee NodeId SockAddr -> TVar (Int.PSQ POSIXTime) -> IO ()
+updateTable client naddr tbl committee sched = do
+        self <- atomically $ R.thisNode <$> readTVar tbl
+        when (nodeIP self /= nodeIP naddr) $ do
+            -- TODO: IP address vote?
+            insertNode (toxKademlia client committee tbl sched) naddr
+
+toxKademlia :: Client -> TriadCommittee NodeId SockAddr -> TVar (R.BucketList NodeInfo) -> TVar (Int.PSQ POSIXTime) -> Kademlia NodeId NodeInfo
+toxKademlia client committee var sched
+    = Kademlia quietInsertions
+               toxSpace
+               (vanillaIO var $ ping client)
+                   { tblTransition = \tr -> do
+                        io1 <- transitionCommittee committee tr
+                        io2 <- touchBucket toxSpace (15*60) var sched tr
+                        return $ do
+                            io1 >> io2
+                            {-
+                            hPutStrLn stderr $ unwords
+                                [ show (transitionedTo tr)
+                                , show (transitioningNode tr)
+                                ]
+                            -}
+                            return ()
+                   }
+
+transitionCommittee :: TriadCommittee NodeId SockAddr -> RoutingTransition NodeInfo -> STM (IO ())
+transitionCommittee committee (RoutingTransition ni Stranger) = do
+    delVote committee (nodeId ni)
+    return $ do
+        -- hPutStrLn stderr $ "delVote "++show (nodeId ni)
+        return ()
+transitionCommittee committee _ = return $ return ()
+