1 files changed, 20 insertions, 9 deletions
diff --git a/Network/Wai/Middleware/Braid.hs b/Network/Wai/Middleware/Braid.hs
index 5dbc7f4..c14e099 100644
--- a/Network/Wai/Middleware/Braid.hs
+++ b/Network/Wai/Middleware/Braid.hs
@@ -57,11 +57,13 @@ import Data.ByteString.Builder (Builder, byteString)
 import qualified Data.ByteString.Char8 as BC
 import qualified Data.ByteString.Lazy as L
 import qualified Data.CaseInsensitive as CI
+import Data.Char (isDigit)
 import Network.HTTP.Types.Header (Header, HeaderName, RequestHeaders)
 import Network.HTTP.Types.Method (methodGet, methodPatch, methodPut)
 import Network.HTTP.Types.Status (Status, mkStatus)
 import qualified Network.Wai as Wai
 import Network.Wai.Middleware.AddHeaders (addHeaders)
+import System.Timeout (timeout)
 
 type Topic = [Text]
 
@@ -101,10 +103,13 @@ hSub = "Subscribe"
 getSubscription :: Wai.Request -> Maybe B.ByteString
 getSubscription req = lookupHeader hSub <| Wai.requestHeaders req
 
-getSubscriptionKeepAliveTime :: Wai.Request -> B.ByteString
+getSubscriptionKeepAliveTime :: Wai.Request -> Maybe Int
 getSubscriptionKeepAliveTime req =
-  let Just s = lookupHeader hSub <| Wai.requestHeaders req
-   in snd <| BC.breakSubstring "=" s
+  lookupHeader hSub (Wai.requestHeaders req) +> \h ->
+    let (_, rest) = BC.breakSubstring "keep-alive=" h
+     in if B.null rest
+          then Nothing
+          else readMaybe <| BC.unpack <| BC.takeWhile isDigit <| B.drop 11 rest
 
 hasSubscription :: Wai.Request -> Bool
 hasSubscription req = isJust <| getSubscription req
@@ -222,15 +227,21 @@ addPatchHeader = Wai.ifRequest isPutRequest <| addHeaders [("Patches", "OK")]
 
 -- |
 --    TODO: look into Chan vs BroadcastChan (https://github.com/merijn/broadcast-chan)
-streamUpdates :: Chan Update -> Topic -> Maybe ByteString -> Wai.StreamingBody
-streamUpdates chan topic client write flush = do
+streamUpdates :: Chan Update -> Topic -> Maybe ByteString -> Maybe Int -> Wai.StreamingBody
+streamUpdates chan topic client keepAlive write flush = do
   flush
   src <- liftIO <| dupChan chan
   fix <| \loop -> do
-    update <- readChan src
-    case updateToBuilder topic client update of
-      Just b -> write b >> flush >> loop
-      Nothing -> loop
+    update <-
+      case keepAlive of
+        Just t -> timeout (t * 1000000) (readChan src)
+        Nothing -> Just </ readChan src
+    case update of
+      Just u ->
+        case updateToBuilder topic client u of
+          Just b -> write b >> flush >> loop
+          Nothing -> loop
+      Nothing -> write (byteString ": \n") >> flush >> loop
 
 braidify :: Chan Update -> Wai.Middleware
 braidify src =