8 files changed, 158 insertions, 10 deletions

diff --git a/CryptoTransport.hs b/CryptoTransport.hs
index 203d4d82..9d4b3d84 100644
--- a/CryptoTransport.hs
+++ b/CryptoTransport.hs
@@ -12,7 +12,7 @@ module CryptoTransport
     ) where
 
 import Crypto.Tox
-import DHTTransport (Cookie)
+import Network.Tox.DHT.Transport (Cookie)
 
 import Network.Socket
 import Data.ByteString
diff --git a/DHTHandlers.hs b/DHTHandlers.hs
index b98a4d6e..87f5e1b6 100644
--- a/DHTHandlers.hs
+++ b/DHTHandlers.hs
@@ -3,7 +3,7 @@
 {-# LANGUAGE TupleSections              #-}
 module DHTHandlers where
 
-import DHTTransport
+import Network.Tox.DHT.Transport as DHTTransport
 import Network.QueryResponse as QR hiding (Client)
 import qualified Network.QueryResponse as QR (Client)
 import Crypto.Tox
diff --git a/OnionHandlers.hs b/OnionHandlers.hs
new file mode 100644
index 00000000..5e0951e8
--- /dev/null
+++ b/OnionHandlers.hs
@@ -0,0 +1,148 @@
+{-# LANGUAGE PatternSynonyms            #-}
+module OnionHandlers where
+
+import Network.Tox.DHT.Transport
+import DHTHandlers hiding (Message,Client)
+import OnionTransport
+import Network.QueryResponse as QR hiding (Client)
+import qualified Network.QueryResponse as QR (Client)
+import Crypto.Tox
+import qualified Data.Wrapper.PSQ             as PSQ
+         ;import Data.Wrapper.PSQ             (PSQ)
+import Crypto.Error.Types                     (CryptoFailable (..),
+                                               throwCryptoError)
+
+import System.IO
+import qualified Data.ByteArray               as BA
+import Data.Serialize as S
+import qualified Data.Wrapper.PSQInt as Int
+import Network.Kademlia
+import Network.Address               (WantIP (..), ipFamily, testIdBit)
+import qualified Network.DHT.Routing as R
+import Control.TriadCommittee
+import qualified Data.MinMaxPSQ               as MinMaxPSQ
+         ;import Data.MinMaxPSQ               (MinMaxPSQ')
+import Network.BitTorrent.DHT.Token           as Token
+
+import Control.Exception hiding (Handler)
+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
+import Data.Ord
+import Data.Functor.Identity
+
+type Client = QR.Client String PacketKind TransactionId OnionToOwner Message
+type Message = OnionMessage Identity
+
+classify :: Message -> MessageClass String PacketKind TransactionId
+classify msg = go msg
+ where
+    go (OnionAnnounce         announce)    = IsQuery AnnounceType
+                                                $ TransactionId (snd $ runIdentity $ assymData announce)
+                                                                (assymNonce announce)
+    go (OnionAnnounceResponse n8 n24 resp) = IsResponse (TransactionId n8 n24)
+    go (OnionToRoute          {})          = IsQuery DataRequestType (TransactionId (Nonce8 0) (Nonce24 zeros24))
+    go (OnionToRouteResponse  {})          = IsResponse              (TransactionId (Nonce8 0) (Nonce24 zeros24))
+
+-- Toxcore generates `ping_id`s by taking a 32 byte sha hash of the current time,
+-- some secret bytes generated when the instance is created, the current time
+-- divided by a 20 second timeout, the public key of the requester and the source
+-- ip/port that the packet was received from. Since the ip/port that the packet
+-- was received from is in the `ping_id`, the announce packets being sent with a
+-- ping id must be sent using the same path as the packet that we received the
+-- `ping_id` from or announcing will fail.
+--
+-- The reason for this 20 second timeout in toxcore is that it gives a reasonable
+-- time (20 to 40 seconds) for a peer to announce himself while taking in count
+-- all the possible delays with some extra seconds.
+announceH :: Routing -> TVar SessionTokens -> TVar AnnouncedKeys -> OnionToOwner -> AnnounceRequest -> IO AnnounceResponse
+announceH routing toks keydb (OnionToOwner naddr retpath) req = do
+    case () of
+        _ | announcePingId req == zeros32
+          -> go False
+
+        _ -> let Nonce32 bs = announcePingId req
+                 tok        = fromPaddedByteString 32 bs
+             in checkToken toks naddr tok >>= go
+        `catch` (\(SomeException e) -> hPutStrLn stderr ("announceH Exception! "++show e) >> throw e)
+ where
+    go withTok = do
+        ns <- getNodesH routing naddr (GetNodes (announceSeeking req))
+        tm <- getPOSIXTime
+        let storing = (nodeId naddr == announceSeeking req)
+        record <- atomically $ do
+            when (withTok && storing) $ do
+                let toxpath = OnionToOwner naddr{ nodeId = announceKey req } retpath
+                    -- Note: The following distance calculation assumes that
+                    -- our nodeid doesn't change and is the same for both
+                    -- routing4 and routing6.
+                    d = xorNodeId (nodeId (tentativeId routing))
+                                  (announceSeeking req)
+                modifyTVar' keydb (insertKey tm (announceSeeking req) toxpath d)
+            ks <- readTVar keydb
+            return $ snd . snd <$> MinMaxPSQ.lookup' (announceSeeking req) (keyAssoc ks)
+        newtok <- if storing
+                    then Nonce32 . toPaddedByteString 32 <$> grantToken toks naddr
+                    else return $ zeros32
+        let k = case record of
+                Nothing                     -> NotStored newtok
+                Just (OnionToOwner {}) | storing -> Acknowledged newtok
+                Just (OnionToOwner ni _)         -> SendBackKey $ id2key (nodeId ni)
+        let response = AnnounceResponse k ns
+        hPutStrLn stderr $ unwords ["Announce:", show req, "-reply->", show response]
+        return response
+
+dataToRouteH ::
+          TVar AnnouncedKeys
+          -> Transport err OnionToOwner (OnionMessage f)
+          -> addr
+          -> OnionMessage f
+          -> IO ()
+dataToRouteH keydb udp _ (OnionToRoute pub assym) = do
+    let k = key2id pub
+    mb <- atomically $ do
+        ks <- readTVar keydb
+        forM (MinMaxPSQ.lookup' k (keyAssoc ks)) $ \(p,(cnt,rpath)) -> do
+            writeTVar keydb $ ks { keyAssoc = MinMaxPSQ.insert' k (cnt + 1, rpath) p (keyAssoc ks) }
+            return rpath
+    forM_ mb $ \rpath -> do
+        -- forward
+        sendMessage udp rpath $ OnionToRouteResponse assym
+        hPutStrLn stderr $ "Forwarding data-to-route -->"++show k
+
+type NodeDistance = NodeId
+
+data AnnouncedKeys = AnnouncedKeys
+    { keyByAge :: !(PSQ NodeId (Down POSIXTime)) -- timeout of 300 seconds
+    , keyAssoc :: !(MinMaxPSQ' NodeId NodeDistance (Int,OnionToOwner))
+    }
+
+
+insertKey :: POSIXTime -> NodeId -> OnionToOwner -> NodeDistance -> AnnouncedKeys -> AnnouncedKeys
+insertKey tm pub toxpath d keydb = AnnouncedKeys
+    { keyByAge = PSQ.insert pub (Down tm) (keyByAge keydb)
+    , keyAssoc = case MinMaxPSQ.lookup' pub (keyAssoc keydb) of
+                    Just (_,(cnt,_)) -> MinMaxPSQ.insert' pub (cnt,toxpath) d (keyAssoc keydb)
+                    Nothing          -> MinMaxPSQ.insert' pub (0  ,toxpath) d (keyAssoc keydb)
+    }
+
+areq :: Message -> Either String AnnounceRequest
+areq (OnionAnnounce assym) = Right $ fst $ runIdentity $ assymData assym
+areq _                     = Left "Unexpected non-announce OnionMessage"
+
+handlers :: Transport err OnionToOwner Message
+            -> Routing
+            -> TVar SessionTokens
+            -> TVar AnnouncedKeys
+            -> PacketKind
+            -> Maybe (MethodHandler String TransactionId OnionToOwner Message)
+handlers net routing toks keydb AnnounceType
+    = Just
+    $ MethodHandler areq (\(TransactionId n8 n24) src dst -> OnionAnnounceResponse n8 n24 . Identity)
+    $ announceH routing toks keydb
+handlers net _ _ keydb _ = Just $ NoReply Right $ dataToRouteH keydb net
+
diff --git a/OnionTransport.hs b/OnionTransport.hs
index 0ba71922..26130a0c 100644
--- a/OnionTransport.hs
+++ b/OnionTransport.hs
@@ -40,8 +40,8 @@ import Network.Address (fromSockAddr,toSockAddr,setPort,either4or6,sockAddrPort)
 import Network.QueryResponse
 import Crypto.Tox hiding (encrypt,decrypt)
 import Network.Tox.Address
-import qualified ToxCrypto
-import DHTTransport (NodeInfo(..),NodeId(..),SendNodes,nodeInfo,DHTPublicKey)
+import qualified Crypto.Tox as ToxCrypto
+import Network.Tox.DHT.Transport (NodeInfo(..),NodeId(..),SendNodes,nodeInfo,DHTPublicKey)
 
 import Debug.Trace
 import Control.Arrow
diff --git a/Tox.hs b/Tox.hs
index 44703cbd..a8ae8a84 100644
--- a/Tox.hs
+++ b/Tox.hs
@@ -86,7 +86,7 @@ import GHC.TypeLits
 import Crypto.Tox hiding (Assym)
 import ToxTransport
 import Network.Tox.Address
-import qualified DHTTransport as DHT
+import qualified Network.Tox.DHT.Transport as DHT
 import qualified DHTHandlers as DHT
 import qualified OnionTransport as Onion
 import qualified OnionHandlers as Onion
diff --git a/ToxTransport.hs b/ToxTransport.hs
index ea7f6492..a401638a 100644
--- a/ToxTransport.hs
+++ b/ToxTransport.hs
@@ -10,7 +10,7 @@ module ToxTransport (toxTransport) where
 
 import Network.QueryResponse
 import Crypto.Tox
-import DHTTransport
+import Network.Tox.DHT.Transport
 import OnionTransport
 import CryptoTransport
 
diff --git a/examples/dhtd.hs b/examples/dhtd.hs
index 5d4b3b9f..e52d494c 100644
--- a/examples/dhtd.hs
+++ b/examples/dhtd.hs
@@ -65,7 +65,7 @@ import Data.Wrapper.PSQ as PSQ (pattern (:->))
 import qualified Data.Wrapper.PSQ as PSQ
 import Data.Ord
 import Data.Time.Clock.POSIX
-import qualified DHTTransport as Tox
+import qualified Network.Tox.DHT.Transport as Tox
 import qualified DHTHandlers as Tox
 import qualified OnionHandlers as Tox
 import Data.Typeable
diff --git a/DHTTransport.hs b/src/Network/Tox/DHT/Transport.hs
index 189ff0ee..5a2d8a84 100644
--- a/DHTTransport.hs
+++ b/src/Network/Tox/DHT/Transport.hs
@@ -6,11 +6,11 @@
 {-# LANGUAGE TupleSections              #-}
 {-# LANGUAGE TypeOperators              #-}
 {-# LANGUAGE UndecidableInstances #-}
-module DHTTransport
+module Network.Tox.DHT.Transport
     ( parseDHTAddr
     , encodeDHTAddr
     , forwardDHTRequests
-    , module ToxAddress
+    , module Network.Tox.Address
     , DHTMessage(..)
     , Ping(..)
     , Pong(..)
@@ -28,7 +28,7 @@ module DHTTransport
 
 import Network.Tox.Address
 import Crypto.Tox hiding (encrypt,decrypt)
-import qualified ToxCrypto
+import qualified Crypto.Tox as ToxCrypto
 import Network.QueryResponse
 
 import Control.Arrow