1 files changed, 306 insertions, 0 deletions

diff --git a/src/Network/Tox/DHT/Handlers.hs b/src/Network/Tox/DHT/Handlers.hs
new file mode 100644
index 00000000..901da99e
--- /dev/null
+++ b/src/Network/Tox/DHT/Handlers.hs
@@ -0,0 +1,306 @@
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE PatternSynonyms            #-}
+{-# LANGUAGE TupleSections              #-}
+module Network.Tox.DHT.Handlers where
+
+import Network.Tox.DHT.Transport as DHTTransport
+import Network.QueryResponse as QR hiding (Client)
+import qualified Network.QueryResponse as QR (Client)
+import Crypto.Tox
+import Network.BitTorrent.DHT.Search
+import qualified Data.Wrapper.PSQInt as Int
+import Network.Kademlia
+import Network.Address               (WantIP (..), ipFamily, testIdBit,fromSockAddr, sockAddrPort)
+import qualified Network.DHT.Routing as R
+import Control.TriadCommittee
+import System.Global6
+
+import qualified Data.ByteArray as BA
+import qualified Data.ByteString.Char8 as C8
+import qualified Data.ByteString.Base16 as Base16
+import Control.Arrow
+import Control.Monad
+import Control.Concurrent.STM
+import Data.Time.Clock.POSIX  (POSIXTime, getPOSIXTime)
+import Network.Socket
+import Data.Hashable
+import Data.IP
+import Data.Ord
+import Data.Maybe
+import Data.Bits
+import Data.Serialize (Serialize)
+import Data.Word
+import System.IO
+
+data TransactionId = TransactionId
+ { transactionKey :: Nonce8  -- ^ Used to lookup pending query.
+ , cryptoNonce    :: Nonce24 -- ^ Used during the encryption layer.
+ }
+ deriving (Eq,Ord,Show)
+
+newtype PacketKind = PacketKind Word8
+ deriving (Eq, Ord, Serialize)
+
+pattern OnionRequest0Type = PacketKind 128 -- 0x80     Onion Request 0
+pattern OnionRequest1Type = PacketKind 129 -- 0x81     Onion Request 1
+pattern OnionRequest2Type = PacketKind 130 -- 0x82     Onion Request 2
+pattern AnnounceType  = PacketKind 131 -- 0x83     Announce Request
+pattern AnnounceResponseType = PacketKind 132 -- 0x84     Announce Response
+
+pattern DataRequestType = PacketKind 133 -- 0x85     Onion Data Request (data to route request packet)
+pattern DataResponseType = PacketKind 134 -- 0x86     Onion Data Response (data to route response packet)
+-- 0x8c     Onion Response 3
+-- 0x8d     Onion Response 2
+pattern OnionResponse3Type = PacketKind 140 -- 0x8c     Onion Response 3
+pattern OnionResponse2Type = PacketKind 141 -- 0x8d     Onion Response 2
+pattern OnionResponse1Type = PacketKind 142 -- 0x8e     Onion Response 1
+-- 0xf0     Bootstrap Info
+
+pattern DHTRequestType = PacketKind 32 -- 0x20     DHT Request
+
+pattern CookieRequestType = PacketKind 0x18
+pattern CookieResponseType = PacketKind 0x19
+
+pattern PingType      = PacketKind 0 -- 0x00     Ping Request
+pattern PongType      = PacketKind 1 -- 0x01     Ping Response
+pattern GetNodesType  = PacketKind 2 -- 0x02     Nodes Request
+pattern SendNodesType = PacketKind 4 -- 0x04     Nodes Response
+
+
+instance Show PacketKind where
+    showsPrec d PingType = mappend "PingType"
+    showsPrec d PongType = mappend "PongType"
+    showsPrec d GetNodesType = mappend "GetNodesType"
+    showsPrec d SendNodesType = mappend "SendNodesType"
+    showsPrec d DHTRequestType = mappend "DHTRequestType"
+    showsPrec d OnionRequest0Type = mappend "OnionRequest0Type"
+    showsPrec d OnionResponse1Type = mappend "OnionResponse1Type"
+    showsPrec d OnionResponse3Type = mappend "OnionResponse3Type"
+    showsPrec d AnnounceType = mappend "AnnounceType"
+    showsPrec d AnnounceResponseType = mappend "AnnounceResponseType"
+    showsPrec d DataRequestType = mappend "DataRequestType"
+    showsPrec d DataResponseType = mappend "DataResponseType"
+    showsPrec d CookieRequestType = mappend "CookieRequestType"
+    showsPrec d CookieResponseType = mappend "CookieResponseType"
+    showsPrec d (PacketKind x) = mappend "PacketKind " . showsPrec (d+1) x
+
+
+classify :: Message -> 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
+    }
+
+newRouting :: SockAddr -> TransportCrypto
+                       -> (TVar (R.BucketList NodeInfo) -> SockAddr -> STM ()) -- ^ invoked on IPv4 change
+                       -> (TVar (R.BucketList NodeInfo) -> SockAddr -> STM ()) -- ^ invoked on IPv6 change
+                       -> IO Routing
+newRouting addr crypto update4 update6 = do
+    let tentative_ip4 = fromMaybe (IPv4 $ toEnum 0) (IPv4 <$> fromSockAddr addr)
+        tentative_ip6 = fromMaybe (IPv6 $ toEnum 0) (IPv6 <$> fromSockAddr addr)
+        tentative_info = NodeInfo
+                { nodeId   = key2id $ transportPublic crypto
+                , nodeIP   = fromMaybe (toEnum 0) (fromSockAddr addr)
+                , nodePort = fromMaybe 0 $ sockAddrPort addr
+                }
+        tentative_info4 = tentative_info { nodeIP = tentative_ip4 }
+    tentative_info6 <-
+        maybe (tentative_info { nodeIP = tentative_ip6 })
+              (\ip6 -> tentative_info { nodeIP = IPv6 ip6 })
+            <$> global6
+    atomically $ do
+        let nobkts = R.defaultBucketCount :: Int
+        tbl4 <- newTVar $ R.nullTable (comparing nodeId) (\s -> hashWithSalt s . nodeId) tentative_info4 nobkts
+        tbl6 <- newTVar $ R.nullTable (comparing nodeId) (\s -> hashWithSalt s . nodeId) tentative_info6 nobkts
+        committee4 <- newTriadCommittee (update4 tbl4) -- $ updateIPVote tbl4 addr4
+        committee6 <- newTriadCommittee (update6 tbl6) -- $ updateIPVote tbl6 addr6
+        sched4 <- newTVar Int.empty
+        sched6 <- newTVar Int.empty
+        return $ Routing tentative_info sched4 tbl4 committee4 sched6 tbl6 committee6
+
+
+-- TODO: This should cover more cases
+isLocal :: IP -> Bool
+isLocal (IPv6 ip6) = (ip6 == toEnum 0)
+isLocal (IPv4 ip4) = (ip4 == toEnum 0)
+
+isGlobal :: IP -> Bool
+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  = testNodeIdBit
+    , R.kademliaXor      = xorNodeId
+    , R.kademliaSample   = sampleNodeId
+    }
+
+
+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 -> (Nonce8 -> dta) -> Assym dta
+wrapAssym (TransactionId n8 n24) src dst dta = Assym
+    { senderKey = id2key $ nodeId src
+    , assymNonce = n24
+    , assymData = dta n8
+    }
+
+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
+
+showHex :: BA.ByteArrayAccess ba => ba -> String
+showHex bs = C8.unpack $ Base16.encode $ BA.convert bs
+
+ping :: Client -> NodeInfo -> IO Bool
+ping client addr = do
+    hPutStrLn stderr $ show addr ++ " <-- ping"
+    reply <- QR.sendQuery client (serializer PingType DHTPing unpong) Ping addr
+    hPutStrLn stderr $ show addr ++ " -pong-> " ++ show reply
+    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
+    hPutStrLn stderr $ show addr ++ " <-- getnodes " ++ show nid
+    reply <- QR.sendQuery client (serializer GetNodesType DHTGetNodes unsendNodes) (GetNodes nid) addr
+    hPutStrLn stderr $ show addr ++ " -sendnodes-> " ++ show reply
+    return $ fmap unwrapNodes $ join reply
+
+updateRouting :: Client -> Routing -> NodeInfo -> Message -> IO ()
+updateRouting client routing naddr msg = do
+    let typ = fst $ dhtMessageType $ fst $ DHTTransport.encrypt (error "updateRouting") msg naddr
+        tid = mapMessage (\n24 (n8,_) -> TransactionId n8 n24) msg
+    hPutStrLn stderr $ "updateRouting "++show (typ,tid)
+    -- 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 ()
+
+type Handler = MethodHandler String TransactionId NodeInfo Message
+
+isPing :: (f Ping -> Ping) ->  DHTMessage f -> Either String Ping
+isPing unpack (DHTPing a) = Right $ unpack $ assymData a
+isPing _      _           = Left "Bad ping"
+
+mkPong :: TransactionId -> NodeInfo -> NodeInfo -> Pong -> DHTMessage ((,) Nonce8)
+mkPong tid src dst pong = DHTPong $ wrapAssym tid src dst (, pong)
+
+isGetNodes :: (f GetNodes -> GetNodes) ->  DHTMessage f -> Either String GetNodes
+isGetNodes unpack (DHTGetNodes a) = Right $ unpack $ assymData a
+isGetNodes _      _               = Left "Bad GetNodes"
+
+mkSendNodes :: TransactionId -> NodeInfo -> NodeInfo -> SendNodes -> DHTMessage ((,) Nonce8)
+mkSendNodes tid src dst sendnodes = DHTSendNodes $ wrapAssym 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
+
+nodeSearch :: Client -> Search NodeId (IP,PortNumber) () NodeInfo NodeInfo
+nodeSearch client  = Search
+    { searchSpace       = toxSpace
+    , searchNodeAddress = nodeIP &&& nodePort
+    , searchQuery       = getNodes client
+    }