From 5443b38525d34f3c9e7c58555e343105e3c65650 Mon Sep 17 00:00:00 2001
From: John Lenz <lenz@math.uic.edu>
Date: Fri, 28 Aug 2015 16:34:09 -0500
Subject: [PATCH] yesod-websockets: accept websocket requests with a
 subprotocol

The websockets library allows passing an AcceptRequest when
upgrading the request to websockets, which at the moment only
contains the subprotocol to use with the client.  The list of
subprotocols that the client is prepared to speak is in the
RequestHead.
---
 yesod-websockets/Yesod/WebSockets.hs | 46 ++++++++++++++++++++--------
 1 file changed, 34 insertions(+), 12 deletions(-)

diff --git a/yesod-websockets/Yesod/WebSockets.hs b/yesod-websockets/Yesod/WebSockets.hs
index d1c93acb..d064022f 100644
--- a/yesod-websockets/Yesod/WebSockets.hs
+++ b/yesod-websockets/Yesod/WebSockets.hs
@@ -5,6 +5,7 @@ module Yesod.WebSockets
     ( -- * Core API
       WebSocketsT
     , webSockets
+    , webSocketsWith
     , receiveData
     , receiveDataE
     , receiveDataMessageE
@@ -55,19 +56,40 @@ type WebSocketsT = ReaderT WS.Connection
 --
 -- Since 0.1.0
 webSockets :: (Y.MonadBaseControl IO m, Y.MonadHandler m) => WebSocketsT m () -> m ()
-webSockets inner = do
+webSockets = webSocketsWith $ const $ return $ Just $ WS.AcceptRequest Nothing
+
+-- | Varient of 'webSockets' which allows you to specify the 'WS.AcceptRequest'
+-- setttings when upgrading to a websocket connection.
+--
+-- Since 0.2.4
+webSocketsWith :: (Y.MonadBaseControl IO m, Y.MonadHandler m)
+               => (WS.RequestHead -> m (Maybe WS.AcceptRequest))
+               -- ^ A Nothing indicates that the websocket upgrade request should not happen
+               -- and instead the rest of the handler will be called instead.  This allows
+               -- you to use 'WS.getRequestSubprotocols' and only accept the request if
+               -- a compatible subprotocol is given.  Also, the action runs before upgrading
+               -- the request to websockets, so you can also use short-circuiting handler
+               -- actions such as 'Y.invalidArgs'.
+               -> WebSocketsT m ()
+               -> m ()
+webSocketsWith buildAr inner = do
     req <- Y.waiRequest
-    when (WaiWS.isWebSocketsReq req) $
-        Y.sendRawResponseNoConduit
-          $ \src sink -> control $ \runInIO -> WaiWS.runWebSockets
-            WS.defaultConnectionOptions
-            (WaiWS.getRequestHead req)
-            (\pconn -> do
-                conn <- WS.acceptRequest pconn
-                WS.forkPingThread conn 30
-                runInIO $ runReaderT inner conn)
-            src
-            sink
+    when (WaiWS.isWebSocketsReq req) $ do
+        let rhead = WaiWS.getRequestHead req
+        mar <- buildAr rhead
+        case mar of
+            Nothing -> return ()
+            Just ar -> do
+                Y.sendRawResponseNoConduit
+                  $ \src sink -> control $ \runInIO -> WaiWS.runWebSockets
+                    WS.defaultConnectionOptions
+                    rhead
+                    (\pconn -> do
+                        conn <- WS.acceptRequestWith pconn ar
+                        WS.forkPingThread conn 30
+                        runInIO $ runReaderT inner conn)
+                    src
+                    sink
 
 -- | Wrapper for capturing exceptions
 wrapWSE :: (MonadIO m, WS.WebSocketsData a) => (WS.Connection -> a -> IO ())-> a -> WebSocketsT m (Either SomeException ())