module Network.WebSockets.Hybi13
( headerVersions
, finishRequest
, finishResponse
, encodeMessages
, decodeMessages
, createRequest
, encodeFrame
) where
import qualified Blaze.ByteString.Builder as B
import Control.Applicative (pure, (<$>))
import Control.Arrow (first)
import Control.Exception (throwIO)
import Control.Monad (forM, liftM, when)
import Data.Binary.Get (Get, getInt64be,
getLazyByteString,
getWord16be, getWord8)
import Data.Binary.Put (putWord16be, runPut)
import Data.Bits ((.&.), (.|.))
import Data.ByteString (ByteString)
import qualified Data.ByteString.Base64 as B64
import Data.ByteString.Char8 ()
import qualified Data.ByteString.Lazy as BL
import Data.Digest.Pure.SHA (bytestringDigest, sha1)
import Data.IORef
import Data.Monoid (mappend, mconcat,
mempty)
import Data.Tuple (swap)
import System.Entropy as R
import System.Random (RandomGen, newStdGen)
import Network.WebSockets.Http
import Network.WebSockets.Hybi13.Demultiplex
import Network.WebSockets.Hybi13.Mask
import Network.WebSockets.Stream (Stream)
import qualified Network.WebSockets.Stream as Stream
import Network.WebSockets.Types
headerVersions :: [ByteString]
headerVersions = ["13"]
finishRequest :: RequestHead
-> Headers
-> Either HandshakeException Response
finishRequest reqHttp headers = do
!key <- getRequestHeader reqHttp "Sec-WebSocket-Key"
let !hash = hashKey key
!encoded = B64.encode hash
return $ response101 (("Sec-WebSocket-Accept", encoded):headers) ""
finishResponse :: RequestHead
-> ResponseHead
-> Either HandshakeException Response
finishResponse request response = do
when (responseCode response /= 101) $ Left $
MalformedResponse response "Wrong response status or message."
key <- getRequestHeader request "Sec-WebSocket-Key"
responseHash <- getResponseHeader response "Sec-WebSocket-Accept"
let challengeHash = B64.encode $ hashKey key
when (responseHash /= challengeHash) $ Left $
MalformedResponse response "Challenge and response hashes do not match."
return $ Response response ""
encodeMessage :: RandomGen g => ConnectionType -> g -> Message -> (g, B.Builder)
encodeMessage conType gen msg = (gen', builder)
where
mkFrame = Frame True False False False
(mask, gen') = case conType of
ServerConnection -> (Nothing, gen)
ClientConnection -> first Just (randomMask gen)
builder = encodeFrame mask $ case msg of
(ControlMessage (Close code pl)) -> mkFrame CloseFrame $
runPut (putWord16be code) `mappend` pl
(ControlMessage (Ping pl)) -> mkFrame PingFrame pl
(ControlMessage (Pong pl)) -> mkFrame PongFrame pl
(DataMessage rsv1 rsv2 rsv3 (Text pl _)) -> Frame True rsv1 rsv2 rsv3 TextFrame pl
(DataMessage rsv1 rsv2 rsv3 (Binary pl)) -> Frame True rsv1 rsv2 rsv3 BinaryFrame pl
encodeMessages
:: ConnectionType
-> Stream
-> IO ([Message] -> IO ())
encodeMessages conType stream = do
genRef <- newIORef =<< newStdGen
return $ \msgs -> do
builders <- forM msgs $ \msg ->
atomicModifyIORef genRef $ \s -> encodeMessage conType s msg
Stream.write stream (B.toLazyByteString $ mconcat builders)
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))
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
opcode = case frameType f of
ContinuationFrame -> 0x00
TextFrame -> 0x01
BinaryFrame -> 0x02
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)
(lenflag, len)
| len' < 126 = (fromIntegral len', mempty)
| len' < 0x10000 = (126, B.fromWord16be (fromIntegral len'))
| otherwise = (127, B.fromWord64be (fromIntegral len'))
decodeMessages
:: Stream
-> IO (IO (Maybe Message))
decodeMessages stream = do
dmRef <- newIORef emptyDemultiplexState
return $ go dmRef
where
go dmRef = do
mbFrame <- Stream.parseBin stream parseFrame
case mbFrame of
Nothing -> return Nothing
Just frame -> do
demultiplexResult <- atomicModifyIORef' dmRef $
\s -> swap $ demultiplex s frame
case demultiplexResult of
DemultiplexError err -> throwIO err
DemultiplexContinue -> go dmRef
DemultiplexSuccess msg -> return (Just msg)
parseFrame :: Get Frame
parseFrame = do
byte0 <- getWord8
let fin = byte0 .&. 0x80 == 0x80
rsv1 = byte0 .&. 0x40 == 0x40
rsv2 = byte0 .&. 0x20 == 0x20
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
len <- case lenflag of
126 -> fromIntegral <$> getWord16be
127 -> getInt64be
_ -> return (fromIntegral lenflag)
masker <- maskPayload <$> if mask then Just <$> parseMask else pure Nothing
chunks <- getLazyByteString len
return $ Frame fin rsv1 rsv2 rsv3 ft (masker chunks)
hashKey :: ByteString -> ByteString
hashKey key = unlazy $ bytestringDigest $ sha1 $ lazy $ key `mappend` guid
where
guid = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
lazy = BL.fromChunks . return
unlazy = mconcat . BL.toChunks
createRequest :: ByteString
-> ByteString
-> Bool
-> Headers
-> IO RequestHead
createRequest hostname path secure customHeaders = do
key <- B64.encode `liftM` getEntropy 16
return $ RequestHead path (headers key ++ customHeaders) secure
where
headers key =
[ ("Host" , hostname )
, ("Connection" , "Upgrade" )
, ("Upgrade" , "websocket" )
, ("Sec-WebSocket-Key" , key )
, ("Sec-WebSocket-Version" , versionNumber)
]
versionNumber = head headerVersions