1 files changed, 43 insertions, 33 deletions

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 }