6 files changed, 70 insertions, 34 deletions

diff --git a/src/Network/BitTorrent/MainlineDHT.hs b/src/Network/BitTorrent/MainlineDHT.hs
index a7359bda..c0413322 100644
--- a/src/Network/BitTorrent/MainlineDHT.hs
+++ b/src/Network/BitTorrent/MainlineDHT.hs
@@ -444,7 +444,7 @@ encodePacket :: Message BValue -> NodeInfo -> (ByteString, SockAddr)
 encodePacket msg ni = ( toStrict $ BE.encode msg
                       , nodeAddr ni )
 
-classify :: Message BValue -> MessageClass String Method TransactionId
+classify :: Message BValue -> MessageClass String Method TransactionId NodeInfo (Message BValue)
 classify (Q { msgID = tid, qryMethod = meth }) = IsQuery meth tid
 classify (R { msgID = tid                   }) = IsResponse tid
 
@@ -566,7 +566,7 @@ newClient swarms addr = do
         outgoingClient = client { clientNet = net { awaitMessage = ($ Nothing) } }
 
         dispatch = DispatchMethods
-            { classifyInbound = classify -- :: x -> MessageClass err meth tid
+            { classifyInbound = classify -- :: x -> MessageClass err meth tid addr x
             , lookupHandler   = handlers -- :: meth -> Maybe (MethodHandler err tid addr x)
             , tableMethods    = mapT     -- :: TransactionMethods tbl tid x
             }
diff --git a/src/Network/QueryResponse.hs b/src/Network/QueryResponse.hs
index fca6d5cc..492b7bb4 100644
--- a/src/Network/QueryResponse.hs
+++ b/src/Network/QueryResponse.hs
@@ -226,11 +226,15 @@ data Client err meth tid addr x = forall tbl. Client
     }
 
 -- | An incoming message can be classified into three cases.
-data MessageClass err meth tid
-    = IsQuery meth tid -- ^ An unsolicited query is handled based on it's /meth/ value.  Any response
-                       -- should include the provided /tid/ value.
-    | IsResponse tid   -- ^ A response to a outgoing query we associated with a /tid/ value.
-    | IsUnknown err    -- ^ None of the above.
+data MessageClass err meth tid addr x
+    = IsQuery meth tid                                    -- ^ An unsolicited query is handled based on it's /meth/ value.  Any response
+                                                          -- should include the provided /tid/ value.
+    | IsResponse tid                                      -- ^ A response to a outgoing query we associated with a /tid/ value.
+    | IsUnsolicited (addr -> addr -> IO (Maybe (x -> x))) -- ^ Transactionless informative packet.  The io action will be invoked
+                                                          -- with the source and destination address of a message.  If it handles the
+                                                          -- message, it should return Nothing. Otherwise, it should return a transform
+                                                          -- (usually /id/) to apply before the next handler examines it.
+    | IsUnknown err                                       -- ^ None of the above.
 
 -- | Handler for an inbound query of type /x/ from an address of type _addr_.
 data MethodHandler err tid addr x = forall a b. MethodHandler
@@ -243,6 +247,7 @@ data MethodHandler err tid addr x = forall a b. MethodHandler
       -- address of the query is provided to the handler.
     , methodAction :: addr -> a -> IO b
     }
+      -- | See also 'IsUnsolicited' which likely makes this constructor unnecessary.
     | forall a. NoReply
     { -- | Parse the query into a more specific type for this method.
       methodParse :: x -> Either err a
@@ -377,7 +382,7 @@ transactionMethods (TableMethods insert delete lookup) generate = TransactionMet
 -- | A set of methods necessary for dispatching incoming packets.
 data DispatchMethods tbl err meth tid addr x = DispatchMethods
     { -- | Classify an inbound packet as a query or response.
-      classifyInbound :: x -> MessageClass err meth tid
+      classifyInbound :: x -> MessageClass err meth tid addr x
       -- | Lookup the handler for a inbound query.
     , lookupHandler :: meth -> Maybe (MethodHandler err tid addr x)
       -- | Methods for handling incoming responses.
@@ -459,6 +464,10 @@ handleMessage (Client net d err pending whoami responseID) addr plain = do
                                          (>>= \m -> do mapM_ (sendMessage net addr) m
                                                        return $! Nothing)
                                          (dispatchQuery m tid' self plain addr)
+            IsUnsolicited action -> do
+                self <- whoami (Just addr)
+                action self addr
+                return Nothing
             IsResponse tid -> do
                 action <- atomically $ do
                     ts0 <- readTVar pending
diff --git a/src/Network/Tox.hs b/src/Network/Tox.hs
index d434360f..d5db6979 100644
--- a/src/Network/Tox.hs
+++ b/src/Network/Tox.hs
@@ -174,7 +174,7 @@ myAddr routing maddr = atomically $ do
 
 newClient :: (DRG g, Show addr, Show meth) =>
               g -> Transport String addr x
-                -> (x -> MessageClass String meth DHT.TransactionId)
+                -> (Client String meth DHT.TransactionId addr x -> x -> MessageClass String meth DHT.TransactionId addr x)
                 -> (Maybe addr -> IO addr)
                 -> (meth -> Maybe (MethodHandler String DHT.TransactionId addr x))
                 -> (forall d. TransactionMethods d DHT.TransactionId addr x -> TransactionMethods d DHT.TransactionId addr x)
@@ -195,8 +195,8 @@ newClient drg net classify selfAddr handlers modifytbl modifynet = do
             let word64mapT = transactionMethods (contramap w64Key w64MapMethods) gen
             map_var <- atomically $ newTVar (drg, Data.Word64Map.empty)
             return $ Left (word64mapT,map_var)
-    let dispatch tbl var handlers = DispatchMethods
-                { classifyInbound = classify
+    let dispatch tbl var handlers client = DispatchMethods
+                { classifyInbound = classify client
                 , lookupHandler = handlers -- var
                 , tableMethods = modifytbl tbl
                 }
@@ -204,7 +204,7 @@ newClient drg net classify selfAddr handlers modifytbl modifynet = do
         mkclient (tbl,var) handlers =
             let client = Client
                     { clientNet           = addHandler eprinter (handleMessage client) $ modifynet client net
-                    , clientDispatcher    = dispatch tbl var handlers
+                    , clientDispatcher    = dispatch tbl var handlers client
                     , clientErrorReporter = eprinter { reportTimeout = reportTimeout ignoreErrors }
                     , clientPending       = var
                     , clientAddress       = selfAddr
@@ -301,7 +301,7 @@ newTox keydb addr = do
         atomically $ newTVar nil { maxInterval = 20 } -- 20 second timeout on announce ping-ids.
     oniondrg <- drgNew
     let onionnet = layerTransportM (Onion.decrypt crypto) (Onion.encrypt crypto) onioncrypt
-    onionclient <- newClient oniondrg onionnet Onion.classify
+    onionclient <- newClient oniondrg onionnet (const Onion.classify)
                     (getOnionAlias crypto $ R.thisNode <$> readTVar (DHT.routing4 routing))
                     (Onion.handlers onionnet routing toks keydb)
                     (hookQueries orouter DHT.transactionKey)
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
diff --git a/src/Network/Tox/Onion/Handlers.hs b/src/Network/Tox/Onion/Handlers.hs
index 76908df8..b06fc2af 100644
--- a/src/Network/Tox/Onion/Handlers.hs
+++ b/src/Network/Tox/Onion/Handlers.hs
@@ -42,7 +42,7 @@ import Data.Functor.Identity
 type Client r = QR.Client String PacketKind TransactionId (OnionDestination r) Message
 type Message = OnionMessage Identity
 
-classify :: Message -> MessageClass String PacketKind TransactionId
+classify :: Message -> MessageClass String PacketKind TransactionId (OnionDestination r) Message
 classify msg = go msg
  where
     go (OnionAnnounce         announce)    = IsQuery AnnounceType