3 files changed, 212 insertions, 0 deletions

diff --git a/src/Network/DHT/Mainline.hs b/src/Network/DHT/Mainline.hs
index 9af42a6d..bdf9e8b2 100644
--- a/src/Network/DHT/Mainline.hs
+++ b/src/Network/DHT/Mainline.hs
@@ -56,6 +56,8 @@
 --   <http://pdos.csail.mit.edu/~petar/papers/maymounkov-kademlia-lncs.pdf>
 --
 {-# LANGUAGE CPP #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE DeriveDataTypeable #-}
 {-# LANGUAGE FlexibleInstances  #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
@@ -94,6 +96,7 @@ module Network.DHT.Mainline
        , checkToken
        ) where
 
+import Data.String
 import Control.Applicative
 import Data.Bool
 #ifdef VERSION_bencoding
@@ -570,3 +573,7 @@ instance Kademlia KMessageOf where
     initializeDHTData = TorrentData
                 <$> newTVarIO def
                 <*> (newTVarIO =<< nullSessionTokens)
+
+deriving instance IsString (QueryMethod dht) => IsString (Method dht param result)
+deriving instance BEncode (QueryMethod dht) => BEncode (Method dht param result)
+
diff --git a/src/Network/DHT/Tox.hs b/src/Network/DHT/Tox.hs
new file mode 100644
index 00000000..d6fc866f
--- /dev/null
+++ b/src/Network/DHT/Tox.hs
@@ -0,0 +1,112 @@
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses #-}
+module Network.DHT.Tox where
+
+import Data.Serialize
+import Data.Default
+import Text.PrettyPrint.HughesPJClass
+
+import Network.DHT.Types
+import Network.DatagramServer.Types
+import qualified Network.DatagramServer.Tox as Tox
+import Network.KRPC.Method
+import Data.Word
+import Data.ByteString (ByteString)
+import Data.IP
+import Data.Bool
+import Data.Maybe
+import Control.Monad
+import System.Random
+
+instance Kademlia Tox.Message where
+    data DHTData Tox.Message ip = ToxData
+    namePing      _   = Tox.Ping
+    nameFindNodes _   = Tox.GetNodes
+    initializeDHTData = return ToxData
+
+instance Pretty (NodeId Tox.Message) where
+    pPrint (Tox.NodeId nid) = encodeHexDoc nid
+
+instance Serialize (Query Tox.Message (Ping Tox.Message)) where
+    get = getToxPing False Network.DHT.Types.Query Tox.QueryNonce
+    put (Network.DHT.Types.Query extra Ping) = putToxPing False (Tox.qryNonce extra)
+instance Serialize (Response Tox.Message (Ping Tox.Message)) where
+    get = getToxPing True Network.DHT.Types.Response Tox.ResponseNonce
+    put (Network.DHT.Types.Response extra Ping) = putToxPing True (Tox.rspNonce extra)
+
+instance Serialize (Query Tox.Message (FindNode Tox.Message ip)) where
+    get = do
+        nid <- get
+        n8 <- get
+        return $ Network.DHT.Types.Query (Tox.QueryNonce n8) (FindNode nid)
+    put (Network.DHT.Types.Query (Tox.QueryNonce n8) (FindNode nid)) = do
+        put nid
+        put n8
+instance Serialize (Response Tox.Message (NodeFound Tox.Message IPv4)) where
+    get = do
+        num <- get :: Get Word8
+        when (num > 4) $ fail "Too many nodes in Tox get-nodes reply"
+        ns0 <- sequence $ replicate (fromIntegral num) (nodeFormatToNodeInfo <$> get)
+        -- TODO: Allow tcp and ipv6. For now filtering to udp ip4...
+        let ns = flip mapMaybe ns0 $ \(NodeInfo nid addr u) -> do
+                    guard $ not u
+                    ip4 <- fromAddr addr
+                    return $ NodeInfo nid ip4 ()
+        n8 <- get
+        return $ Network.DHT.Types.Response (Tox.ResponseNonce n8) $ NodeFound ns
+    put (Network.DHT.Types.Response (Tox.ResponseNonce n8) (NodeFound ns)) = do
+        put ( fromIntegral (length ns) :: Word8 )
+        forM_ ns $ \(NodeInfo nid ip4 ()) -> do
+            put Tox.NodeFormat { nodePublicKey = nid
+                               , nodeIsTCP     = False
+                               , nodeIP        = IPv4 (nodeHost ip4)
+                               , nodePort      = nodePort ip4
+                               }
+        put n8
+
+instance KRPC Tox.Message (Query    Tox.Message (FindNode  Tox.Message IPv4))
+                          (Response Tox.Message (NodeFound Tox.Message IPv4)) where
+    method = Method Tox.GetNodes
+    validateExchange = validateToxExchange
+    makeQueryExtra _ _ _ _ = Tox.QueryNonce <$> randomIO
+    makeResponseExtra _ _ q _ = return $ Tox.ResponseNonce $ Tox.qryNonce $ queryExtra q
+    messageSender q _ = Tox.msgClient q
+    messageResponder _ r = Tox.msgClient r
+
+instance KRPC Tox.Message (Query    Tox.Message (Ping Tox.Message))
+                          (Response Tox.Message (Ping Tox.Message)) where
+    method = Method Tox.Ping
+    validateExchange = validateToxExchange
+    makeQueryExtra _ _ _ _ = Tox.QueryNonce <$> randomIO
+    makeResponseExtra _ _ q _ = return $ Tox.ResponseNonce $ Tox.qryNonce $ queryExtra q
+    messageSender q _ = Tox.msgClient q
+    messageResponder _ r = Tox.msgClient r
+
+instance DataHandlers ByteString Tox.Message
+
+instance Default Bool where def = False
+
+getToxPing isPong c n = do
+    q'r <- get :: Get Word8
+    when (bool 0 1 isPong /= q'r) $
+        fail "Tox ping/pong parse fail."
+    n8 <- get :: Get Tox.Nonce8
+    return $ c (n n8) Ping
+
+putToxPing isPong n8 = do
+    put (bool 0 1 isPong :: Word8)
+    put n8
+
+validateToxExchange q r = qnonce == rnonce
+      where
+        qnonce = Tox.qryNonce . queryExtra    . Tox.msgPayload $ q
+        rnonce = Tox.rspNonce . responseExtra . Tox.msgPayload $ r
+
+
+nodeFormatToNodeInfo nf = NodeInfo nid addr u
+ where
+    u    = Tox.nodeIsTCP nf
+    addr = NodeAddr (Tox.nodeIP nf) (Tox.nodePort nf)
+    nid  = Tox.nodePublicKey nf
+
+-- instance Default Bool where def = False
diff --git a/src/Network/DHT/Types.hs b/src/Network/DHT/Types.hs
index 47f98ebe..287d5680 100644
--- a/src/Network/DHT/Types.hs
+++ b/src/Network/DHT/Types.hs
@@ -1,3 +1,6 @@
+{-# LANGUAGE DefaultSignatures #-}
+{-# LANGUAGE GADTs, MultiParamTypeClasses #-}
+{-# LANGUAGE FunctionalDependencies #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE StandaloneDeriving #-}
@@ -14,6 +17,11 @@ import Network.DatagramServer.Types
 import Network.DHT.Routing
 import Data.Typeable
 import GHC.Generics
+import Data.Serialize
+import Data.Hashable
+import Data.String
+import Data.Monoid
+import Data.Char
 
 data TableParameters msg ip u = TableParameters
     { maxBuckets :: Int
@@ -81,3 +89,88 @@ class Kademlia dht where
     namePing :: Proxy dht -> QueryMethod dht
     nameFindNodes :: Proxy dht -> QueryMethod dht
     initializeDHTData :: IO (DHTData dht ip)
+data MethodHandler raw dht ip =
+        forall a b. ( SerializableTo raw (Response dht b)
+                    , SerializableTo raw (Query dht a)
+                    , KRPC dht (Query dht a) (Response dht b)
+                    ) => MethodHandler (QueryMethod dht) (NodeAddr ip -> a -> IO b)
+
+class DataHandlers raw dht where
+    dataHandlers ::
+        ( Ord ip , Hashable ip, Typeable ip, Serialize ip) =>
+         (NodeId dht  -> IO [NodeInfo dht ip ()])
+        -> DHTData dht ip
+        -> [MethodHandler raw dht ip]
+    dataHandlers _ _ = []
+
+-- | Method datatype used to describe method name, parameters and
+-- return values of procedure. Client use a method to /invoke/, server
+-- /implements/ the method to make the actual work.
+--
+--   We use the following fantom types to ensure type-safiety:
+--
+--     * param: Type of method parameters.
+--
+--     * result: Type of return value of the method.
+--
+newtype Method dht param result = Method { methodName :: QueryMethod dht }
+
+deriving instance Eq (QueryMethod dht) => Eq (Method dht param result)
+deriving instance Ord (QueryMethod dht) => Ord (Method dht param result)
+
+-- | Example:
+--
+--   @show (Method \"concat\" :: [Int] Int) == \"concat :: [Int] -> Int\"@
+--
+instance (Show (QueryMethod dht), Typeable a, Typeable b) => Show (Method dht a b) where
+  showsPrec _ = showsMethod
+
+showsMethod :: forall dht a b. ( Show (QueryMethod dht), Typeable a , Typeable b ) => Method dht a b -> ShowS
+showsMethod (Method name) =
+    -- showString (BC.unpack name) <>
+    shows (show name) <>
+    showString " :: " <>
+    shows paramsTy <>
+    showString " -> " <>
+    shows valuesTy
+  where
+    impossible = error "KRPC.showsMethod: impossible"
+    paramsTy = typeOf (impossible :: a)
+    valuesTy = typeOf (impossible :: b)
+
+-- | In order to perform or handle KRPC query you need to provide
+--   corresponding 'KRPC' class.
+--
+--   Example:
+--
+--   @
+--   data Ping = Ping Text deriving BEncode
+--   data Pong = Pong Text deriving BEncode
+--
+--   instance 'KRPC' Ping Pong where
+--     method = \"ping\"
+--   @
+--
+class ( Typeable req, Typeable resp, Envelope dht)
+    => KRPC dht req resp | req -> resp, resp -> req where
+
+  -- | Method name. Default implementation uses lowercased @req@
+  -- datatype name.
+  --
+  method :: Method dht req resp
+
+  -- TODO add underscores
+  default method :: (IsString (QueryMethod dht), Typeable req) => Method dht req resp
+  method = Method $ fromString $ map toLower $ show $ typeOf hole
+    where
+      hole = error "krpc.method: impossible" :: req
+
+
+  validateExchange :: dht req -> dht resp -> Bool
+  validateExchange _ _ = True
+
+  makeQueryExtra :: DHTData dht ip -> NodeId dht -> Proxy req -> Proxy resp -> IO (QueryExtra dht)
+  makeResponseExtra :: DHTData dht ip -> NodeId dht -> req -> Proxy resp -> IO (ResponseExtra dht)
+
+  messageSender :: dht req -> Proxy resp -> NodeId dht
+  messageResponder :: Proxy req -> dht resp -> NodeId dht