4 files changed, 76 insertions, 49 deletions

diff --git a/src/Data/Tox/Relay.hs b/src/Data/Tox/Relay.hs
index d1e9fb99..c563db8d 100644
--- a/src/Data/Tox/Relay.hs
+++ b/src/Data/Tox/Relay.hs
@@ -4,6 +4,7 @@
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE KindSignatures             #-}
 {-# LANGUAGE MultiParamTypeClasses      #-}
+{-# LANGUAGE PatternSynonyms            #-}
 {-# LANGUAGE StandaloneDeriving         #-}
 {-# LANGUAGE UndecidableInstances       #-}
 module Data.Tox.Relay where
@@ -61,9 +62,15 @@ data RelayPacket
     | RelayData ByteString ConId
  deriving (Eq,Ord,Show,Data)
 
-packetNumber :: RelayPacket -> Word8
-packetNumber (RelayData _ (ConId conid)) = conid -- 0 to 15 not allowed.
-packetNumber rp                          = fromIntegral $ pred $ constrIndex $ toConstr rp
+newtype PacketNumber = PacketNumber { packetNumberToWord8 :: Word8 }
+ deriving (Eq,Ord,Show)
+
+pattern PingPacket    = PacketNumber 4
+pattern OnionPacketID = PacketNumber 8
+
+packetNumber :: RelayPacket -> PacketNumber
+packetNumber (RelayData _ (ConId conid)) = PacketNumber $ conid -- 0 to 15 not allowed.
+packetNumber rp                          = PacketNumber $ fromIntegral $ pred $ constrIndex $ toConstr rp
 
 instance Sized RelayPacket where
     size = mappend (ConstSize 1) $ VarSize $ \x -> case x of
@@ -101,7 +108,7 @@ instance Serialize RelayPacket where
             conid -> (`RelayData` ConId conid) <$> (remaining >>= getBytes)
 
     put rp = do
-        putWord8 $ packetNumber rp
+        putWord8 $ packetNumberToWord8 $ packetNumber rp
         case rp of
             RoutingRequest k             -> putPublicKey k
             RoutingResponse rpid k       -> put rpid >> putPublicKey k
diff --git a/src/Network/QueryResponse/TCP.hs b/src/Network/QueryResponse/TCP.hs
index 328a19e1..efeab305 100644
--- a/src/Network/QueryResponse/TCP.hs
+++ b/src/Network/QueryResponse/TCP.hs
@@ -19,9 +19,12 @@ import Data.Time.Clock.POSIX
 import Data.Word
 import Network.BSD
 import Network.Socket
+import System.Timeout
 import System.IO
 import System.IO.Error
 
+import DebugTag
+import DPut
 import Connection.Tcp (socketFamily)
 import qualified Data.MinMaxPSQ as MM
 import Network.QueryResponse
@@ -61,11 +64,12 @@ acquireConnection :: MVar (Maybe (Either a (x, addr)))
                            -> TCPCache (SessionProtocol x y)
                            -> StreamHandshake addr x y
                            -> addr
+                           -> Bool
                            -> IO (Maybe (y -> IO ()))
-acquireConnection mvar tcpcache stream addr = do
+acquireConnection mvar tcpcache stream addr bDoCon = do
     cache <- atomically $ readTVar (lru tcpcache)
     case MM.lookup' (TCPAddress $ streamAddr stream addr) cache of
-        Nothing -> do
+        Nothing -> fmap join $ forM (guard bDoCon) $ \() -> do
             proto <- getProtocolNumber "tcp"
             mh <- catchIOError (do sock <- socket (socketFamily $ streamAddr stream addr) Stream proto
                                    connect sock (streamAddr stream addr) `catchIOError` (\e -> close sock)
@@ -77,10 +81,12 @@ acquireConnection mvar tcpcache stream addr = do
             t <- getPOSIXTime
             rthread <- forkIO $ fix $ \loop -> do
                 x <- streamDecode st
-                putMVar mvar $ fmap (\u -> Right (u, addr)) x
                 case x of
-                    Just _ -> loop
+                    Just u -> do
+                        timeout (1000000) $ putMVar mvar $ Just $ Right (u, addr)
+                        loop
                     Nothing -> do
+                        dput XTCP $ "TCP disconnected: " ++ show (streamAddr stream addr)
                         atomically $ modifyTVar' (lru tcpcache)
                                    $ MM.delete (TCPAddress $ streamAddr stream addr)
                         hClose h
@@ -91,14 +97,18 @@ acquireConnection mvar tcpcache stream addr = do
                     , tcpState  = st
                     , tcpThread = rthread
                     }
-            let (retires,cache') = MM.takeView (tcpMax tcpcache)
-                                    $ MM.insert' (TCPAddress $ streamAddr stream addr) v (Down t) cache
+            retires <- atomically $ do
+                c <- readTVar (lru tcpcache)
+                let (rs,c') = MM.takeView (tcpMax tcpcache)
+                            $ MM.insert' (TCPAddress $ streamAddr stream addr) v (Down t) c
+                writeTVar (lru tcpcache) c'
+                return rs
             forM_ retires $ \(MM.Binding (TCPAddress k) r _) -> void $ forkIO $ do
                 myThreadId >>= flip labelThread ("tcp-close:"++show k)
+                dput XTCP $ "TCP dropped: " ++ show k
                 killThread (tcpThread r)
                 streamGoodbye st
                 hClose (tcpHandle r)
-                atomically $ modifyTVar' (lru tcpcache) $ MM.delete (TCPAddress k)
 
             return $ Just $ streamEncode st
         Just (tm,v) -> do
@@ -119,15 +129,15 @@ closeAll tcpcache stream = do
 
 tcpTransport :: Int -- ^ maximum number of TCP links to maintain.
                 -> StreamHandshake addr x y
-                -> IO (TransportA err addr x y)
+                -> IO (TransportA err addr x (Bool,y))
 tcpTransport maxcon stream = do
     msgvar <- newEmptyMVar
     tcpcache <- atomically $ (`TCPCache` maxcon) <$> newTVar (MM.empty)
     return Transport
-        { awaitMessage   = (takeMVar msgvar >>=)
-        , sendMessage    = \addr y -> do
+        { awaitMessage   = \f -> takeMVar msgvar >>= \x -> f x `catchIOError` (\e -> dput XTCP ("TCP transport stopped. " ++ show e) >> f Nothing)
+        , sendMessage    = \addr (bDoCon,y) -> do
                                 t <- forkIO $ do
-                                    msock <- acquireConnection msgvar tcpcache stream addr
+                                    msock <- acquireConnection msgvar tcpcache stream addr bDoCon
                                     mapM_ ($ y) msock
                                   `catchIOError` \e -> return ()
                                 labelThread t "tcp-send"
diff --git a/src/Network/Tox.hs b/src/Network/Tox.hs
index d0304170..c14339e4 100644
--- a/src/Network/Tox.hs
+++ b/src/Network/Tox.hs
@@ -328,7 +328,7 @@ newToxOverTransport keydb addr onNewSession suppliedDHTKey udp tcp = do
     (orouter,otbl) <- newOnionRouter crypto (dput XRoutes)
     (cryptonet,dhtcrypt,onioncrypt,dtacrypt,handshakes)
         <- toxTransport crypto orouter lookupClose udp
-                (sendMessage (clientNet $ tcpClient $ tcpKademliaClient orouter))
+                (\dst x -> sendMessage (clientNet $ tcpClient $ tcpKademliaClient orouter) dst (True,x))
                 tcp
     sessions <- initSessions (sendMessage cryptonet)
 
diff --git a/src/Network/Tox/TCP.hs b/src/Network/Tox/TCP.hs
index a7881c24..71050fe8 100644
--- a/src/Network/Tox/TCP.hs
+++ b/src/Network/Tox/TCP.hs
@@ -7,6 +7,7 @@ module Network.Tox.TCP
     , NodeInfo(..)
     ) where
 
+import Debug.Trace
 import Control.Arrow
 import Control.Concurrent
 import Control.Concurrent.STM
@@ -27,6 +28,7 @@ import qualified Data.Vector as Vector
 import Network.Socket (SockAddr(..))
 import qualified Text.ParserCombinators.ReadP as RP
 import System.IO.Error
+import System.Timeout
 
 import ControlMaybe
 import Crypto.Tox
@@ -113,15 +115,16 @@ tcpStream crypto = StreamHandshake
                                     dput XTCP $ "TCP: Failed to decode packet."
                                     return Nothing
                                 Right x -> do
-                                    n24 <- takeMVar nread
-                                    let r = decrypt (noncef' n24) x >>= decodePlain
-                                    putMVar nread (incrementNonce24 n24)
-                                    either (dput XTCP)
-                                           (\x' -> do
-                                                dput XTCP $ "TCP:" ++ show addr ++ " --> " ++ show x'
-                                                return ())
-                                           r
-                                    return $ either (const Nothing) Just r
+                                    m24 <- timeout 100000 (takeMVar nread)
+                                    fmap join $ forM m24 $ \n24 -> do
+                                        let r = decrypt (noncef' n24) x >>= decodePlain
+                                        putMVar nread (incrementNonce24 n24)
+                                        either (dput XTCP)
+                                               (\x' -> do
+                                                    dput XTCP $ "TCP:" ++ show addr ++ " --> " ++ show x'
+                                                    return ())
+                                               r
+                                        return $ either (const Nothing) Just r
                  in go `catchIOError` \e -> do
                             dput XTCP $ "TCP exception: " ++ show e
                             return Nothing
@@ -129,14 +132,14 @@ tcpStream crypto = StreamHandshake
                 n24 <- takeMVar nsend
                 dput XTCP $ "TCP:" ++ show addr ++ " <-- " ++ show y
                 let bs = encode $ encrypt (noncef' n24) $ encodePlain y
-                hPut h $ encode (fromIntegral $ Data.ByteString.length bs :: Word16)
-                            <> bs
+                hPut h (encode (fromIntegral $ Data.ByteString.length bs :: Word16) <> bs)
+                    `catchIOError` \e -> dput XTCP $ "TCP write exception: " ++ show e
                 putMVar nsend (incrementNonce24 n24)
             }
     , streamAddr = nodeAddr
     }
 
-toxTCP :: TransportCrypto -> IO (Transport err NodeInfo RelayPacket)
+toxTCP :: TransportCrypto -> IO (TransportA err NodeInfo RelayPacket (Bool,RelayPacket))
 toxTCP crypto = tcpTransport 30 (tcpStream crypto)
 
 tcpSpace :: KademliaSpace NodeId NodeInfo
@@ -153,7 +156,7 @@ nodeSearch tcp = Search
 
 data TCPClient err meth tid = TCPClient
     { tcpCrypto     :: TransportCrypto
-    , tcpClient     :: Client err () tid NodeInfo RelayPacket
+    , tcpClient     :: Client err PacketNumber tid NodeInfo (Bool,RelayPacket)
     , tcpGetGateway :: UDP.NodeInfo -> STM (Maybe NodeInfo)
     }
 
@@ -200,8 +203,8 @@ getUDPNodes' tcp seeking dst0 = do
         wrap0 <- lookupNonceFunction (tcpCrypto tcp) (transportSecret $ tcpCrypto tcp) (UDP.id2key $ UDP.nodeId dst)
         let meth = MethodSerializer -- MethodSerializer Nonce8 NodeInfo RelayPacket meth AnnounceRequest (Either String AnnounceResponse)
                     { methodTimeout = \tid addr -> return (addr,12000000) -- 12 second timeout
-                    , method = () -- meth
-                    , wrapQuery = \n8 src gateway x ->
+                    , method = OnionPacketID -- meth
+                    , wrapQuery = \n8 src gateway x -> (,) True $
                         OnionPacket n24 $ Addressed (UDP.nodeAddr dst)
                             $ wrapOnionPure b (wrap2 n24) (nodeAddr gateway')
                                 $ wrapOnionPure c (wrap1 n24) (UDP.nodeAddr dst)
@@ -212,7 +215,7 @@ getUDPNodes' tcp seeking dst0 = do
                                             , asymmData  = pure (x,n8)
                                             }
                     , unwrapResponse = \case
-                        OnionPacketResponse (OnionAnnounceResponse _ n24' r)
+                        (_,OnionPacketResponse (OnionAnnounceResponse _ n24' r))
                             -> decrypt (wrap0 n24') r >>= decodePlain
                         x   -> Left $ "getUDPNodes: unwrapResponse fail " ++ show x
                     }
@@ -222,26 +225,28 @@ getUDPNodes' tcp seeking dst0 = do
             return ( (ns,ns, const () <$> mb), gateway )
 
 
-handleOOB :: PublicKey -> ByteString -> NodeInfo -> NodeInfo -> IO (Maybe (RelayPacket -> RelayPacket))
+handleOOB :: PublicKey -> ByteString -> NodeInfo -> NodeInfo -> IO (Maybe (x -> x))
 handleOOB k bs src dst = do
     dput XMisc $ "TODO: handleOOB " ++ show src
     return Nothing
 
-handle2route :: OnionMessage Encrypted -> NodeInfo -> NodeInfo -> IO (Maybe (RelayPacket -> RelayPacket))
+handle2route :: OnionMessage Encrypted -> NodeInfo -> NodeInfo -> IO (Maybe (x -> x))
 handle2route o src dst = do
     dput XMisc $ "TODO: handle2route " ++ show src
     return Nothing
 
-tcpPing :: Client err () Nonce8 addr RelayPacket -> addr -> IO (Maybe ())
-tcpPing client dst = sendQuery client meth () dst
+tcpPing :: Show addr => Client err PacketNumber Nonce8 addr (Bool,RelayPacket) -> addr -> IO (Maybe ())
+tcpPing client dst = do
+    dput XTCP $ "tcpPing " ++ show dst
+    sendQuery client meth () dst
  where meth = MethodSerializer
-                { wrapQuery      = \n8 src dst () -> RelayPing n8
+                { wrapQuery      = \n8 src dst () -> (True,RelayPing n8)
                 , unwrapResponse = \_ -> ()
                 , methodTimeout  = \n8 dst -> return (dst,5000000)
-                , method         = ()
+                , method         = PingPacket
                 }
 
-type RelayClient = Client String () Nonce8 NodeInfo RelayPacket
+type RelayClient = Client String PacketNumber Nonce8 NodeInfo (Bool,RelayPacket)
 
 -- | Create a new TCP relay client.  Because polymorphic existential record
 -- updates are currently hard with GHC, this function accepts parameters for
@@ -249,29 +254,34 @@ type RelayClient = Client String () Nonce8 NodeInfo RelayPacket
 -- defaults are 'id' and 'tryPutMVar'.  The resulting customized table state
 -- will be returned to the caller along with the new client.
 newClient :: TransportCrypto
-               -> (MVar RelayPacket -> a)       -- ^ store mvar for query
-               -> (a -> RelayPacket -> IO void) -- ^ load mvar for query
+               -> (MVar (Bool,RelayPacket) -> a) -- ^ store mvar for query
+               -> (a -> RelayPacket -> IO void)  -- ^ load mvar for query
                -> IO ( TVar (ChaChaDRG, Data.Word64Map.Word64Map a)
-                     , Client String () Nonce8 NodeInfo RelayPacket)
+                     , Client String PacketNumber Nonce8 NodeInfo (Bool,RelayPacket))
 newClient crypto store load = do
     net <- toxTCP crypto
     drg <- drgNew
     map_var <- atomically $ newTVar (drg, Data.Word64Map.empty)
     return $ (,) map_var Client
-        { clientNet = net
+        { clientNet = {- XXX: Client type forces this pointless layering. -} layerTransport ((Right .) . (,) . (,) False) (,) net
         , clientDispatcher = DispatchMethods
-            { classifyInbound = \case
-                RelayPing n           -> IsQuery () n
+            { classifyInbound = (. snd) $ \case
+                RelayPing n           -> IsQuery PingPacket n
                 RelayPong n           -> IsResponse n
                 OnionPacketResponse (OnionAnnounceResponse n8 n24 ciphered) -> IsResponse n8
                 OnionPacketResponse o@(OnionToRouteResponse _) -> IsUnsolicited $ handle2route o
                 OOBRecv k bs          -> IsUnsolicited $ handleOOB k bs
-            , lookupHandler   = \() -> Just MethodHandler
-                    { methodParse     = \(RelayPing n8) -> Right ()
-                    , methodSerialize = \n8 src dst () -> RelayPong n8
+            , lookupHandler   = \case
+                PingPacket -> Just MethodHandler
+                    { methodParse     = \(_,RelayPing n8) -> Right ()
+                    , methodSerialize = \n8 src dst () -> (False, RelayPong n8)
                     , methodAction    = \src () -> return ()
                     }
-            , tableMethods = transactionMethods' store load (contramap (\(Nonce8 w64) -> w64) w64MapMethods)
+                w -> trace ("tcp-lookupHandler: "++show w) $ Just NoReply
+                            { methodParse = \x -> Left "tcp-lookuphandler?" -- :: x -> Either err a
+                            , noreplyAction = \addr a -> dput XTCP $ "tcp-lookupHandler: "++show w
+                            }
+            , tableMethods = transactionMethods' store (\x -> load x . snd) (contramap (\(Nonce8 w64) -> w64) w64MapMethods)
                                 $ first (either error Nonce8 . decode) . randomBytesGenerate 8
             }
         , clientErrorReporter = logErrors { reportTimeout = reportTimeout ignoreErrors }