4 files changed, 124 insertions, 74 deletions

diff --git a/OnionRouter.hs b/OnionRouter.hs
index 018a6314..0bbcd20d 100644
--- a/OnionRouter.hs
+++ b/OnionRouter.hs
@@ -79,7 +79,7 @@ data OnionRouter = OnionRouter
       -- discarded and replaced with a fresh one.
     , pendingRoutes   :: TArray Int Int
       -- | Parameters used to implement Kademlia for TCP relays.
-    , tcpKademliaClient :: TCP.TCPClient String () Nonce8
+    , tcpKademliaClient :: TCP.TCPClient String Nonce8
       -- | This thread maintains the TCP relay table.
     , tcpKademliaThread :: ThreadId
     , tcpProber :: TCP.TCPProber
@@ -275,7 +275,7 @@ updateTCP or addr x =
 selectGateway :: TVar (R.BucketList TCP.NodeInfo) -> NodeInfo -> STM (Maybe TCP.NodeInfo)
 selectGateway tbl ni = do
     ns <- kclosest TCP.tcpSpace 2 (nodeId ni) <$> readTVar tbl
-    return $ listToMaybe $ dropWhile (\n -> TCP.nodeId n == nodeId ni) ns
+    return $ listToMaybe ns -- $ dropWhile (\n -> TCP.nodeId n == nodeId ni) ns
 
 quitRouteBuilder :: OnionRouter -> IO ()
 quitRouteBuilder or = do
@@ -473,7 +473,7 @@ handleEvent getnodes or e@(BuildRoute (RouteId rid)) = do
                         forkIO $ sendq csel cq 2 >>= atomically . writeTVar cv . Just
                         -- This timeout should be unnecessary... But I'm paranoid.
                         -- Note: 10 seconds should be sufficient for typical get-nodes queries.
-                        tm <- timeout 20000000 $ atomically $ do -- Wait for all 3 results.
+                        tm <- timeout 30000000 $ atomically $ do -- Wait for all 3 results.
                             rs <- catMaybes <$> sequence [readTVar av,readTVar bv,readTVar cv]
                             case rs of [_,_,_] -> do
                                                     return $ catMaybes $ rs
diff --git a/TCPProber.hs b/TCPProber.hs
index db6d0217..506ef142 100644
--- a/TCPProber.hs
+++ b/TCPProber.hs
@@ -145,7 +145,8 @@ runProbeQueue prober client maxjobs = do
             labelThread t ("probe."++show ni)
         loop
 
-getNodes :: TCPProber -> TCP.TCPClient err () Nonce8 -> NodeId -> TCP.NodeInfo -> IO (Maybe ([TCP.NodeInfo],[TCP.NodeInfo],Maybe ()))
+
+getNodes :: TCPProber -> TCP.TCPClient err Nonce8 -> NodeId -> TCP.NodeInfo -> IO (Maybe ([TCP.NodeInfo],[TCP.NodeInfo],Maybe ()))
 getNodes prober tcp seeking dst = do
     r <- TCP.getUDPNodes' tcp seeking (TCP.udpNodeInfo dst)
     dput XTCP $ "Got via TCP nodes: " ++ show r
@@ -162,11 +163,11 @@ getNodes prober tcp seeking dst = do
                 return $ Just ts
         return $ Just ts
 
-nodeSearch :: TCPProber -> TCP.TCPClient err () Nonce8 -> Search NodeId (IP, PortNumber) () TCP.NodeInfo TCP.NodeInfo
+nodeSearch :: TCPProber -> TCP.TCPClient err Nonce8 -> Search NodeId (IP, PortNumber) () TCP.NodeInfo TCP.NodeInfo
 nodeSearch prober tcp = Search
     { searchSpace       = TCP.tcpSpace
     , searchNodeAddress = TCP.nodeIP &&& TCP.tcpPort
     , searchQuery       = getNodes prober tcp
-    , searchAlpha       = 8
+    , searchAlpha       = 1
     , searchK           = 16
     }
diff --git a/src/Network/QueryResponse/TCP.hs b/src/Network/QueryResponse/TCP.hs
index efeab305..a606c51d 100644
--- a/src/Network/QueryResponse/TCP.hs
+++ b/src/Network/QueryResponse/TCP.hs
@@ -14,6 +14,7 @@ import Control.Monad
 import Data.ByteString (ByteString,hPut)
 import Data.Function
 import Data.Hashable
+import Data.Maybe
 import Data.Ord
 import Data.Time.Clock.POSIX
 import Data.Word
@@ -29,11 +30,13 @@ import Connection.Tcp (socketFamily)
 import qualified Data.MinMaxPSQ as MM
 import Network.QueryResponse
 
-data TCPSession st = TCPSession
-    { tcpHandle :: Handle
-    , tcpState  :: st
-    , tcpThread :: ThreadId
-    }
+data TCPSession st
+    = PendingTCPSession
+    | TCPSession
+        { tcpHandle :: Handle
+        , tcpState  :: st
+        , tcpThread :: ThreadId
+        }
 
 newtype TCPAddress = TCPAddress SockAddr
  deriving (Eq,Ord)
@@ -60,6 +63,10 @@ data StreamHandshake addr x y = StreamHandshake
     , streamAddr  :: addr -> SockAddr
     }
 
+killSession :: TCPSession st -> IO ()
+killSession PendingTCPSession       = return ()
+killSession TCPSession{tcpThread=t} = killThread t
+
 acquireConnection :: MVar (Maybe (Either a (x, addr)))
                            -> TCPCache (SessionProtocol x y)
                            -> StreamHandshake addr x y
@@ -67,65 +74,95 @@ acquireConnection :: MVar (Maybe (Either a (x, addr)))
                            -> Bool
                            -> IO (Maybe (y -> IO ()))
 acquireConnection mvar tcpcache stream addr bDoCon = do
-    cache <- atomically $ readTVar (lru tcpcache)
-    case MM.lookup' (TCPAddress $ streamAddr stream addr) cache of
+    now <- getPOSIXTime
+    entry <- atomically $ do
+        c <- readTVar (lru tcpcache)
+        let v = MM.lookup' (TCPAddress $ streamAddr stream addr) c
+        case v of
+            Nothing | bDoCon -> writeTVar (lru tcpcache)
+                                $ MM.insert' (TCPAddress $ streamAddr stream addr) PendingTCPSession (Down now) c
+                    | otherwise -> return ()
+            Just (tm, v) -> modifyTVar' (lru tcpcache) $ MM.insert' (TCPAddress $ streamAddr stream addr) v (Down now)
+        return v
+    case entry of
         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)
-                                   h <- socketToHandle sock ReadWriteMode
-                                   return $ Just h)
+            mh <- catchIOError (do h <- timeout 10000000 $ do
+                                       sock <- socket (socketFamily $ streamAddr stream addr) Stream proto
+                                       connect sock (streamAddr stream addr) `catchIOError` (\e -> close sock)
+                                       h <- socketToHandle sock ReadWriteMode
+                                       hSetBuffering h NoBuffering
+                                       return h
+                                   return h)
                                $ \e -> return Nothing
-            fmap join $ forM mh $ \h -> do
-            st <- streamHello stream addr h
-            t <- getPOSIXTime
-            rthread <- forkIO $ fix $ \loop -> do
-                x <- streamDecode st
-                case x of
-                    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
-            let showAddr a = show (streamAddr stream a)
-            labelThread rthread ("tcp:"++showAddr addr)
-            let v = TCPSession
-                    { tcpHandle = h
-                    , tcpState  = st
-                    , tcpThread = rthread
-                    }
-            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)
+            ret <- fmap join $ forM mh $ \h -> do
+                st <- streamHello stream addr h
+                dput XTCP $ "TCP Connected! " ++ show (streamAddr stream addr)
+                rthread <- forkIO $ fix $ \loop -> do
+                    x <- streamDecode st
+                    dput XTCP $ "TCP streamDecode " ++ show (streamAddr stream addr) ++ " --> " ++ maybe "Nothing" (const "got") x
+                    case x of
+                        Just u -> do
+                            m <- timeout (1000000) $ putMVar mvar $ Just $ Right (u, addr)
+                            when (isNothing m) $ do
+                                dput XTCP $ "TCP "++show (streamAddr stream addr) ++ "dropped packet."
+                            loop
+                        Nothing -> do
+                            dput XTCP $ "TCP disconnected: " ++ show (streamAddr stream addr)
+                            atomically $ modifyTVar' (lru tcpcache)
+                                       $ MM.delete (TCPAddress $ streamAddr stream addr)
+                            c <- atomically $ readTVar (lru tcpcache)
+                            now <- getPOSIXTime
+                            forM_ (zip [1..] $ MM.toList c) $ \(i,MM.Binding (TCPAddress addr) r (Down tm)) -> do
+                                let stat = case r of
+                                        PendingTCPSession -> "pending."
+                                        TCPSession {}     -> "established."
+                                dput XTCP $ unwords [show i ++ ".", "Still connected:",  show addr, show (now - tm), stat]
+                            hClose h
+                let showAddr a = show (streamAddr stream a)
+                labelThread rthread ("tcp:"++showAddr addr)
+                let v = TCPSession
+                        { tcpHandle = h
+                        , tcpState  = st
+                        , tcpThread = rthread
+                        }
+                t <- getPOSIXTime
+                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
+                    killSession r
+                    case r of TCPSession {tcpState=st,tcpHandle=h} -> do
+                                    streamGoodbye st
+                                    hClose h
+                              _ -> return ()
 
-            return $ Just $ streamEncode st
-        Just (tm,v) -> do
-            t <- getPOSIXTime
-            let TCPSession { tcpHandle = h, tcpState = st } = v
-                cache' = MM.insert' (TCPAddress $ streamAddr stream addr) v (Down t) cache
-            atomically $ writeTVar (lru tcpcache) cache'
-            return $ Just $ streamEncode st
+                return $ Just $ streamEncode st
+            when (isNothing ret) $ do
+                atomically $ modifyTVar' (lru tcpcache) $ MM.delete (TCPAddress $ streamAddr stream addr)
+            return ret
+        Just (tm, PendingTCPSession) -> do
+            fmap join $ timeout 10000000 $ atomically $ do
+                c <- readTVar (lru tcpcache)
+                let v = MM.lookup' (TCPAddress $ streamAddr stream addr) c
+                case v of
+                    Just (_,TCPSession{tcpState=st}) -> return $ Just $ streamEncode st
+                    Nothing                          -> return Nothing
+                    _                                -> retry
+        Just (tm, v@TCPSession {tcpState=st}) -> return $ Just $ streamEncode st
 
 closeAll :: TCPCache (SessionProtocol x y) -> StreamHandshake addr x y -> IO ()
 closeAll tcpcache stream = do
     cache <- atomically $ swapTVar (lru tcpcache) MM.empty
     forM_ (MM.toList cache) $ \(MM.Binding (TCPAddress addr) r tm) -> do
-        let st = tcpState r
-        killThread (tcpThread r)
-        streamGoodbye st
-        hClose (tcpHandle r)
+        killSession r
+        case r of TCPSession{tcpState=st,tcpHandle=h} -> streamGoodbye st >> hClose h
+                  _                                   -> return ()
 
 tcpTransport :: Int -- ^ maximum number of TCP links to maintain.
                 -> StreamHandshake addr x y
diff --git a/src/Network/Tox/TCP.hs b/src/Network/Tox/TCP.hs
index 71050fe8..36200586 100644
--- a/src/Network/Tox/TCP.hs
+++ b/src/Network/Tox/TCP.hs
@@ -11,6 +11,7 @@ import Debug.Trace
 import Control.Arrow
 import Control.Concurrent
 import Control.Concurrent.STM
+import Control.Exception
 import Control.Monad
 import Crypto.Random
 import Data.Aeson (ToJSON(..),FromJSON(..))
@@ -92,7 +93,7 @@ tcpStream crypto = StreamHandshake
             nil = SessionProtocol
                 { streamGoodbye = return ()
                 , streamDecode = return Nothing
-                , streamEncode = \y -> return ()
+                , streamEncode = \y -> dput XTCP $ "TCP nil <-- " ++ show y
                 }
         either (\_ -> return nil) id $ mwelcome <&> \welcome -> do
         dput XTCP $ "TCP:" ++ show addr ++ " --> " ++ show welcome
@@ -100,12 +101,13 @@ tcpStream crypto = StreamHandshake
         nsend <- newMVar (sessionBaseNonce $ runIdentity $ helloData hello)
         nread <- newMVar (sessionBaseNonce $ runIdentity $ welcomeData welcome)
         let them = sessionPublicKey $ runIdentity $ welcomeData welcome
+        hvar <- newMVar h
         return SessionProtocol
             { streamGoodbye = do
                 dput XTCP $ "Closing " ++ show addr
                 return () -- No goodbye packet?  Seems rude.
             , streamDecode =
-                let go = decode <$> hGet h 2 >>= \case
+                let go h = decode <$> hGet h 2 >>= \case
                         Left e    -> do
                             dput XTCP $ "TCP: Failed to get length: " ++ e
                             return Nothing
@@ -115,25 +117,27 @@ tcpStream crypto = StreamHandshake
                                     dput XTCP $ "TCP: Failed to decode packet."
                                     return Nothing
                                 Right x -> do
-                                    m24 <- timeout 100000 (takeMVar nread)
+                                    m24 <- timeout 1000000 (takeMVar nread)
                                     fmap join $ forM m24 $ \n24 -> do
                                         let r = decrypt (noncef' n24) x >>= decodePlain
                                         putMVar nread (incrementNonce24 n24)
-                                        either (dput XTCP)
+                                        either (dput XTCP . ("TCP decryption: " ++))
                                                (\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
+                 in bracket (takeMVar hvar) (putMVar hvar)
+                     $ \h -> go h `catchIOError` \e -> do
+                                        dput XTCP $ "TCP exception: " ++ show e
+                                        return Nothing
             , streamEncode = \y -> do
                 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)
-                    `catchIOError` \e -> dput XTCP $ "TCP write exception: " ++ show e
+                ($ h) -- bracket (takeMVar hvar) (putMVar hvar)
+                    $ \h -> 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
@@ -154,7 +158,7 @@ nodeSearch tcp = Search
     }
 -}
 
-data TCPClient err meth tid = TCPClient
+data TCPClient err tid = TCPClient
     { tcpCrypto     :: TransportCrypto
     , tcpClient     :: Client err PacketNumber tid NodeInfo (Bool,RelayPacket)
     , tcpGetGateway :: UDP.NodeInfo -> STM (Maybe NodeInfo)
@@ -181,10 +185,10 @@ getTCPNodes tcp seeking dst = do
         return $ Just ts
 -}
 
-getUDPNodes :: TCPClient err () Nonce8 -> NodeId -> UDP.NodeInfo -> IO (Maybe ([UDP.NodeInfo], [UDP.NodeInfo], Maybe ()))
+getUDPNodes :: TCPClient err Nonce8 -> NodeId -> UDP.NodeInfo -> IO (Maybe ([UDP.NodeInfo], [UDP.NodeInfo], Maybe ()))
 getUDPNodes tcp seeking dst = fmap fst <$> getUDPNodes' tcp seeking dst
 
-getUDPNodes' :: TCPClient err () Nonce8 -> NodeId -> UDP.NodeInfo -> IO (Maybe (([UDP.NodeInfo], [UDP.NodeInfo], Maybe ()), NodeInfo))
+getUDPNodes' :: TCPClient err Nonce8 -> NodeId -> UDP.NodeInfo -> IO (Maybe (([UDP.NodeInfo], [UDP.NodeInfo], Maybe ()), NodeInfo))
 getUDPNodes' tcp seeking dst0 = do
     mgateway <- atomically $ tcpGetGateway tcp dst0
     fmap join $ forM mgateway $ \gateway -> do
@@ -201,7 +205,14 @@ getUDPNodes' tcp seeking dst0 = do
         wrap2 <- lookupNonceFunction (tcpCrypto tcp) b (UDP.id2key $ UDP.nodeId dst)
         wrap1 <- lookupNonceFunction (tcpCrypto tcp) c (UDP.id2key $ nodeId gateway)
         wrap0 <- lookupNonceFunction (tcpCrypto tcp) (transportSecret $ tcpCrypto tcp) (UDP.id2key $ UDP.nodeId dst)
-        let meth = MethodSerializer -- MethodSerializer Nonce8 NodeInfo RelayPacket meth AnnounceRequest (Either String AnnounceResponse)
+        let meth :: MethodSerializer
+                        Nonce8
+                        a -- NodeInfo
+                        (Bool, RelayPacket)
+                        PacketNumber
+                        AnnounceRequest
+                        (Either String AnnounceResponse)
+            meth = MethodSerializer
                     { methodTimeout = \tid addr -> return (addr,12000000) -- 12 second timeout
                     , method = OnionPacketID -- meth
                     , wrapQuery = \n8 src gateway x -> (,) True $
@@ -271,6 +282,7 @@ newClient crypto store load = do
                 OnionPacketResponse (OnionAnnounceResponse n8 n24 ciphered) -> IsResponse n8
                 OnionPacketResponse o@(OnionToRouteResponse _) -> IsUnsolicited $ handle2route o
                 OOBRecv k bs          -> IsUnsolicited $ handleOOB k bs
+                wut                   -> IsUnknown (show wut)
             , lookupHandler   = \case
                 PingPacket -> Just MethodHandler
                     { methodParse     = \(_,RelayPing n8) -> Right ()