2 files changed, 22 insertions, 6 deletions
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)