diff --git a/CHANGELOG b/CHANGELOG --- a/CHANGELOG +++ b/CHANGELOG @@ -1,3 +1,9 @@ +- 0.12.0.0 + * Add limit options for frame and message size to prevent against (D)DoS + attacks + * Fix space leak in encodeMessages (by Roman Borschel) + * Stricter frame/encoding decoding for ping/close frames (by Lars Petersen) + - 0.11.2.0 * Fix 0-width reason phrase parsing * Change receive buffer from 1024 to 8192 bytes (by Ondrej Palkovsky) diff --git a/src/Network/WebSockets.hs b/src/Network/WebSockets.hs --- a/src/Network/WebSockets.hs +++ b/src/Network/WebSockets.hs @@ -25,6 +25,9 @@ , PermessageDeflate (..) , defaultPermessageDeflate + -- ** Protection limits + , SizeLimit (..) + -- * Sending and receiving messages , receive , receiveDataMessage diff --git a/src/Network/WebSockets/Client.hs b/src/Network/WebSockets/Client.hs --- a/src/Network/WebSockets/Client.hs +++ b/src/Network/WebSockets/Client.hs @@ -102,9 +102,11 @@ "Network.WebSockets.Client.runClientWithStream: no handshake " ++ "response from server" void $ either throwIO return $ finishResponse protocol request response - parse <- decodeMessages protocol stream - write <- encodeMessages protocol ClientConnection stream - sentRef <- newIORef False + parse <- decodeMessages protocol + (connectionFramePayloadSizeLimit opts) + (connectionMessageDataSizeLimit opts) stream + write <- encodeMessages protocol ClientConnection stream + sentRef <- newIORef False app Connection { connectionOptions = opts diff --git a/src/Network/WebSockets/Connection.hs b/src/Network/WebSockets/Connection.hs --- a/src/Network/WebSockets/Connection.hs +++ b/src/Network/WebSockets/Connection.hs @@ -37,6 +37,8 @@ , CompressionOptions (..) , PermessageDeflate (..) , defaultPermessageDeflate + + , SizeLimit (..) ) where @@ -65,6 +67,7 @@ -------------------------------------------------------------------------------- +import Network.WebSockets.Connection.Options import Network.WebSockets.Extensions as Extensions import Network.WebSockets.Extensions.PermessageDeflate import Network.WebSockets.Extensions.StrictUnicode @@ -142,7 +145,7 @@ pmdExt <- case connectionCompressionOptions (pendingOptions pc) of NoCompression -> return Nothing PermessageDeflateCompression pmd0 -> - case negotiateDeflate (Just pmd0) rqExts of + case negotiateDeflate (connectionMessageDataSizeLimit options) (Just pmd0) rqExts of Left err -> do rejectRequestWith pc defaultRejectRequest {rejectMessage = B8.pack err} throwIO NotSupported @@ -162,7 +165,11 @@ either throwIO (sendResponse pc) response - parseRaw <- decodeMessages protocol (pendingStream pc) + parseRaw <- decodeMessages + protocol + (connectionFramePayloadSizeLimit options) + (connectionMessageDataSizeLimit options) + (pendingStream pc) writeRaw <- encodeMessages protocol ServerConnection (pendingStream pc) write <- foldM (\x ext -> extWrite ext x) writeRaw exts @@ -170,7 +177,7 @@ sentRef <- newIORef False let connection = Connection - { connectionOptions = pendingOptions pc + { connectionOptions = options , connectionType = ServerConnection , connectionProtocol = protocol , connectionParse = parse @@ -181,6 +188,7 @@ pendingOnAccept pc connection return connection where + options = pendingOptions pc request = pendingRequest pc versionHeader = [("Sec-WebSocket-Version", B.intercalate ", " $ concatMap headerVersions protocols)] @@ -249,42 +257,6 @@ -- the server is in charge of closing the TCP connection. This IORef tracks -- if we have sent a close message and are waiting for the peer to respond. } - - --------------------------------------------------------------------------------- --- | Set options for a 'Connection'. -data ConnectionOptions = ConnectionOptions - { connectionOnPong :: !(IO ()) - -- ^ Whenever a 'pong' is received, this IO action is executed. It can be - -- used to tickle connections or fire missiles. - , connectionCompressionOptions :: !CompressionOptions - -- ^ Enable 'PermessageDeflate'. - , connectionStrictUnicode :: !Bool - -- ^ Enable strict unicode on the connection. This means that if a client - -- (or server) sends invalid UTF-8, we will throw a 'UnicodeException' - -- rather than replacing it by the unicode replacement character U+FFFD. - } - - --------------------------------------------------------------------------------- --- | The default connection options: --- --- * Nothing happens when a pong is received. --- * Compression is disabled. --- * Lenient unicode decoding. -defaultConnectionOptions :: ConnectionOptions -defaultConnectionOptions = ConnectionOptions - { connectionOnPong = return () - , connectionCompressionOptions = NoCompression - , connectionStrictUnicode = False - } - - --------------------------------------------------------------------------------- -data CompressionOptions - = NoCompression - | PermessageDeflateCompression PermessageDeflate - deriving (Eq, Show) -------------------------------------------------------------------------------- diff --git a/src/Network/WebSockets/Connection/Options.hs b/src/Network/WebSockets/Connection/Options.hs new file mode 100644 --- /dev/null +++ b/src/Network/WebSockets/Connection/Options.hs @@ -0,0 +1,119 @@ +-------------------------------------------------------------------------------- +module Network.WebSockets.Connection.Options + ( ConnectionOptions (..) + , defaultConnectionOptions + + , CompressionOptions (..) + , PermessageDeflate (..) + , defaultPermessageDeflate + + , SizeLimit (..) + , atMostSizeLimit + ) where + + +-------------------------------------------------------------------------------- +import Data.Int (Int64) + + +-------------------------------------------------------------------------------- +-- | Set options for a 'Connection'. Please do not use this constructor +-- directly, but rather use 'defaultConnectionOptions' and then set the fields +-- you want, e.g.: +-- +-- > myOptions = defaultConnectionOptions {connectionStrictUnicode = True} +-- +-- This way your code does not break if the library introduces new fields. +data ConnectionOptions = ConnectionOptions + { connectionOnPong :: !(IO ()) + -- ^ Whenever a 'pong' is received, this IO action is executed. It can be + -- used to tickle connections or fire missiles. + , connectionCompressionOptions :: !CompressionOptions + -- ^ Enable 'PermessageDeflate'. + , connectionStrictUnicode :: !Bool + -- ^ Enable strict unicode on the connection. This means that if a client + -- (or server) sends invalid UTF-8, we will throw a 'UnicodeException' + -- rather than replacing it by the unicode replacement character U+FFFD. + , connectionFramePayloadSizeLimit :: !SizeLimit + -- ^ The maximum size for incoming frame payload size in bytes. If a + -- frame exceeds this limit, a 'ParseException' is thrown. + , connectionMessageDataSizeLimit :: !SizeLimit + -- ^ 'connectionFrameSizeLimit' is often not enough since a malicious + -- client can send many small frames to create a huge message. This limit + -- allows you to protect from that. If a message exceeds this limit, a + -- 'ParseException' is thrown. + -- + -- Note that, if compression is enabled, we check the size of the + -- compressed messages, as well as the size of the uncompressed messages + -- as we are deflating them to ensure we don't use too much memory in any + -- case. + } + + +-------------------------------------------------------------------------------- +-- | The default connection options: +-- +-- * Nothing happens when a pong is received. +-- * Compression is disabled. +-- * Lenient unicode decoding. +defaultConnectionOptions :: ConnectionOptions +defaultConnectionOptions = ConnectionOptions + { connectionOnPong = return () + , connectionCompressionOptions = NoCompression + , connectionStrictUnicode = False + , connectionFramePayloadSizeLimit = mempty + , connectionMessageDataSizeLimit = mempty + } + + +-------------------------------------------------------------------------------- +data CompressionOptions + = NoCompression + | PermessageDeflateCompression PermessageDeflate + deriving (Eq, Show) + + +-------------------------------------------------------------------------------- +-- | Four extension parameters are defined for "permessage-deflate" to +-- help endpoints manage per-connection resource usage. +-- +-- - "server_no_context_takeover" +-- - "client_no_context_takeover" +-- - "server_max_window_bits" +-- - "client_max_window_bits" +data PermessageDeflate = PermessageDeflate + { serverNoContextTakeover :: Bool + , clientNoContextTakeover :: Bool + , serverMaxWindowBits :: Int + , clientMaxWindowBits :: Int + , pdCompressionLevel :: Int + } deriving (Eq, Show) + + +-------------------------------------------------------------------------------- +defaultPermessageDeflate :: PermessageDeflate +defaultPermessageDeflate = PermessageDeflate False False 15 15 8 + + +-------------------------------------------------------------------------------- +-- | A size limit, in bytes. The 'Monoid' instance takes the minimum limit. +data SizeLimit + = NoSizeLimit + | SizeLimit !Int64 + deriving (Eq, Show) + + +-------------------------------------------------------------------------------- +instance Monoid SizeLimit where + mempty = NoSizeLimit + + mappend NoSizeLimit y = y + mappend x NoSizeLimit = x + mappend (SizeLimit x) (SizeLimit y) = SizeLimit (min x y) + + +-------------------------------------------------------------------------------- +atMostSizeLimit :: Int64 -> SizeLimit -> Bool +atMostSizeLimit _ NoSizeLimit = True +atMostSizeLimit s (SizeLimit l) = s <= l +{-# INLINE atMostSizeLimit #-} diff --git a/src/Network/WebSockets/Extensions/PermessageDeflate.hs b/src/Network/WebSockets/Extensions/PermessageDeflate.hs --- a/src/Network/WebSockets/Extensions/PermessageDeflate.hs +++ b/src/Network/WebSockets/Extensions/PermessageDeflate.hs @@ -7,48 +7,32 @@ ( defaultPermessageDeflate , PermessageDeflate(..) , negotiateDeflate + + -- * Considered internal + , makeMessageInflater + , makeMessageDeflater ) where -------------------------------------------------------------------------------- import Control.Applicative ((<$>)) import Control.Exception (throwIO) -import Control.Monad (foldM) +import Control.Monad (foldM, unless) import qualified Data.ByteString as B import qualified Data.ByteString.Char8 as B8 import qualified Data.ByteString.Lazy as BL import qualified Data.ByteString.Lazy.Char8 as BL8 import qualified Data.ByteString.Lazy.Internal as BL +import Data.Int (Int64) import Data.Monoid import qualified Data.Streaming.Zlib as Zlib +import Network.WebSockets.Connection.Options import Network.WebSockets.Extensions import Network.WebSockets.Extensions.Description import Network.WebSockets.Http import Network.WebSockets.Types -import Text.Read (readMaybe) import Prelude - - --------------------------------------------------------------------------------- --- | Four extension parameters are defined for "permessage-deflate" to --- help endpoints manage per-connection resource usage. --- --- - "server_no_context_takeover" --- - "client_no_context_takeover" --- - "server_max_window_bits" --- - "client_max_window_bits" -data PermessageDeflate = PermessageDeflate - { serverNoContextTakeover :: Bool - , clientNoContextTakeover :: Bool - , serverMaxWindowBits :: Int - , clientMaxWindowBits :: Int - , pdCompressionLevel :: Int - } deriving (Eq, Show) - - --------------------------------------------------------------------------------- -defaultPermessageDeflate :: PermessageDeflate -defaultPermessageDeflate = PermessageDeflate False False 15 15 8 +import Text.Read (readMaybe) -------------------------------------------------------------------------------- @@ -77,13 +61,14 @@ -------------------------------------------------------------------------------- -negotiateDeflate :: Maybe PermessageDeflate -> NegotiateExtension -negotiateDeflate pmd0 exts0 = do +negotiateDeflate + :: SizeLimit -> Maybe PermessageDeflate -> NegotiateExtension +negotiateDeflate messageLimit pmd0 exts0 = do (headers, pmd1) <- negotiateDeflateOpts exts0 pmd0 return Extension { extHeaders = headers , extParse = \parseRaw -> do - inflate <- makeMessageInflater pmd1 + inflate <- makeMessageInflater messageLimit pmd1 return $ do msg <- parseRaw case msg of @@ -230,10 +215,11 @@ deflateBody :: Zlib.Deflate -> BL.ByteString -> IO BL.ByteString deflateBody ptr = fmap maybeStrip . go . BL.toChunks where - go [] = dePopper (Zlib.flushDeflate ptr) + go [] = + dePopper (Zlib.flushDeflate ptr) go (c : cs) = do - bl <- Zlib.feedDeflate ptr c >>= dePopper - (bl <>) <$> go cs + chunk <- Zlib.feedDeflate ptr c >>= dePopper + (chunk <>) <$> go cs -------------------------------------------------------------------------------- @@ -245,9 +231,11 @@ -------------------------------------------------------------------------------- -makeMessageInflater :: Maybe PermessageDeflate -> IO (Message -> IO Message) -makeMessageInflater Nothing = return rejectExtensions -makeMessageInflater (Just pmd) +makeMessageInflater + :: SizeLimit -> Maybe PermessageDeflate + -> IO (Message -> IO Message) +makeMessageInflater _ Nothing = return rejectExtensions +makeMessageInflater messageLimit (Just pmd) | clientNoContextTakeover pmd = return $ \msg -> do ptr <- initInflate pmd @@ -280,9 +268,21 @@ ---------------------------------------------------------------------------- inflateBody :: Zlib.Inflate -> BL.ByteString -> IO BL.ByteString inflateBody ptr = - go . BL.toChunks . (<> appTailL) + go 0 . BL.toChunks . (<> appTailL) where - go [] = BL.fromStrict <$> Zlib.flushInflate ptr - go (c : cs) = do - bl <- Zlib.feedInflate ptr c >>= dePopper - (bl <>) <$> go cs + go :: Int64 -> [B.ByteString] -> IO BL.ByteString + go size0 [] = do + chunk <- Zlib.flushInflate ptr + checkSize (fromIntegral (B.length chunk) + size0) + return (BL.fromStrict chunk) + go size0 (c : cs) = do + chunk <- Zlib.feedInflate ptr c >>= dePopper + let size1 = size0 + BL.length chunk + checkSize size1 + (chunk <>) <$> go size1 cs + + + ---------------------------------------------------------------------------- + checkSize :: Int64 -> IO () + checkSize size = unless (atMostSizeLimit size messageLimit) $ throwIO $ + ParseException $ "Message of size " ++ show size ++ " exceeded limit" diff --git a/src/Network/WebSockets/Hybi13.hs b/src/Network/WebSockets/Hybi13.hs --- a/src/Network/WebSockets/Hybi13.hs +++ b/src/Network/WebSockets/Hybi13.hs @@ -5,12 +5,14 @@ ( headerVersions , finishRequest , finishResponse + , encodeMessage , encodeMessages , decodeMessages , createRequest -- Internal (used for testing) , encodeFrame + , parseFrame ) where @@ -19,7 +21,8 @@ import Control.Applicative (pure, (<$>)) import Control.Arrow (first) import Control.Exception (throwIO) -import Control.Monad (forM, liftM, when) +import Control.Monad (forM, liftM, unless, + when) import Data.Binary.Get (Get, getInt64be, getLazyByteString, getWord16be, getWord8) @@ -39,6 +42,7 @@ -------------------------------------------------------------------------------- +import Network.WebSockets.Connection.Options import Network.WebSockets.Http import Network.WebSockets.Hybi13.Demultiplex import Network.WebSockets.Hybi13.Mask @@ -112,7 +116,7 @@ genRef <- newIORef =<< newStdGen return $ \msgs -> do builders <- forM msgs $ \msg -> - atomicModifyIORef genRef $ \s -> encodeMessage conType s msg + atomicModifyIORef' genRef $ \s -> encodeMessage conType s msg Stream.write stream (B.toLazyByteString $ mconcat builders) @@ -120,13 +124,21 @@ encodeFrame :: Maybe Mask -> Frame -> B.Builder encodeFrame mask f = B.fromWord8 byte0 `mappend` B.fromWord8 byte1 `mappend` len `mappend` maskbytes `mappend` - B.fromLazyByteString (maskPayload mask (framePayload f)) + B.fromLazyByteString (maskPayload mask payload) where + byte0 = fin .|. rsv1 .|. rsv2 .|. rsv3 .|. opcode fin = if frameFin f then 0x80 else 0x00 rsv1 = if frameRsv1 f then 0x40 else 0x00 rsv2 = if frameRsv2 f then 0x20 else 0x00 rsv3 = if frameRsv3 f then 0x10 else 0x00 + payload = case frameType f of + ContinuationFrame -> framePayload f + TextFrame -> framePayload f + BinaryFrame -> framePayload f + CloseFrame -> BL.take 125 $ framePayload f + PingFrame -> BL.take 125 $ framePayload f + PongFrame -> BL.take 125 $ framePayload f opcode = case frameType f of ContinuationFrame -> 0x00 TextFrame -> 0x01 @@ -134,13 +146,12 @@ CloseFrame -> 0x08 PingFrame -> 0x09 PongFrame -> 0x0a - (maskflag, maskbytes) = case mask of Nothing -> (0x00, mempty) Just m -> (0x80, encodeMask m) byte1 = maskflag .|. lenflag - len' = BL.length (framePayload f) + len' = BL.length payload (lenflag, len) | len' < 126 = (fromIntegral len', mempty) | len' < 0x10000 = (126, B.fromWord16be (fromIntegral len')) @@ -149,19 +160,21 @@ -------------------------------------------------------------------------------- decodeMessages - :: Stream + :: SizeLimit + -> SizeLimit + -> Stream -> IO (IO (Maybe Message)) -decodeMessages stream = do +decodeMessages frameLimit messageLimit stream = do dmRef <- newIORef emptyDemultiplexState return $ go dmRef where go dmRef = do - mbFrame <- Stream.parseBin stream parseFrame + mbFrame <- Stream.parseBin stream (parseFrame frameLimit) case mbFrame of Nothing -> return Nothing Just frame -> do demultiplexResult <- atomicModifyIORef' dmRef $ - \s -> swap $ demultiplex s frame + \s -> swap $ demultiplex messageLimit s frame case demultiplexResult of DemultiplexError err -> throwIO err DemultiplexContinue -> go dmRef @@ -170,8 +183,8 @@ -------------------------------------------------------------------------------- -- | Parse a frame -parseFrame :: Get Frame -parseFrame = do +parseFrame :: SizeLimit -> Get Frame +parseFrame frameSizeLimit = do byte0 <- getWord8 let fin = byte0 .&. 0x80 == 0x80 rsv1 = byte0 .&. 0x40 == 0x40 @@ -179,15 +192,6 @@ rsv3 = byte0 .&. 0x10 == 0x10 opcode = byte0 .&. 0x0f - ft <- case opcode of - 0x00 -> return ContinuationFrame - 0x01 -> return TextFrame - 0x02 -> return BinaryFrame - 0x08 -> return CloseFrame - 0x09 -> return PingFrame - 0x0a -> return PongFrame - _ -> fail $ "Unknown opcode: " ++ show opcode - byte1 <- getWord8 let mask = byte1 .&. 0x80 == 0x80 lenflag = byte1 .&. 0x7f @@ -197,11 +201,30 @@ 127 -> getInt64be _ -> return (fromIntegral lenflag) + -- Check size against limit. + unless (atMostSizeLimit len frameSizeLimit) $ + fail $ "Frame of size " ++ show len ++ " exceeded limit" + + ft <- case opcode of + 0x00 -> return ContinuationFrame + 0x01 -> return TextFrame + 0x02 -> return BinaryFrame + 0x08 -> enforceControlFrameRestrictions len fin >> return CloseFrame + 0x09 -> enforceControlFrameRestrictions len fin >> return PingFrame + 0x0a -> enforceControlFrameRestrictions len fin >> return PongFrame + _ -> fail $ "Unknown opcode: " ++ show opcode + masker <- maskPayload <$> if mask then Just <$> parseMask else pure Nothing chunks <- getLazyByteString len return $ Frame fin rsv1 rsv2 rsv3 ft (masker chunks) + + where + enforceControlFrameRestrictions len fin + | not fin = fail "Control Frames must not be fragmented!" + | len > 125 = fail "Control Frames must not carry payload > 125 bytes!" + | otherwise = pure () -------------------------------------------------------------------------------- hashKey :: ByteString -> ByteString diff --git a/src/Network/WebSockets/Hybi13/Demultiplex.hs b/src/Network/WebSockets/Hybi13/Demultiplex.hs --- a/src/Network/WebSockets/Hybi13/Demultiplex.hs +++ b/src/Network/WebSockets/Hybi13/Demultiplex.hs @@ -13,16 +13,15 @@ -------------------------------------------------------------------------------- -import Blaze.ByteString.Builder (Builder) -import qualified Blaze.ByteString.Builder as B -import Control.Exception (Exception) -import Data.Binary.Get (getWord16be, runGet) -import qualified Data.ByteString.Lazy as BL -import Data.Monoid (mappend) -import Data.Typeable (Typeable) - - --------------------------------------------------------------------------------- +import Blaze.ByteString.Builder (Builder) +import qualified Blaze.ByteString.Builder as B +import Control.Exception (Exception) +import Data.Binary.Get (getWord16be, runGet) +import qualified Data.ByteString.Lazy as BL +import Data.Int (Int64) +import Data.Monoid (mappend) +import Data.Typeable (Typeable) +import Network.WebSockets.Connection.Options import Network.WebSockets.Types @@ -64,7 +63,7 @@ -- | Internal state used by the demultiplexer data DemultiplexState = EmptyDemultiplexState - | DemultiplexState !Builder !(Builder -> Message) + | DemultiplexState !Int64 !Builder !(Builder -> Message) -------------------------------------------------------------------------------- @@ -81,20 +80,21 @@ -------------------------------------------------------------------------------- -demultiplex :: DemultiplexState +demultiplex :: SizeLimit + -> DemultiplexState -> Frame -> (DemultiplexResult, DemultiplexState) -demultiplex state (Frame True False False False PingFrame pl) +demultiplex _ state (Frame True False False False PingFrame pl) | BL.length pl > 125 = (DemultiplexError $ CloseRequest 1002 "Protocol Error", emptyDemultiplexState) | otherwise = (DemultiplexSuccess $ ControlMessage (Ping pl), state) -demultiplex state (Frame True False False False PongFrame pl) = +demultiplex _ state (Frame True False False False PongFrame pl) = (DemultiplexSuccess (ControlMessage (Pong pl)), state) -demultiplex _ (Frame True False False False CloseFrame pl) = +demultiplex _ _ (Frame True False False False CloseFrame pl) = (DemultiplexSuccess (ControlMessage (uncurry Close parsedClose)), emptyDemultiplexState) where -- The Close frame MAY contain a body (the "Application data" portion of the @@ -114,32 +114,44 @@ | BL.length pl == 1 = (1002, BL.empty) | otherwise = (1000, BL.empty) -demultiplex EmptyDemultiplexState (Frame fin rsv1 rsv2 rsv3 tp pl) = case tp of +demultiplex sizeLimit EmptyDemultiplexState (Frame fin rsv1 rsv2 rsv3 tp pl) = case tp of + _ | not (atMostSizeLimit size sizeLimit) -> + ( DemultiplexError $ ParseException $ + "Message of size " ++ show size ++ " exceeded limit" + , emptyDemultiplexState + ) TextFrame | fin -> (DemultiplexSuccess (text pl), emptyDemultiplexState) | otherwise -> - (DemultiplexContinue, DemultiplexState plb (text . B.toLazyByteString)) + (DemultiplexContinue, DemultiplexState size plb (text . B.toLazyByteString)) BinaryFrame | fin -> (DemultiplexSuccess (binary pl), emptyDemultiplexState) - | otherwise -> (DemultiplexContinue, DemultiplexState plb (binary . B.toLazyByteString)) + | otherwise -> (DemultiplexContinue, DemultiplexState size plb (binary . B.toLazyByteString)) _ -> (DemultiplexError $ CloseRequest 1002 "Protocol Error", emptyDemultiplexState) where - plb = B.fromLazyByteString pl + size = BL.length pl + plb = B.fromLazyByteString pl text x = DataMessage rsv1 rsv2 rsv3 (Text x Nothing) binary x = DataMessage rsv1 rsv2 rsv3 (Binary x) -demultiplex (DemultiplexState b f) (Frame fin False False False ContinuationFrame pl) - | fin = (DemultiplexSuccess (f b'), emptyDemultiplexState) - | otherwise = (DemultiplexContinue, DemultiplexState b' f) +demultiplex sizeLimit (DemultiplexState size0 b f) (Frame fin False False False ContinuationFrame pl) + | not (atMostSizeLimit size1 sizeLimit) = + ( DemultiplexError $ ParseException $ + "Message of size " ++ show size1 ++ " exceeded limit" + , emptyDemultiplexState + ) + | fin = (DemultiplexSuccess (f b'), emptyDemultiplexState) + | otherwise = (DemultiplexContinue, DemultiplexState size1 b' f) where - b' = b `mappend` plb - plb = B.fromLazyByteString pl + size1 = size0 + BL.length pl + b' = b `mappend` plb + plb = B.fromLazyByteString pl -demultiplex _ _ = +demultiplex _ _ _ = (DemultiplexError (CloseRequest 1002 "Protocol Error"), emptyDemultiplexState) diff --git a/src/Network/WebSockets/Protocol.hs b/src/Network/WebSockets/Protocol.hs --- a/src/Network/WebSockets/Protocol.hs +++ b/src/Network/WebSockets/Protocol.hs @@ -16,14 +16,15 @@ -------------------------------------------------------------------------------- -import Data.ByteString (ByteString) -import qualified Data.ByteString as B +import Data.ByteString (ByteString) +import qualified Data.ByteString as B -------------------------------------------------------------------------------- +import Network.WebSockets.Connection.Options import Network.WebSockets.Http -import qualified Network.WebSockets.Hybi13 as Hybi13 -import Network.WebSockets.Stream (Stream) +import qualified Network.WebSockets.Hybi13 as Hybi13 +import Network.WebSockets.Stream (Stream) import Network.WebSockets.Types @@ -77,9 +78,10 @@ -------------------------------------------------------------------------------- decodeMessages - :: Protocol -> Stream + :: Protocol -> SizeLimit -> SizeLimit -> Stream -> IO (IO (Maybe Message)) -decodeMessages Hybi13 = Hybi13.decodeMessages +decodeMessages Hybi13 frameLimit messageLimit = + Hybi13.decodeMessages frameLimit messageLimit -------------------------------------------------------------------------------- diff --git a/src/Network/WebSockets/Stream.hs b/src/Network/WebSockets/Stream.hs --- a/src/Network/WebSockets/Stream.hs +++ b/src/Network/WebSockets/Stream.hs @@ -20,7 +20,7 @@ import qualified Data.Attoparsec.ByteString as Atto import qualified Data.ByteString as B import qualified Data.ByteString.Lazy as BL -import Data.IORef (IORef, atomicModifyIORef, +import Data.IORef (IORef, atomicModifyIORef', newIORef, readIORef, writeIORef) import qualified Network.Socket as S @@ -75,7 +75,7 @@ return $ Stream (receive' ref receiveLock) (send' ref sendLock) ref where closeRef :: IORef StreamState -> IO () - closeRef ref = atomicModifyIORef ref $ \state -> case state of + closeRef ref = atomicModifyIORef' ref $ \state -> case state of Open buf -> (Closed buf, ()) Closed buf -> (Closed buf, ()) diff --git a/src/Network/WebSockets/Types.hs b/src/Network/WebSockets/Types.hs --- a/src/Network/WebSockets/Types.hs +++ b/src/Network/WebSockets/Types.hs @@ -160,7 +160,7 @@ -- | The client sent invalid UTF-8. Note that this exception will only be -- thrown if strict decoding is set in the connection options. | UnicodeException String - deriving (Show, Typeable) + deriving (Eq, Show, Typeable) -------------------------------------------------------------------------------- diff --git a/tests/haskell/Network/WebSockets/Extensions/PermessageDeflate/Tests.hs b/tests/haskell/Network/WebSockets/Extensions/PermessageDeflate/Tests.hs new file mode 100644 --- /dev/null +++ b/tests/haskell/Network/WebSockets/Extensions/PermessageDeflate/Tests.hs @@ -0,0 +1,46 @@ +-------------------------------------------------------------------------------- +{-# LANGUAGE OverloadedStrings #-} +module Network.WebSockets.Extensions.PermessageDeflate.Tests + ( tests + ) where + + +-------------------------------------------------------------------------------- +import Control.Exception (try) +import qualified Data.ByteString.Lazy as BL +import Network.WebSockets.Extensions.PermessageDeflate +import Network.WebSockets.Types +import Network.WebSockets.Connection.Options +import Test.Framework (Test, + testGroup) +import Test.Framework.Providers.HUnit (testCase) +import Test.HUnit (Assertion, + (@?=)) + + +-------------------------------------------------------------------------------- +tests :: Test +tests = testGroup "Network.WebSockets.Extensions.PermessageDeflate.Tests" + [ testCase "OK 1" $ do + inflater <- makeMessageInflater + (SizeLimit 100) (Just defaultPermessageDeflate) + message <- inflater $ DataMessage True False False (Binary deflated100) + message @?= + DataMessage False False False (Binary inflated100) + , testCase "Exceed 1" $ do + inflater <- makeMessageInflater + (SizeLimit 99) (Just defaultPermessageDeflate) + assertParseException $ + inflater $ DataMessage True False False (Binary deflated100) + ] + where + assertParseException :: IO a -> Assertion + assertParseException io = do + errOrX <- try io + case errOrX of + Left (ParseException _) -> return () + _ -> fail "Excepted ParseException" + + -- This inflates to 100 bytes. + deflated100 = "b`\160=\NUL\NUL" + inflated100 = BL.replicate 100 0 diff --git a/tests/haskell/Network/WebSockets/Hybi13/Demultiplex/Tests.hs b/tests/haskell/Network/WebSockets/Hybi13/Demultiplex/Tests.hs new file mode 100644 --- /dev/null +++ b/tests/haskell/Network/WebSockets/Hybi13/Demultiplex/Tests.hs @@ -0,0 +1,72 @@ +-------------------------------------------------------------------------------- +{-# LANGUAGE OverloadedStrings #-} +module Network.WebSockets.Hybi13.Demultiplex.Tests + ( tests + ) where + + +-------------------------------------------------------------------------------- +import Control.Applicative ((<$>)) +import qualified Data.ByteString.Lazy as BL +import Network.WebSockets +import Network.WebSockets.Hybi13.Demultiplex +import Prelude +import Test.Framework (Test, testGroup) +import Test.Framework.Providers.HUnit (testCase) +import Test.HUnit (Assertion, (@=?)) + + +-------------------------------------------------------------------------------- +tests :: Test +tests = testGroup "Network.WebSockets.Hybi13.Demultiplex.Tests" + [ testMessageDataSizeLimit + ] + + +-------------------------------------------------------------------------------- +testMessageDataSizeLimit :: Test +testMessageDataSizeLimit = testGroup "testMessageDataSizeLimit Hybi13" + [ testCase "OK 1" $ + Right [DataMessage False False False (Binary (mkZeroes 100))] @=? + testDemultiplex (SizeLimit 100) (fragmented 5 20) + , testCase "Exceeds 1" $ + assertLeft $ + testDemultiplex (SizeLimit 99) (fragmented 5 20) + , testCase "Exceeds 2" $ + assertLeft $ + testDemultiplex (SizeLimit 100) (fragmented 6 20) + , testCase "Exceeds 3" $ + assertLeft $ + testDemultiplex (SizeLimit 100) (fragmented 101 1) + , testCase "Exceeds 4" $ + assertLeft $ + testDemultiplex (SizeLimit 100) (fragmented 1 101) + ] + where + fragmented :: Int -> Int -> [Frame] + fragmented n size = + let payload = mkZeroes size in + [Frame False False False False BinaryFrame payload] ++ + replicate (n - 2) (Frame False False False False ContinuationFrame payload) ++ + [Frame True False False False ContinuationFrame payload] + + mkZeroes :: Int -> BL.ByteString + mkZeroes size = BL.replicate (fromIntegral size) 0 + + assertLeft :: Either a b -> Assertion + assertLeft (Left _) = return () + assertLeft (Right _) = fail "Expecting test to fail" + + +-------------------------------------------------------------------------------- +testDemultiplex + :: SizeLimit + -> [Frame] + -> Either ConnectionException [Message] +testDemultiplex messageLimit = go emptyDemultiplexState + where + go _state0 [] = return [] + go state0 (frame : frames) = case demultiplex messageLimit state0 frame of + (DemultiplexContinue, state1) -> go state1 frames + (DemultiplexError err, _) -> Left err + (DemultiplexSuccess m, state1) -> (m :) <$> go state1 frames diff --git a/tests/haskell/Network/WebSockets/Tests.hs b/tests/haskell/Network/WebSockets/Tests.hs --- a/tests/haskell/Network/WebSockets/Tests.hs +++ b/tests/haskell/Network/WebSockets/Tests.hs @@ -12,19 +12,10 @@ import Control.Concurrent (forkIO) import Control.Exception (try) import Control.Monad (replicateM) +import Data.Binary.Get (runGetOrFail) import qualified Data.ByteString.Lazy as BL import Data.List (intersperse) import Data.Maybe (catMaybes) -import Test.Framework (Test, testGroup) -import Test.Framework.Providers.QuickCheck2 (testProperty) -import Test.HUnit ((@=?)) -import Test.QuickCheck (Arbitrary (..), Gen, - Property) -import qualified Test.QuickCheck as QC -import qualified Test.QuickCheck.Monadic as QC - - --------------------------------------------------------------------------------- import Network.WebSockets import qualified Network.WebSockets.Hybi13 as Hybi13 import Network.WebSockets.Hybi13.Demultiplex @@ -32,6 +23,15 @@ import qualified Network.WebSockets.Stream as Stream import Network.WebSockets.Tests.Util import Network.WebSockets.Types +import Test.Framework (Test, testGroup) +import Test.Framework.Providers.HUnit (testCase) +import Test.Framework.Providers.QuickCheck2 (testProperty) +import Test.HUnit ((@=?)) +import Test.QuickCheck (Arbitrary (..), Gen, + Property) +import qualified Test.QuickCheck as QC +import qualified Test.QuickCheck.Monadic as QC +import Prelude -------------------------------------------------------------------------------- @@ -39,15 +39,17 @@ tests = testGroup "Network.WebSockets.Test" [ testProperty "simple encode/decode Hybi13" (testSimpleEncodeDecode Hybi13) , testProperty "fragmented Hybi13" testFragmentedHybi13 + , testRfc_6455_5_5_1 + , testRfc_6455_5_5_2 + , testFramePayloadSizeLimit ] - -------------------------------------------------------------------------------- testSimpleEncodeDecode :: Protocol -> Property testSimpleEncodeDecode protocol = QC.monadicIO $ QC.forAllM QC.arbitrary $ \msgs -> QC.run $ do echo <- Stream.makeEchoStream - parse <- decodeMessages protocol echo + parse <- decodeMessages protocol mempty mempty echo write <- encodeMessages protocol ClientConnection echo _ <- forkIO $ write msgs msgs' <- catMaybes <$> replicateM (length msgs) parse @@ -60,7 +62,7 @@ testFragmentedHybi13 = QC.monadicIO $ QC.forAllM QC.arbitrary $ \fragmented -> QC.run $ do echo <- Stream.makeEchoStream - parse <- Hybi13.decodeMessages echo + parse <- Hybi13.decodeMessages mempty mempty echo -- is' <- Streams.filter isDataMessage =<< Hybi13.decodeMessages is -- Simple hacky encoding of all frames @@ -86,6 +88,58 @@ Left _ -> return [] Right (Just msg) -> (msg :) <$> parseAll parse Right Nothing -> return [] + +-------------------------------------------------------------------------------- +testRfc_6455_5_5_1 :: Test +testRfc_6455_5_5_1 = + testCase "RFC 6455, 5.5: Frame encoder shall truncate control frame payload to 125 bytes" $ do + 260 @=? BL.length (encodedFrame ContinuationFrame) + 260 @=? BL.length (encodedFrame TextFrame) + 260 @=? BL.length (encodedFrame BinaryFrame) + 127 @=? BL.length (encodedFrame CloseFrame) + 127 @=? BL.length (encodedFrame PingFrame) + 127 @=? BL.length (encodedFrame PongFrame) + where + payload256 = BL.replicate 256 0 + encodedFrame ft + = Builder.toLazyByteString + $ Hybi13.encodeFrame Nothing (Frame True False False False ft payload256) + +-------------------------------------------------------------------------------- +testRfc_6455_5_5_2 :: Test +testRfc_6455_5_5_2 = + testCase "RFC 6455, 5.5: Frame decoder shall fail if control frame payload length > 125 bytes" $ + Left (BL.drop 4 ping126, 4, errMsg) @=? runGetOrFail (Hybi13.parseFrame mempty) ping126 + where + errMsg = "Control Frames must not carry payload > 125 bytes!" + ping126 = mconcat + [ "\137\254\NUL~\219\SI\190\252\219\SI\190\252\219\SI\190\252\219" + , "\SI\190\252\219\SI\190\252\219\SI\190\252\219\SI\190\252\219" + , "\SI\190\252\219\SI\190\252\219\SI\190\252\219\SI\190\252\219" + , "\SI\190\252\219\SI\190\252\219\SI\190\252\219\SI\190\252\219" + , "\SI\190\252\219\SI\190\252\219\SI\190\252\219\SI\190\252\219" + , "\SI\190\252\219\SI\190\252\219\SI\190\252\219\SI\190\252\219" + , "\SI\190\252\219\SI\190\252\219\SI\190\252\219\SI\190\252\219" + , "\SI\190\252\219\SI\190\252\219\SI\190\252\219\SI\190\252\219" + , "\SI\190\252\219\SI" + ] + +testFramePayloadSizeLimit :: Test +testFramePayloadSizeLimit = testGroup "FramePayloadSizeLimit Hybi13" + [ testCase "OK 1" $ case parse (frame 99) of + Right _ -> return () + Left _ -> fail "Expecting successful parse." + , testCase "OK 2" $ case parse (frame 100) of + Right _ -> return () + Left _ -> fail "Expecting successful parse." + , testCase "Exceed" $ case parse (frame 101) of + Right _ -> fail "Expecting parse to fail." + Left _ -> return () + ] + where + parse = runGetOrFail (Hybi13.parseFrame (SizeLimit 100)) + frame n = Builder.toLazyByteString $ Hybi13.encodeFrame Nothing $ + Frame True False False False BinaryFrame (BL.replicate n 20) -------------------------------------------------------------------------------- diff --git a/tests/haskell/TestSuite.hs b/tests/haskell/TestSuite.hs --- a/tests/haskell/TestSuite.hs +++ b/tests/haskell/TestSuite.hs @@ -1,19 +1,23 @@ -------------------------------------------------------------------------------- import qualified Network.WebSockets.Extensions.Tests +import qualified Network.WebSockets.Extensions.PermessageDeflate.Tests import qualified Network.WebSockets.Handshake.Tests import qualified Network.WebSockets.Http.Tests +import qualified Network.WebSockets.Hybi13.Demultiplex.Tests import qualified Network.WebSockets.Mask.Tests import qualified Network.WebSockets.Server.Tests import qualified Network.WebSockets.Tests -import Test.Framework (defaultMain) +import Test.Framework (defaultMain) -------------------------------------------------------------------------------- main :: IO () main = defaultMain [ Network.WebSockets.Extensions.Tests.tests + , Network.WebSockets.Extensions.PermessageDeflate.Tests.tests , Network.WebSockets.Handshake.Tests.tests , Network.WebSockets.Http.Tests.tests + , Network.WebSockets.Hybi13.Demultiplex.Tests.tests , Network.WebSockets.Server.Tests.tests , Network.WebSockets.Mask.Tests.tests , Network.WebSockets.Tests.tests diff --git a/websockets.cabal b/websockets.cabal --- a/websockets.cabal +++ b/websockets.cabal @@ -1,5 +1,5 @@ Name: websockets -Version: 0.11.2.0 +Version: 0.12.0.0 Synopsis: A sensible and clean way to write WebSocket-capable servers in Haskell. @@ -62,6 +62,7 @@ Other-modules: Network.WebSockets.Client + Network.WebSockets.Connection.Options Network.WebSockets.Extensions.Description Network.WebSockets.Extensions.PermessageDeflate Network.WebSockets.Extensions.StrictUnicode @@ -100,9 +101,11 @@ Network.WebSockets Network.WebSockets.Client Network.WebSockets.Connection + Network.WebSockets.Connection.Options Network.WebSockets.Extensions Network.WebSockets.Extensions.Description Network.WebSockets.Extensions.PermessageDeflate + Network.WebSockets.Extensions.PermessageDeflate.Tests Network.WebSockets.Extensions.StrictUnicode Network.WebSockets.Extensions.Tests Network.WebSockets.Handshake.Tests @@ -110,6 +113,7 @@ Network.WebSockets.Http.Tests Network.WebSockets.Hybi13 Network.WebSockets.Hybi13.Demultiplex + Network.WebSockets.Hybi13.Demultiplex.Tests Network.WebSockets.Hybi13.Mask Network.WebSockets.Mask.Tests Network.WebSockets.Protocol