2 files changed, 47 insertions, 20 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
diff --git a/src/Network/Tox/DHT/Transport.hs b/src/Network/Tox/DHT/Transport.hs
index dd2838f2..736e84d1 100644
--- a/src/Network/Tox/DHT/Transport.hs
+++ b/src/Network/Tox/DHT/Transport.hs
@@ -68,6 +68,7 @@ data DHTMessage (f :: * -> *)
     | DHTCookieRequest (Asymm (f CookieRequest))
     | DHTCookie Nonce24 (f Cookie)
     | DHTDHTRequest PublicKey (Asymm (f DHTRequest))
+    | DHTLanDiscovery NodeId
 
 deriving instance ( Show (f Cookie)
                   , Show (Asymm (f Ping))
@@ -78,14 +79,15 @@ deriving instance ( Show (f Cookie)
                   , Show (Asymm (f DHTRequest))
                   ) => Show (DHTMessage f)
 
-mapMessage :: forall f b. (forall a. Nonce24 -> f a -> b) -> DHTMessage f -> b
-mapMessage f (DHTPing          a) = f (asymmNonce a) (asymmData a)
-mapMessage f (DHTPong          a) = f (asymmNonce a) (asymmData a)
-mapMessage f (DHTGetNodes      a) = f (asymmNonce a) (asymmData a)
-mapMessage f (DHTSendNodes     a) = f (asymmNonce a) (asymmData a)
-mapMessage f (DHTCookieRequest a) = f (asymmNonce a) (asymmData a)
-mapMessage f (DHTDHTRequest  _ a) = f (asymmNonce a) (asymmData a)
-mapMessage f (DHTCookie nonce fcookie) = f nonce fcookie
+mapMessage :: forall f b. (forall a. Nonce24 -> f a -> b) -> DHTMessage f -> Maybe b
+mapMessage f (DHTPing          a)      = Just $ f (asymmNonce a) (asymmData a)
+mapMessage f (DHTPong          a)      = Just $ f (asymmNonce a) (asymmData a)
+mapMessage f (DHTGetNodes      a)      = Just $ f (asymmNonce a) (asymmData a)
+mapMessage f (DHTSendNodes     a)      = Just $ f (asymmNonce a) (asymmData a)
+mapMessage f (DHTCookieRequest a)      = Just $ f (asymmNonce a) (asymmData a)
+mapMessage f (DHTDHTRequest  _ a)      = Just $ f (asymmNonce a) (asymmData a)
+mapMessage f (DHTCookie nonce fcookie) = Just $ f nonce fcookie
+mapMessage f (DHTLanDiscovery nid)     = Nothing
 
 
 instance Sized Ping where size = ConstSize 1
@@ -109,12 +111,20 @@ parseDHTAddr crypto (msg,saddr)
                     either (const Nothing) Just $ nodeInfo (key2id key) saddr
         left $ fanGet bs getCookie (uncurry DHTCookie) (const $ ni)
     0x20 -> left $ fanGet bs getDHTReqest (uncurry DHTDHTRequest) (asymNodeInfo saddr . snd)
+    0x21 -> left $ do
+               nid <- runGet get bs
+               ni <- nodeInfo nid saddr
+               return (DHTLanDiscovery nid, ni)
     _ -> right
 
 encodeDHTAddr :: (DHTMessage Encrypted8,NodeInfo) -> IO (ByteString, SockAddr)
 encodeDHTAddr (msg,ni) = return (runPut $ putMessage msg, nodeAddr ni)
 
-dhtMessageType :: DHTMessage Encrypted8 -> ( Word8, Put )
+dhtMessageType :: ( Serialize (f DHTRequest)
+                  , Serialize (f Cookie), Serialize (f CookieRequest)
+                  , Serialize (f SendNodes), Serialize (f GetNodes)
+                  , Serialize (f Pong), Serialize (f Ping)
+                  ) => DHTMessage f -> (Word8, Put)
 dhtMessageType (DHTPing           a) = (0x00, putAsymm a)
 dhtMessageType (DHTPong           a) = (0x01, putAsymm a)
 dhtMessageType (DHTGetNodes       a) = (0x02, putAsymm a)
@@ -122,6 +132,7 @@ dhtMessageType (DHTSendNodes      a) = (0x04, putAsymm a)
 dhtMessageType (DHTCookieRequest  a) = (0x18, putAsymm a)
 dhtMessageType (DHTCookie       n x) = (0x19, put n >> put x)
 dhtMessageType (DHTDHTRequest   k a) = (0x20, putPublicKey k >> putAsymm a)
+dhtMessageType (DHTLanDiscovery nid) = (0x21, put nid)
 
 putMessage :: DHTMessage Encrypted8 -> Put
 putMessage msg = case dhtMessageType msg of (typ,p) -> put typ >> p
@@ -452,6 +463,7 @@ sequenceMessage (DHTSendNodes asym)         = fmap DHTSendNodes     $ sequenceA
 sequenceMessage (DHTCookieRequest asym)     = fmap DHTCookieRequest $ sequenceA $ fmap uncomposed asym
 sequenceMessage (DHTCookie n dta)           = fmap (DHTCookie n)    $ uncomposed dta
 sequenceMessage (DHTDHTRequest pubkey asym) = fmap (DHTDHTRequest pubkey) $ sequenceA $ fmap uncomposed asym
+sequenceMessage (DHTLanDiscovery nid)       = pure $ DHTLanDiscovery nid
 
 transcode :: forall f g. (forall a. Serialize a => Nonce24 -> Either (f a) (Asymm (f a)) -> g a) -> DHTMessage f -> DHTMessage g
 transcode f (DHTPing asym)              = DHTPing              $ asym { asymmData = f (asymmNonce asym) (Right asym) }
@@ -461,3 +473,4 @@ transcode f (DHTSendNodes asym)         = DHTSendNodes         $ asym { asymmDat
 transcode f (DHTCookieRequest asym)     = DHTCookieRequest     $ asym { asymmData = f (asymmNonce asym) (Right asym) }
 transcode f (DHTCookie n dta)           = DHTCookie n          $ f n  $ Left dta
 transcode f (DHTDHTRequest pubkey asym) = DHTDHTRequest pubkey $ asym { asymmData = f (asymmNonce asym) (Right asym) }
+transcode f (DHTLanDiscovery nid)       = DHTLanDiscovery nid