1 files changed, 7 insertions, 1 deletions

diff --git a/dht/src/Network/Tox.hs b/dht/src/Network/Tox.hs
index b396c2ea..97b97bad 100644
--- a/dht/src/Network/Tox.hs
+++ b/dht/src/Network/Tox.hs
@@ -52,6 +52,7 @@ import System.Global6
 import Network.Address                        (WantIP (..),IP,getBindAddress)
 import qualified Network.Kademlia.Routing          as R
 import Network.QueryResponse
+import Network.StreamServer                  (ServerHandle,quitListening)
 import Crypto.Tox
 import Data.Word64Map                        (fitsInInt)
 import qualified Data.Word64Map              (empty)
@@ -183,6 +184,7 @@ data Tox extra  = Tox
     , toxContactInfo    :: ContactInfo extra
     , toxAnnounceToLan  :: IO ()
     , toxBindAddress    :: SockAddr
+    , toxRelayServer    :: Maybe ServerHandle
     }
 
 
@@ -291,7 +293,9 @@ newTox keydb bindspecs onsess suppliedDHTKey tcp = do
     addr <- getSocketName sock
     (relay,sendTCP) <- tcpRelay addr (\a x -> sendMessage udp a $ S.runPut $ Onion.putRequest x)
     tox <- newToxOverTransport keydb addr onsess suppliedDHTKey udp sendTCP
-    return tox { toxAnnounceToLan = announceToLan sock (key2id $ transportPublic $ toxCryptoKeys tox) }
+    return tox { toxAnnounceToLan = announceToLan sock (key2id $ transportPublic $ toxCryptoKeys tox)
+               , toxRelayServer   = Just relay
+               }
 
 -- | This version of 'newTox' is useful for automated tests using 'testPairTransport'.
 newToxOverTransport :: TVar Onion.AnnouncedKeys
@@ -375,6 +379,7 @@ newToxOverTransport keydb addr onNewSession suppliedDHTKey udp tcp = do
         , toxContactInfo    = roster
         , toxAnnounceToLan  = return ()
         , toxBindAddress    = addr
+        , toxRelayServer    = Nothing
         }
 
 onionTimeout :: Tox extra -> DHT.TransactionId -> Onion.OnionDestination RouteId -> STM (Onion.OnionDestination RouteId, Int)
@@ -435,6 +440,7 @@ forkTox tox with_avahi = do
                 quitRouteBuilder (toxOnionRoutes tox)
                 quitToRoute
                 quitHs
+                mapM_ quitListening (toxRelayServer tox)
            , bootstrap (DHT.refresher4 $ toxRouting tox)
            , bootstrap (DHT.refresher6 $ toxRouting tox)
            )