module Network.WebSockets.Connection
( makeSocket
, closeSocket
, PendingConnection (..)
, makePendingConnection
, AcceptRequest(..)
, acceptRequest
, acceptRequestWith
, rejectRequest
, Connection (..)
, ConnectionOptions (..)
, defaultConnectionOptions
, receive
, receiveDataMessage
, receiveData
, send
, sendDataMessage
, sendTextData
, sendBinaryData
, sendClose
, sendCloseCode
, sendPing
) where
import Blaze.ByteString.Builder (Builder)
import qualified Blaze.ByteString.Builder as Builder
import Control.Exception (throw)
import Control.Monad (unless)
import qualified Data.ByteString as B
import Data.List (find)
import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import Data.Word (Word16)
import System.IO.Streams (InputStream, OutputStream)
import qualified System.IO.Streams as Streams
import qualified System.IO.Streams.Attoparsec as Streams
import Network.Socket (Socket)
import qualified Network.Socket as S
import Network.WebSockets.Http
import Network.WebSockets.Protocol
import Network.WebSockets.Types
makeSocket :: String -> Int -> IO Socket
makeSocket host port = do
sock <- S.socket S.AF_INET S.Stream S.defaultProtocol
_ <- S.setSocketOption sock S.ReuseAddr 1
host' <- S.inet_addr host
S.bindSocket sock (S.SockAddrInet (fromIntegral port) host')
S.listen sock 5
return sock
closeSocket :: Socket -> IO ()
closeSocket sock = S.sClose sock
data PendingConnection = PendingConnection
{ pendingOptions :: ConnectionOptions
, pendingRequest :: RequestHead
, pendingOnAccept :: Connection -> IO ()
, pendingIn :: InputStream B.ByteString
, pendingOut :: OutputStream Builder
}
makePendingConnection :: Socket -> IO PendingConnection
makePendingConnection sock = do
(conn, _) <- S.accept sock
(sIn, sOut) <- Streams.socketToStreams conn
bOut <- Streams.builderStream sOut
request <- Streams.parseFromStream (decodeRequestHead False) sIn
let pc = PendingConnection
{ pendingOptions = defaultConnectionOptions
, pendingRequest = request
, pendingOnAccept = \_ -> return ()
, pendingIn = sIn
, pendingOut = bOut
}
return pc
data AcceptRequest = AcceptRequest
{ acceptSubprotocol :: Maybe B.ByteString
}
sendResponse :: PendingConnection -> Response -> IO ()
sendResponse pc rsp = do
Streams.write (Just (encodeResponse rsp)) (pendingOut pc)
Streams.write (Just Builder.flush) (pendingOut pc)
acceptRequest :: PendingConnection -> IO Connection
acceptRequest pc = acceptRequestWith pc $ AcceptRequest Nothing
acceptRequestWith :: PendingConnection -> AcceptRequest -> IO Connection
acceptRequestWith pc ar = case find (flip compatible request) protocols of
Nothing -> do
sendResponse pc $ response400 versionHeader ""
throw NotSupported
Just protocol -> do
let subproto = maybe [] (\p -> [("Sec-WebSocket-Protocol", p)]) $ acceptSubprotocol ar
response = finishRequest protocol request subproto
sendResponse pc response
msgIn <- decodeMessages protocol (pendingIn pc)
msgOut <- encodeMessages protocol ServerConnection (pendingOut pc)
sentRef <- newIORef False
let connection = Connection
{ connectionOptions = pendingOptions pc
, connectionType = ServerConnection
, connectionProtocol = protocol
, connectionIn = msgIn
, connectionOut = msgOut
, connectionSentClose = sentRef
}
pendingOnAccept pc connection
return connection
where
request = pendingRequest pc
versionHeader = [("Sec-WebSocket-Version",
B.intercalate ", " $ concatMap headerVersions protocols)]
rejectRequest :: PendingConnection -> B.ByteString -> IO ()
rejectRequest pc message = sendResponse pc $ response400 [] message
data Connection = Connection
{ connectionOptions :: ConnectionOptions
, connectionType :: ConnectionType
, connectionProtocol :: Protocol
, connectionIn :: InputStream Message
, connectionOut :: OutputStream Message
, connectionSentClose :: IORef Bool
}
data ConnectionOptions = ConnectionOptions
{ connectionOnPong :: IO ()
}
defaultConnectionOptions :: ConnectionOptions
defaultConnectionOptions = ConnectionOptions
{ connectionOnPong = return ()
}
receive :: Connection -> IO Message
receive conn = do
mmsg <- Streams.read (connectionIn conn)
case mmsg of
Nothing -> throw ConnectionClosed
Just msg -> return msg
receiveDataMessage :: Connection -> IO DataMessage
receiveDataMessage conn = do
msg <- receive conn
case msg of
DataMessage am -> return am
ControlMessage cm -> case cm of
Close i closeMsg -> do
hasSentClose <- readIORef $ connectionSentClose conn
unless hasSentClose $ send conn msg
throw $ CloseRequest i closeMsg
Pong _ -> do
connectionOnPong (connectionOptions conn)
receiveDataMessage conn
Ping pl -> do
send conn (ControlMessage (Pong pl))
receiveDataMessage conn
receiveData :: WebSocketsData a => Connection -> IO a
receiveData conn = do
dm <- receiveDataMessage conn
case dm of
Text x -> return (fromLazyByteString x)
Binary x -> return (fromLazyByteString x)
send :: Connection -> Message -> IO ()
send conn msg = do
case msg of
(ControlMessage (Close _ _)) -> writeIORef (connectionSentClose conn) True
_ -> return ()
Streams.write (Just msg) (connectionOut conn)
sendDataMessage :: Connection -> DataMessage -> IO ()
sendDataMessage conn = send conn . DataMessage
sendTextData :: WebSocketsData a => Connection -> a -> IO ()
sendTextData conn = sendDataMessage conn . Text . toLazyByteString
sendBinaryData :: WebSocketsData a => Connection -> a -> IO ()
sendBinaryData conn = sendDataMessage conn . Binary . toLazyByteString
sendClose :: WebSocketsData a => Connection -> a -> IO ()
sendClose conn = sendCloseCode conn 1000
sendCloseCode :: WebSocketsData a => Connection -> Word16 -> a -> IO ()
sendCloseCode conn code =
send conn . ControlMessage . Close code . toLazyByteString
sendPing :: WebSocketsData a => Connection -> a -> IO ()
sendPing conn = send conn . ControlMessage . Ping . toLazyByteString