diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,13 @@ +Copyright 2018 Yuji Yamamoto + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,8 @@ +# Skews + +A very quick-and-dirty WebSocket server. Intended for use with unit tests. +YOU SHOULD NOT USE THIS FOR PRODUCTION!! + +## Examples + +See [the spec file](https://github.com/iij-ii/skews/blob/master/test/Network/WebSockets/SkewsSpec.hs). diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/skews.cabal b/skews.cabal new file mode 100644 --- /dev/null +++ b/skews.cabal @@ -0,0 +1,44 @@ +name: skews +version: 0.1.0.0 +synopsis: A very quick-and-dirty WebSocket server. +description: A very quick-and-dirty WebSocket server. Intended for use with unit tests. +homepage: https://github.com/iij-ii/skews#readme +license: Apache-2.0 +license-file: LICENSE +author: Yuji Yamamoto +maintainer: yuji-yamamoto@iij.ad.jp +copyright: 2018 Yuji Yamamoto +category: Web, Test +build-type: Simple +extra-source-files: README.md +cabal-version: >=1.10 + +library + hs-source-dirs: src + exposed-modules: Network.WebSockets.Skews + build-depends: base >= 4.7 && < 5 + , bytestring + , deque + , websockets + default-language: Haskell2010 + +test-suite mock-ws-server-test + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: Spec.hs + other-modules: Network.WebSockets.SkewsSpec + build-depends: base + , skews + , async >= 2.1.1 + , bytestring + , deque + , envy + , hspec + , network + , websockets + ghc-options: -threaded -rtsopts -with-rtsopts=-N + default-language: Haskell2010 + +source-repository head + type: git + location: https://github.com/iij-ii/skews diff --git a/src/Network/WebSockets/Skews.hs b/src/Network/WebSockets/Skews.hs new file mode 100644 --- /dev/null +++ b/src/Network/WebSockets/Skews.hs @@ -0,0 +1,212 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RecordWildCards #-} + +module Network.WebSockets.Skews + ( -- * The server object type and related types. + Server + , Args(..) + , RequestHandler + + -- * Controling the server's lifecycle. + , start + , reinit + , threadId + + -- * Configuring how the server responds/sends WebSocket messages to the client. + , enqueRequestHandler + , enqueResponse + , replaceRequestHandlers + , setDefaultRequestHandler + , setDefaultResponse + , respondWith + , doNothing + , sendToClients + , forgetDefaultRequestHandler + , forgetDefaultResponse + , forgetReceivedRequests + + -- * Checking the server's status. + , listeningPort + , listeningHost + , recentlyReceived + , countConnectedClients + ) where + + +import Control.Applicative ((<|>)) +import Control.Concurrent (ThreadId, forkIO) +import Control.Exception (handleJust) +import Control.Monad (mapM_) +import qualified Data.ByteString.Lazy as B +import qualified Data.IORef as IOR +import Data.Monoid (mempty) +import qualified Deque as Q +import qualified Network.WebSockets as WS + + +-- | The Server object. +-- You can easily configure the behavior (how the server responds with some requests) +-- after creating by 'start' function. +data Server = + Server + { requestHandlerQueue :: !(IOR.IORef (Q.Deque RequestHandler)) + , defaultRequestHandler :: !(IOR.IORef (Maybe RequestHandler)) + , recentlyReceivedRef :: !(IOR.IORef (Q.Deque B.ByteString)) + , clientConnections :: !(IOR.IORef [WS.Connection]) + , threadId :: !ThreadId -- ^ Call 'Control.Concurrent.killThread' to stop the server. + , listeningPort :: !Int + , listeningHost :: !String + } + + +-- | Used to configure the server's behavior when receiving 'WS.DataMessage' as a 'B.ByteString'. +-- If returns @Nothing@, the server does nothing. +type RequestHandler = B.ByteString -> IO (Maybe B.ByteString) + +-- | Maybe often used 'RequestHandler'. Always respond with the given 'WS.Message'. +respondWith :: B.ByteString -> RequestHandler +respondWith = const . return . Just + + +-- | Maybe often used 'RequestHandler'. Do nothing. +doNothing :: RequestHandler +doNothing = const $ return Nothing + + +-- TODO: Get an unused port number automatically +data Args = Args + { host :: String + , portNumber :: Int + } + + +-- | Start the server by the given hostname and port number as 'Args' object. +start :: Args -> IO Server +start (Args listeningHost listeningPort) = do + requestHandlerQueue <- IOR.newIORef mempty + defaultRequestHandler <- IOR.newIORef Nothing + recentlyReceivedRef <- IOR.newIORef mempty + clientConnections <- IOR.newIORef mempty + threadId <- forkIO $ WS.runServer listeningHost listeningPort $ \pc -> do + c <- WS.acceptRequest pc + addClientConnection clientConnections c + + let + loop = do + m <- WS.receive c + case m of + WS.DataMessage _ _ _ dat -> do + mbrh <- deque requestHandlerQueue + dmbrh <- IOR.readIORef defaultRequestHandler + let bs = WS.fromDataMessage dat + case mbrh <|> dmbrh of + Just rh -> maybe (return ()) (WS.sendBinaryData c) =<< rh bs + Nothing -> return () + IOR.modifyIORef recentlyReceivedRef (Q.snoc bs) + loop + WS.ControlMessage (WS.Close _ _) -> + WS.sendClose c ("Bye" :: B.ByteString) + _other -> loop + + ignoreConnectionClosed loop + + return Server {..} + + +addClientConnection :: IOR.IORef [WS.Connection] -> WS.Connection -> IO () +addClientConnection ccsr c = + IOR.atomicModifyIORef' ccsr (\ccs -> (c : ccs, ())) + + +-- | Configure the request handler called when the server receives next. +-- 'RequestHandler's configured with this function and other non-@Default@ functions are "dequeued". +-- So the server responds with the request handler only once. +-- +-- If you need the server to respond always with the same response, +-- use 'setDefaultResponse' and 'setDefaultRequestHandler'. +enqueRequestHandler :: Server -> RequestHandler -> IO () +enqueRequestHandler Server {..} rh = + IOR.atomicModifyIORef' requestHandlerQueue (\q -> (Q.snoc rh q, ())) + + +-- | Configure the response called when the server receives next. +enqueResponse :: Server -> B.ByteString -> IO () +enqueResponse s = enqueRequestHandler s . respondWith + + +-- | Configure the request handler called when no request handlers are queued. +setDefaultRequestHandler :: Server -> RequestHandler -> IO () +setDefaultRequestHandler Server {..} = + IOR.atomicWriteIORef defaultRequestHandler . Just + + +-- | Reset the request handler queue. +replaceRequestHandlers :: Server -> [RequestHandler] -> IO () +replaceRequestHandlers Server {..} = + IOR.atomicWriteIORef requestHandlerQueue . Q.fromList + + +-- | Configure the response called when no request handlers are queued. +setDefaultResponse :: Server -> B.ByteString -> IO () +setDefaultResponse s = setDefaultRequestHandler s . respondWith + + +-- | Delete the default request handler. After calling this function, +-- the 'Server' object doesn't respond to any message if the request handler queue is empty. +forgetDefaultRequestHandler :: Server -> IO () +forgetDefaultRequestHandler Server {..} = + IOR.atomicWriteIORef defaultRequestHandler Nothing + + +-- | Alias for 'forgetDefaultRequestHandler' +forgetDefaultResponse :: Server -> IO () +forgetDefaultResponse = forgetDefaultRequestHandler + + +-- | Forget recently received requests. +forgetReceivedRequests :: Server -> IO () +forgetReceivedRequests Server {..} = + IOR.atomicWriteIORef recentlyReceivedRef mempty + + +-- | Send the given 'Message' immediately to the all connected clients. +sendToClients :: Server -> B.ByteString -> IO () +sendToClients Server {..} msg = + mapM_ (`WS.sendBinaryData` msg) =<< IOR.readIORef clientConnections + + +-- | Close all connections, forget recently received requests, +-- and delete any configured 'RequestHandler's. +reinit :: Server -> IO () +reinit Server {..} = do + IOR.atomicWriteIORef requestHandlerQueue mempty + IOR.atomicWriteIORef defaultRequestHandler Nothing + IOR.atomicWriteIORef recentlyReceivedRef mempty + + mapM_ (ignoreConnectionClosed . (`WS.sendClose` ("Bye" :: B.ByteString))) + =<< IOR.readIORef clientConnections + IOR.atomicWriteIORef clientConnections [] + + +-- | Retrieve any messages sent by the clients. +recentlyReceived :: Server -> IO (Q.Deque B.ByteString) +recentlyReceived Server {..} = IOR.readIORef recentlyReceivedRef + + +deque :: IOR.IORef (Q.Deque a) -> IO (Maybe a) +deque qr = IOR.atomicModifyIORef' qr $ \q -> case Q.uncons q of + Just (x, qLeft) -> (qLeft, Just x) + _ -> (mempty, Nothing) + + +ignoreConnectionClosed :: IO () -> IO () +ignoreConnectionClosed = handleJust selectClosed (const $ return ()) + where + selectClosed :: WS.ConnectionException -> Maybe () + selectClosed WS.ConnectionClosed = Just () + selectClosed _ = Nothing + + +-- | For debugging +countConnectedClients :: Server -> IO Int +countConnectedClients Server {..} = length <$> IOR.readIORef clientConnections diff --git a/test/Network/WebSockets/SkewsSpec.hs b/test/Network/WebSockets/SkewsSpec.hs new file mode 100644 --- /dev/null +++ b/test/Network/WebSockets/SkewsSpec.hs @@ -0,0 +1,97 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Network.WebSockets.SkewsSpec + ( spec + ) where + +import Control.Applicative (empty, (<|>)) +import Control.Concurrent (threadDelay) +import Control.Concurrent.Async (async, replicateConcurrently, + replicateConcurrently_, wait) +import Data.ByteString.Char8 () +import Data.Foldable (toList) +import Data.Traversable (for) +import Deque () +import System.Envy (FromEnv, decodeEnv, env, fromEnv) +import Test.Hspec +import Text.Read (readMaybe) + +import qualified Network.WebSockets as WS +import qualified Network.WebSockets.Skews as Skews + +-- import Debug.Trace + +newtype PortNumber = PortNumber Int deriving (Eq, Show) + +instance FromEnv PortNumber where + fromEnv = do + mpn <- (readMaybe <$> env "SKEWS_TEST_PORT") <|> pure (Just 8614) + PortNumber <$> maybe empty pure mpn + + +spec :: Spec +spec = context "After running Skews.start and starting a client" $ do + Right (PortNumber pn) <- runIO decodeEnv + + let host = "127.0.0.1" + server <- runIO $ Skews.start $ Skews.Args host pn + + let runClient = WS.runClient host pn "/" + + before_ (Skews.reinit server) $ do + it + "after enqueResponses and setDefaultResponse, responds with the given responses and records request" + $ do + let requests = + ["client1", "client2", "client3", "client4", "client5"] + -- Close code 1000: normal closure: https://tools.ietf.org/html/rfc6455#section-7.4 + lastRequest = WS.ControlMessage $ WS.Close 1000 "bye by client" + responses = ["response1", "response2", "response3"] + + defaultResponse = "default response" + Skews.setDefaultResponse server defaultResponse + + mapM_ (Skews.enqueResponse server) responses + + actuallyResponded <- runClient $ \conn -> do + ress <- for requests $ \request -> do + WS.sendBinaryData conn request + WS.receiveData conn + + Skews.forgetDefaultResponse server + WS.send conn lastRequest + {- FIXME: + -- timeout function doesn't kill `WS.receive conn`. + -- So I can't test the exact behaviour of forgetDefaultResponse + resLast <- timeout 1 $ WS.receive conn + return (ress, resLast) + -} + return ress + + toList <$> Skews.recentlyReceived server `shouldReturn` requests + actuallyResponded + `shouldBe` responses + ++ [defaultResponse, defaultResponse] + + it "sendToClients delivers the given message to all connected clients" $ do + let msg = "message" + clientCount = 10 + clientAction = runClient $ \conn -> do + threadDelay $ 100 * 1000 + WS.receiveData conn + + actuallyReceived <- async $ replicateConcurrently clientCount clientAction + threadDelay $ 50 * 1000 + Skews.sendToClients server msg + + wait actuallyReceived `shouldReturn` replicate clientCount msg + + it "can record requests sent concurrently" $ do + let msg = "client message" + clientCount = 10 + clientAction = runClient (`WS.sendBinaryData` msg) + replicateConcurrently_ clientCount clientAction + threadDelay $ 50 * 1000 + toList + <$> Skews.recentlyReceived server + `shouldReturn` replicate clientCount msg diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,1 @@ +{-# OPTIONS_GHC -F -pgmF hspec-discover #-}