1--------------------------------------------------------------------------------
2-- | This module exposes connection internals and should only be used if you
3-- really know what you are doing.
4{-# LANGUAGE OverloadedStrings #-}
5module Network.WebSockets.Connection
6    ( PendingConnection (..)
7    , acceptRequest
8    , AcceptRequest(..)
9    , defaultAcceptRequest
10    , acceptRequestWith
11    , rejectRequest
12    , RejectRequest(..)
13    , defaultRejectRequest
14    , rejectRequestWith
15
16    , Connection (..)
17
18    , ConnectionOptions (..)
19    , defaultConnectionOptions
20
21    , receive
22    , receiveDataMessage
23    , receiveData
24    , send
25    , sendDataMessage
26    , sendDataMessages
27    , sendTextData
28    , sendTextDatas
29    , sendBinaryData
30    , sendBinaryDatas
31    , sendClose
32    , sendCloseCode
33    , sendPing
34
35    , withPingThread
36    , forkPingThread
37    , pingThread
38
39    , CompressionOptions (..)
40    , PermessageDeflate (..)
41    , defaultPermessageDeflate
42
43    , SizeLimit (..)
44    ) where
45
46
47--------------------------------------------------------------------------------
48import           Control.Applicative                             ((<$>))
49import           Control.Concurrent                              (forkIO,
50                                                                  threadDelay)
51import qualified Control.Concurrent.Async                        as Async
52import           Control.Exception                               (AsyncException,
53                                                                  fromException,
54                                                                  handle,
55                                                                  throwIO)
56import           Control.Monad                                   (foldM, unless,
57                                                                  when)
58import qualified Data.ByteString                                 as B
59import qualified Data.ByteString.Builder                         as Builder
60import qualified Data.ByteString.Char8                           as B8
61import           Data.IORef                                      (IORef,
62                                                                  newIORef,
63                                                                  readIORef,
64                                                                  writeIORef)
65import           Data.List                                       (find)
66import           Data.Maybe                                      (catMaybes)
67import qualified Data.Text                                       as T
68import           Data.Word                                       (Word16)
69import           Prelude
70
71
72--------------------------------------------------------------------------------
73import           Network.WebSockets.Connection.Options
74import           Network.WebSockets.Extensions                   as Extensions
75import           Network.WebSockets.Extensions.PermessageDeflate
76import           Network.WebSockets.Extensions.StrictUnicode
77import           Network.WebSockets.Http
78import           Network.WebSockets.Protocol
79import           Network.WebSockets.Stream                       (Stream)
80import qualified Network.WebSockets.Stream                       as Stream
81import           Network.WebSockets.Types
82
83
84--------------------------------------------------------------------------------
85-- | A new client connected to the server. We haven't accepted the connection
86-- yet, though.
87data PendingConnection = PendingConnection
88    { pendingOptions  :: !ConnectionOptions
89    -- ^ Options, passed as-is to the 'Connection'
90    , pendingRequest  :: !RequestHead
91    -- ^ Useful for e.g. inspecting the request path.
92    , pendingOnAccept :: !(Connection -> IO ())
93    -- ^ One-shot callback fired when a connection is accepted, i.e., *after*
94    -- the accepting response is sent to the client.
95    , pendingStream   :: !Stream
96    -- ^ Input/output stream
97    }
98
99
100--------------------------------------------------------------------------------
101-- | This datatype allows you to set options for 'acceptRequestWith'.  It is
102-- strongly recommended to use 'defaultAcceptRequest' and then modify the
103-- various fields, that way new fields introduced in the library do not break
104-- your code.
105data AcceptRequest = AcceptRequest
106    { acceptSubprotocol :: !(Maybe B.ByteString)
107    -- ^ The subprotocol to speak with the client.  If 'pendingSubprotcols' is
108    -- non-empty, 'acceptSubprotocol' must be one of the subprotocols from the
109    -- list.
110    , acceptHeaders     :: !Headers
111    -- ^ Extra headers to send with the response.
112    }
113
114
115--------------------------------------------------------------------------------
116defaultAcceptRequest :: AcceptRequest
117defaultAcceptRequest = AcceptRequest Nothing []
118
119
120--------------------------------------------------------------------------------
121-- | Utility
122sendResponse :: PendingConnection -> Response -> IO ()
123sendResponse pc rsp = Stream.write (pendingStream pc)
124    (Builder.toLazyByteString (encodeResponse rsp))
125
126
127--------------------------------------------------------------------------------
128-- | Accept a pending connection, turning it into a 'Connection'.
129acceptRequest :: PendingConnection -> IO Connection
130acceptRequest pc = acceptRequestWith pc defaultAcceptRequest
131
132
133--------------------------------------------------------------------------------
134-- | This function is like 'acceptRequest' but allows you to set custom options
135-- using the 'AcceptRequest' datatype.
136acceptRequestWith :: PendingConnection -> AcceptRequest -> IO Connection
137acceptRequestWith pc ar = case find (flip compatible request) protocols of
138    Nothing       -> do
139        sendResponse pc $ response400 versionHeader ""
140        throwIO NotSupported
141    Just protocol -> do
142
143        -- Get requested list of exceptions from client.
144        rqExts <- either throwIO return $
145            getRequestSecWebSocketExtensions request
146
147        -- Set up permessage-deflate extension if configured.
148        pmdExt <- case connectionCompressionOptions (pendingOptions pc) of
149            NoCompression                     -> return Nothing
150            PermessageDeflateCompression pmd0 ->
151                case negotiateDeflate (connectionMessageDataSizeLimit options) (Just pmd0) rqExts of
152                    Left err   -> do
153                        rejectRequestWith pc defaultRejectRequest {rejectMessage = B8.pack err}
154                        throwIO NotSupported
155                    Right pmd1 -> return (Just pmd1)
156
157        -- Set up strict utf8 extension if configured.
158        let unicodeExt =
159                if connectionStrictUnicode (pendingOptions pc)
160                    then Just strictUnicode else Nothing
161
162        -- Final extension list.
163        let exts = catMaybes [pmdExt, unicodeExt]
164
165        let subproto = maybe [] (\p -> [("Sec-WebSocket-Protocol", p)]) $ acceptSubprotocol ar
166            headers = subproto ++ acceptHeaders ar ++ concatMap extHeaders exts
167            response = finishRequest protocol request headers
168
169        either throwIO (sendResponse pc) response
170
171        parseRaw <- decodeMessages
172            protocol
173            (connectionFramePayloadSizeLimit options)
174            (connectionMessageDataSizeLimit options)
175            (pendingStream pc)
176        writeRaw <- encodeMessages protocol ServerConnection (pendingStream pc)
177
178        write <- foldM (\x ext -> extWrite ext x) writeRaw exts
179        parse <- foldM (\x ext -> extParse ext x) parseRaw exts
180
181        sentRef    <- newIORef False
182        let connection = Connection
183                { connectionOptions   = options
184                , connectionType      = ServerConnection
185                , connectionProtocol  = protocol
186                , connectionParse     = parse
187                , connectionWrite     = write
188                , connectionSentClose = sentRef
189                }
190
191        pendingOnAccept pc connection
192        return connection
193  where
194    options       = pendingOptions pc
195    request       = pendingRequest pc
196    versionHeader = [("Sec-WebSocket-Version",
197        B.intercalate ", " $ concatMap headerVersions protocols)]
198
199
200--------------------------------------------------------------------------------
201-- | Parameters that allow you to tweak how a request is rejected.  Please use
202-- 'defaultRejectRequest' and modify fields using record syntax so your code
203-- will not break when new fields are added.
204data RejectRequest = RejectRequest
205    { -- | The status code, 400 by default.
206      rejectCode    :: !Int
207    , -- | The message, "Bad Request" by default
208      rejectMessage :: !B.ByteString
209    , -- | Extra headers to be sent with the response.
210      rejectHeaders :: Headers
211    , -- | Reponse body of the rejection.
212      rejectBody    :: !B.ByteString
213    }
214
215
216--------------------------------------------------------------------------------
217defaultRejectRequest :: RejectRequest
218defaultRejectRequest = RejectRequest
219    { rejectCode    = 400
220    , rejectMessage = "Bad Request"
221    , rejectHeaders = []
222    , rejectBody    = ""
223    }
224
225
226--------------------------------------------------------------------------------
227rejectRequestWith
228    :: PendingConnection  -- ^ Connection to reject
229    -> RejectRequest      -- ^ Params on how to reject the request
230    -> IO ()
231rejectRequestWith pc reject = sendResponse pc $ Response
232    ResponseHead
233        { responseCode    = rejectCode reject
234        , responseMessage = rejectMessage reject
235        , responseHeaders = rejectHeaders reject
236        }
237    (rejectBody reject)
238
239
240--------------------------------------------------------------------------------
241rejectRequest
242    :: PendingConnection  -- ^ Connection to reject
243    -> B.ByteString       -- ^ Rejection response body
244    -> IO ()
245rejectRequest pc body = rejectRequestWith pc
246    defaultRejectRequest {rejectBody = body}
247
248
249--------------------------------------------------------------------------------
250data Connection = Connection
251    { connectionOptions   :: !ConnectionOptions
252    , connectionType      :: !ConnectionType
253    , connectionProtocol  :: !Protocol
254    , connectionParse     :: !(IO (Maybe Message))
255    , connectionWrite     :: !([Message] -> IO ())
256    , connectionSentClose :: !(IORef Bool)
257    -- ^ According to the RFC, both the client and the server MUST send
258    -- a close control message to each other.  Either party can initiate
259    -- the first close message but then the other party must respond.  Finally,
260    -- the server is in charge of closing the TCP connection.  This IORef tracks
261    -- if we have sent a close message and are waiting for the peer to respond.
262    }
263
264
265--------------------------------------------------------------------------------
266receive :: Connection -> IO Message
267receive conn = do
268    mbMsg <- connectionParse conn
269    case mbMsg of
270        Nothing  -> throwIO ConnectionClosed
271        Just msg -> return msg
272
273
274--------------------------------------------------------------------------------
275-- | Receive an application message. Automatically respond to control messages.
276--
277-- When the peer sends a close control message, an exception of type 'CloseRequest'
278-- is thrown.  The peer can send a close control message either to initiate a
279-- close or in response to a close message we have sent to the peer.  In either
280-- case the 'CloseRequest' exception will be thrown.  The RFC specifies that
281-- the server is responsible for closing the TCP connection, which should happen
282-- after receiving the 'CloseRequest' exception from this function.
283--
284-- This will throw 'ConnectionClosed' if the TCP connection dies unexpectedly.
285receiveDataMessage :: Connection -> IO DataMessage
286receiveDataMessage conn = do
287    msg <- receive conn
288    case msg of
289        DataMessage _ _ _ am -> return am
290        ControlMessage cm    -> case cm of
291            Close i closeMsg -> do
292                hasSentClose <- readIORef $ connectionSentClose conn
293                unless hasSentClose $ send conn msg
294                throwIO $ CloseRequest i closeMsg
295            Pong _    -> do
296                connectionOnPong (connectionOptions conn)
297                receiveDataMessage conn
298            Ping pl   -> do
299                send conn (ControlMessage (Pong pl))
300                receiveDataMessage conn
301
302
303--------------------------------------------------------------------------------
304-- | Receive a message, converting it to whatever format is needed.
305receiveData :: WebSocketsData a => Connection -> IO a
306receiveData conn = fromDataMessage <$> receiveDataMessage conn
307
308
309--------------------------------------------------------------------------------
310send :: Connection -> Message -> IO ()
311send conn = sendAll conn . return
312
313--------------------------------------------------------------------------------
314sendAll :: Connection -> [Message] -> IO ()
315sendAll _    []   = return ()
316sendAll conn msgs = do
317    when (any isCloseMessage msgs) $
318      writeIORef (connectionSentClose conn) True
319    connectionWrite conn msgs
320  where
321    isCloseMessage (ControlMessage (Close _ _)) = True
322    isCloseMessage _                            = False
323
324--------------------------------------------------------------------------------
325-- | Send a 'DataMessage'.  This allows you send both human-readable text and
326-- binary data.  This is a slightly more low-level interface than 'sendTextData'
327-- or 'sendBinaryData'.
328sendDataMessage :: Connection -> DataMessage -> IO ()
329sendDataMessage conn = sendDataMessages conn . return
330
331--------------------------------------------------------------------------------
332-- | Send a collection of 'DataMessage's.  This is more efficient than calling
333-- 'sendDataMessage' many times.
334sendDataMessages :: Connection -> [DataMessage] -> IO ()
335sendDataMessages conn = sendAll conn . map (DataMessage False False False)
336
337--------------------------------------------------------------------------------
338-- | Send a textual message.  The message will be encoded as UTF-8.  This should
339-- be the default choice for human-readable text-based protocols such as JSON.
340sendTextData :: WebSocketsData a => Connection -> a -> IO ()
341sendTextData conn = sendTextDatas conn . return
342
343--------------------------------------------------------------------------------
344-- | Send a number of textual messages.  This is more efficient than calling
345-- 'sendTextData' many times.
346sendTextDatas :: WebSocketsData a => Connection -> [a] -> IO ()
347sendTextDatas conn =
348    sendDataMessages conn .
349    map (\x -> Text (toLazyByteString x) Nothing)
350
351--------------------------------------------------------------------------------
352-- | Send a binary message.  This is useful for sending binary blobs, e.g.
353-- images, data encoded with MessagePack, images...
354sendBinaryData :: WebSocketsData a => Connection -> a -> IO ()
355sendBinaryData conn = sendBinaryDatas conn . return
356
357--------------------------------------------------------------------------------
358-- | Send a number of binary messages.  This is more efficient than calling
359-- 'sendBinaryData' many times.
360sendBinaryDatas :: WebSocketsData a => Connection -> [a] -> IO ()
361sendBinaryDatas conn = sendDataMessages conn . map (Binary . toLazyByteString)
362
363--------------------------------------------------------------------------------
364-- | Send a friendly close message.  Note that after sending this message,
365-- you should still continue calling 'receiveDataMessage' to process any
366-- in-flight messages.  The peer will eventually respond with a close control
367-- message of its own which will cause 'receiveDataMessage' to throw the
368-- 'CloseRequest' exception.  This exception is when you can finally consider
369-- the connection closed.
370sendClose :: WebSocketsData a => Connection -> a -> IO ()
371sendClose conn = sendCloseCode conn 1000
372
373
374--------------------------------------------------------------------------------
375-- | Send a friendly close message and close code.  Similar to 'sendClose',
376-- you should continue calling 'receiveDataMessage' until you receive a
377-- 'CloseRequest' exception.
378--
379-- See <http://tools.ietf.org/html/rfc6455#section-7.4> for a list of close
380-- codes.
381sendCloseCode :: WebSocketsData a => Connection -> Word16 -> a -> IO ()
382sendCloseCode conn code =
383    send conn . ControlMessage . Close code . toLazyByteString
384
385
386--------------------------------------------------------------------------------
387-- | Send a ping
388sendPing :: WebSocketsData a => Connection -> a -> IO ()
389sendPing conn = send conn . ControlMessage . Ping . toLazyByteString
390
391
392--------------------------------------------------------------------------------
393-- | Forks a ping thread, sending a ping message every @n@ seconds over the
394-- connection.  The thread is killed when the inner IO action is finished.
395--
396-- This is useful to keep idle connections open through proxies and whatnot.
397-- Many (but not all) proxies have a 60 second default timeout, so based on that
398-- sending a ping every 30 seconds is a good idea.
399withPingThread
400    :: Connection
401    -> Int    -- ^ Second interval in which pings should be sent.
402    -> IO ()  -- ^ Repeat this after sending a ping.
403    -> IO a   -- ^ Application to wrap with a ping thread.
404    -> IO a   -- ^ Executes application and kills ping thread when done.
405withPingThread conn n action app =
406    Async.withAsync (pingThread conn n action) (\_ -> app)
407
408
409--------------------------------------------------------------------------------
410-- | DEPRECATED: Use 'withPingThread' instead.
411--
412-- Forks a ping thread, sending a ping message every @n@ seconds over the
413-- connection.  The thread dies silently if the connection crashes or is closed.
414--
415-- This is useful to keep idle connections open through proxies and whatnot.
416-- Many (but not all) proxies have a 60 second default timeout, so based on that
417-- sending a ping every 30 seconds is a good idea.
418forkPingThread :: Connection -> Int -> IO ()
419forkPingThread conn n = do
420    _ <- forkIO $ pingThread conn n (return ())
421    return ()
422{-# DEPRECATED forkPingThread "Use 'withPingThread' instead" #-}
423
424
425--------------------------------------------------------------------------------
426-- | Use this if you want to run the ping thread yourself.
427--
428-- See also 'withPingThread'.
429pingThread :: Connection -> Int -> IO () -> IO ()
430pingThread conn n action
431    | n <= 0    = return ()
432    | otherwise = ignore `handle` go 1
433  where
434    go :: Int -> IO ()
435    go i = do
436        threadDelay (n * 1000 * 1000)
437        sendPing conn (T.pack $ show i)
438        action
439        go (i + 1)
440
441    ignore e = case fromException e of
442        Just async -> throwIO (async :: AsyncException)
443        Nothing    -> return ()
444