From a60988db82aed029b71531e8a4bc5698fa247c02 Mon Sep 17 00:00:00 2001 From: Bin Jin Date: Tue, 14 Mar 2017 21:18:00 +0800 Subject: Tun: add timestamp for incoming packets --- src/Network/WireGuard/Internal/State.hs | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) (limited to 'src/Network/WireGuard/Internal/State.hs') diff --git a/src/Network/WireGuard/Internal/State.hs b/src/Network/WireGuard/Internal/State.hs index 38866e8..f7b1ca0 100644 --- a/src/Network/WireGuard/Internal/State.hs +++ b/src/Network/WireGuard/Internal/State.hs @@ -193,12 +193,18 @@ getSession peer = do [] -> return Nothing (s:_) -> return (Just s) -waitForSession :: Peer -> STM Session -waitForSession peer = do - sessions' <- readTVar (sessions peer) - case sessions' of - [] -> retry - (s:_) -> return s +waitForSession :: Int -> Peer -> IO (Maybe Session) +waitForSession timelimit peer = do + getTimeout <- registerDelay timelimit + atomically $ do + sessions' <- readTVar (sessions peer) + case sessions' of + [] -> do + timeout <- readTVar getTimeout + if timeout + then return Nothing + else retry + (s:_) -> return (Just s) findSession :: Peer -> Index -> STM (Maybe (Either ResponderWait Session)) findSession peer index = do -- cgit v1.2.3-59-g8ed1b