Allow sending extra headers when accepting request

src/Network/WebSockets/Connection.hs

@@ -73,6 +73,8 @@ data AcceptRequest = AcceptRequest
     -- ^ The subprotocol to speak with the client.  If 'pendingSubprotcols' is
     -- non-empty, 'acceptSubprotocol' must be one of the subprotocols from the
     -- list.
+    , acceptHeaders :: !Headers
+    -- ^ Extra headers to send with the response.
     }
 
 
@@ -85,7 +87,7 @@ sendResponse pc rsp = Stream.write (pendingStream pc)
 
 --------------------------------------------------------------------------------
 acceptRequest :: PendingConnection -> IO Connection
-acceptRequest pc = acceptRequestWith pc $ AcceptRequest Nothing
+acceptRequest pc = acceptRequestWith pc $ AcceptRequest Nothing []
 
 
 --------------------------------------------------------------------------------
@@ -96,7 +98,8 @@ acceptRequestWith pc ar = case find (flip compatible request) protocols of
         throwIO NotSupported
     Just protocol -> do
         let subproto = maybe [] (\p -> [("Sec-WebSocket-Protocol", p)]) $ acceptSubprotocol ar
-            response = finishRequest protocol request subproto
+            headers = subproto ++ acceptHeaders ar
+            response = finishRequest protocol request headers
         sendResponse pc response
         parse <- decodeMessages protocol (pendingStream pc)
         write <- encodeMessages protocol ServerConnection (pendingStream pc)

tests/haskell/Network/WebSockets/Handshake/Tests.hs

@@ -29,6 +29,8 @@ tests :: Test
 tests = testGroup "Network.WebSockets.Handshake.Test"
     [ testCase "handshake Hybi13"                   testHandshakeHybi13
     , testCase "handshake Hybi13 with subprotocols" testHandshakeHybi13WithProto
+    , testCase "handshake Hybi13 with headers"      testHandshakeHybi13WithHeaders
+    , testCase "handshake Hybi13 with subprotocols and headers" testHandshakeHybi13WithProtoAndHeaders
     , testCase "handshake reject"                   testHandshakeReject
     , testCase "handshake Hybi9000"                 testHandshakeHybi9000
     ]
@@ -90,7 +92,7 @@ testHandshakeHybi13WithProto = do
     ResponseHead code message headers <- testHandshake rq13 $ \pc -> do
         getRequestSubprotocols (pendingRequest pc) @?= ["chat", "superchat"]
         acceptRequestWith pc {pendingOnAccept = \_ -> writeIORef onAcceptFired True}
-                          (AcceptRequest $ Just "superchat")
+                          (AcceptRequest (Just "superchat") [])
 
     readIORef onAcceptFired >>= assert
     code @?= 101
@@ -99,6 +101,40 @@ testHandshakeHybi13WithProto = do
     headers ! "Connection"           @?= "Upgrade"
     headers ! "Sec-WebSocket-Protocol" @?= "superchat"
 
+--------------------------------------------------------------------------------
+testHandshakeHybi13WithHeaders :: Assertion
+testHandshakeHybi13WithHeaders = do
+    onAcceptFired                     <- newIORef False
+    ResponseHead code message headers <- testHandshake rq13 $ \pc -> do
+        getRequestSubprotocols (pendingRequest pc) @?= ["chat", "superchat"]
+        acceptRequestWith pc {pendingOnAccept = \_ -> writeIORef onAcceptFired True}
+                          (AcceptRequest Nothing [("Set-Cookie","sid=foo")])
+
+    readIORef onAcceptFired >>= assert
+    code @?= 101
+    message @?= "WebSocket Protocol Handshake"
+    headers ! "Sec-WebSocket-Accept" @?= "HSmrc0sMlYUkAGmm5OPpG2HaGWk="
+    headers ! "Connection"           @?= "Upgrade"
+    headers ! "Set-Cookie"           @?= "sid=foo"
+    lookup "Sec-WebSocket-Protocol" headers @?= Nothing
+
+--------------------------------------------------------------------------------
+testHandshakeHybi13WithProtoAndHeaders :: Assertion
+testHandshakeHybi13WithProtoAndHeaders = do
+    onAcceptFired                     <- newIORef False
+    ResponseHead code message headers <- testHandshake rq13 $ \pc -> do
+        getRequestSubprotocols (pendingRequest pc) @?= ["chat", "superchat"]
+        acceptRequestWith pc {pendingOnAccept = \_ -> writeIORef onAcceptFired True}
+                          (AcceptRequest (Just "superchat") [("Set-Cookie","sid=foo")])
+
+    readIORef onAcceptFired >>= assert
+    code @?= 101
+    message @?= "WebSocket Protocol Handshake"
+    headers ! "Sec-WebSocket-Accept" @?= "HSmrc0sMlYUkAGmm5OPpG2HaGWk="
+    headers ! "Connection"           @?= "Upgrade"
+    headers ! "Sec-WebSocket-Protocol" @?= "superchat"
+    headers ! "Set-Cookie"           @?= "sid=foo"
+
 --------------------------------------------------------------------------------
 testHandshakeReject :: Assertion
 testHandshakeReject = do