1 files changed, 25 insertions, 11 deletions

diff --git a/src/Network/Tox/DHT/Handlers.hs b/src/Network/Tox/DHT/Handlers.hs
index 6bbcfb43..d010f36d 100644
--- a/src/Network/Tox/DHT/Handlers.hs
+++ b/src/Network/Tox/DHT/Handlers.hs
@@ -22,6 +22,7 @@ import qualified Data.ByteString.Char8 as C8
 import qualified Data.ByteString.Base16 as Base16
 import Control.Arrow
 import Control.Monad
+import Control.Concurrent.Lifted.Instrument
 import Control.Concurrent.STM
 import Data.Time.Clock.POSIX  (POSIXTime, getPOSIXTime)
 import Network.Socket
@@ -87,9 +88,17 @@ instance Show PacketKind where
     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
+msgType :: ( Serialize (f DHTRequest)
+           , Serialize (f Cookie), Serialize (f CookieRequest)
+           , Serialize (f SendNodes), Serialize (f GetNodes)
+           , Serialize (f Pong), Serialize (f Ping)
+           ) => DHTMessage f -> PacketKind
+msgType msg = PacketKind $ fst $ dhtMessageType msg
+
+classify :: Client -> Message -> MessageClass String PacketKind TransactionId NodeInfo Message
+classify client (DHTLanDiscovery {}) = IsUnsolicited (lanDiscoveryH client)
+classify client msg = fromMaybe (IsUnknown "unknown")
+                      $ mapMessage (\nonce24 (nonce8,_) -> go msg (TransactionId nonce8 nonce24)) msg
  where
     go (DHTPing          {}) = IsQuery PingType
     go (DHTGetNodes      {}) = IsQuery GetNodesType
@@ -202,6 +211,13 @@ cookieRequestH crypto ni (CookieRequest remoteUserKey) = do
         edta = encryptSymmetric sym n24 dta
     return $ Cookie n24 edta
 
+lanDiscoveryH :: Client -> NodeInfo -> NodeInfo -> IO (Maybe (Message -> Message))
+lanDiscoveryH client _ ni = do
+    forkIO $ do
+        myThreadId >>= flip labelThread "lan-discover-ping"
+        ping client ni
+        return ()
+    return Nothing
 
 type Message = DHTMessage ((,) Nonce8)
 
@@ -295,11 +311,9 @@ getNodes client nid addr = do
     return $ fmap unwrapNodes $ join reply
 
 updateRouting :: Client -> Routing -> OnionRouter -> NodeInfo -> Message -> IO ()
-updateRouting client routing orouter 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
+updateRouting client routing orouter naddr msg
+    | PacketKind 0x21 <- msgType msg = return () -- ignore lan discovery
+    | otherwise                      = do
     case prefer4or6 naddr Nothing of
         Want_IP4 -> updateTable client naddr orouter (routing4 routing) (committee4 routing) (sched4 routing)
         Want_IP6 -> updateTable client naddr orouter (routing6 routing) (committee6 routing) (sched6 routing)
@@ -365,10 +379,10 @@ 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 _      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
-handlers _      _       _ = error "TODO handlers"
+handlers _      _       _                 = error "TODO handlers"
 
 nodeSearch :: Client -> Search NodeId (IP,PortNumber) () NodeInfo NodeInfo
 nodeSearch client  = Search