From a60988db82aed029b71531e8a4bc5698fa247c02 Mon Sep 17 00:00:00 2001
From: Bin Jin <bjin@ctrl-d.org>
Date: Tue, 14 Mar 2017 21:18:00 +0800
Subject: Tun: add timestamp for incoming packets

---
 src/Network/WireGuard/Internal/State.hs | 18 ++++++++++++------
 src/Network/WireGuard/Internal/Util.hs  | 10 ++++++++++
 2 files changed, 22 insertions(+), 6 deletions(-)

(limited to 'src/Network/WireGuard/Internal')

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
diff --git a/src/Network/WireGuard/Internal/Util.hs b/src/Network/WireGuard/Internal/Util.hs
index f7ecde5..6aefee7 100644
--- a/src/Network/WireGuard/Internal/Util.hs
+++ b/src/Network/WireGuard/Internal/Util.hs
@@ -7,6 +7,7 @@ module Network.WireGuard.Internal.Util
   , catchIOExceptionAnd
   , catchSomeExceptionAnd
   , withJust
+  , dropUntilM
   , zeroMemory
   , copyMemory
   ) where
@@ -52,6 +53,15 @@ withJust mma func = do
         Nothing -> return ()
         Just a  -> func a
 
+dropUntilM :: Monad m => (a -> Bool) -> m a -> m a
+dropUntilM cond ma = loop
+  where
+    loop = do
+        a <- ma
+        if cond a
+          then return a
+          else loop
+
 zeroMemory :: Ptr a -> CSize -> IO ()
 zeroMemory dest nbytes = memset dest 0 (fromIntegral nbytes)
 
-- 
cgit v1.2.3-59-g8ed1b