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

@@ -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 ())