[feat] websockets

Allow the user to set up an lsp server that talks to the user using
the websocket protocol instead of standard IO and do enough refactoring
to make the setup with websockets possible

Author: MangoIV

lsp/lsp.cabal

@@ -79,6 +79,7 @@ library
     , unliftio              ^>=0.2
     , unliftio-core         ^>=0.2
     , unordered-containers  ^>=0.2
+    , websockets            ^>=0.13
 
 executable lsp-demo-reactor-server
   import:             warnings

lsp/src/Language/LSP/Server/Control.hs

@@ -1,18 +1,35 @@
+{-# LANGUAGE DerivingStrategies #-}
 {-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE RecordWildCards #-}
 
 module Language.LSP.Server.Control (
   -- * Running
-  runServer,
   runServerWith,
-  runServerWithHandles,
+  runServerWithConfig,
+  ServerConfig (..),
   LspServerLog (..),
+
+  -- ** Using standard 'IO' 'Handle's
+  runServer,
+
+  -- ** Using 'Handle's
+  runServerWithHandles,
+  prependHeader,
+  parseHeaders,
+
+  -- **  Using websockets
+  WebsocketConfig (..),
+  withWebsocket,
+  withWebsocketRunServer,
 ) where
 
 import Colog.Core (LogAction (..), Severity (..), WithSeverity (..), (<&))
 import Colog.Core qualified as L
 import Control.Applicative ((<|>))
 import Control.Concurrent
+import Control.Concurrent.Async
 import Control.Concurrent.STM.TChan
+import Control.Exception (finally)
 import Control.Monad
 import Control.Monad.IO.Class
 import Control.Monad.STM
@@ -31,6 +48,7 @@ import Language.LSP.Protocol.Message
 import Language.LSP.Server.Core
 import Language.LSP.Server.Processing qualified as Processing
 import Language.LSP.VFS
+import Network.WebSockets qualified as WS
 import Prettyprinter
 import System.IO
 
@@ -42,6 +60,7 @@ data LspServerLog
   | Starting
   | ParsedMsg T.Text
   | SendMsg TL.Text
+  | WebsocketLog WebsocketLog
   deriving (Show)
 
 instance Pretty LspServerLog where
@@ -60,6 +79,7 @@ instance Pretty LspServerLog where
   pretty Starting = "Starting server"
   pretty (ParsedMsg msg) = "---> " <> pretty msg
   pretty (SendMsg msg) = "<--2-- " <> pretty msg
+  pretty (WebsocketLog msg) = "Websocket:" <+> pretty msg
 
 -- ---------------------------------------------------------------------
 
@@ -71,18 +91,20 @@ instance Pretty LspServerLog where
 runServer :: forall config. ServerDefinition config -> IO Int
 runServer =
   runServerWithHandles
-    ioLogger
-    lspLogger
+    defaultIOLogger
+    defaultLspLogger
     stdin
     stdout
+
+defaultIOLogger :: LogAction IO (WithSeverity LspServerLog)
+defaultIOLogger = L.cmap (show . prettyMsg) L.logStringStderr
  where
   prettyMsg l = "[" <> viaShow (L.getSeverity l) <> "] " <> pretty (L.getMsg l)
-  ioLogger :: LogAction IO (WithSeverity LspServerLog)
-  ioLogger = L.cmap (show . prettyMsg) L.logStringStderr
-  lspLogger :: LogAction (LspM config) (WithSeverity LspServerLog)
-  lspLogger =
-    let clientLogger = L.cmap (fmap (T.pack . show . pretty)) defaultClientLogger
-     in clientLogger <> L.hoistLogAction liftIO ioLogger
+
+defaultLspLogger :: LogAction (LspM config) (WithSeverity LspServerLog)
+defaultLspLogger =
+  let clientLogger = L.cmap (fmap (T.pack . show . pretty)) defaultClientLogger
+   in clientLogger <> L.hoistLogAction liftIO defaultIOLogger
 
 {- | Starts a language server over the specified handles.
  This function will return once the @exit@ notification is received.
@@ -116,27 +138,54 @@ runServerWithHandles ioLogger logger hin hout serverDefinition = do
 
 {- | Starts listening and sending requests and responses
  using the specified I/O.
+
+ Assumes that the client sends (and wants to receive) the Content-Length
+ header. If you do not want this to be the case, use 'runServerWithConfig'
 -}
 runServerWith ::
   -- | The logger to use outside the main body of the server where we can't assume the ability to send messages.
   LogAction IO (WithSeverity LspServerLog) ->
   -- | The logger to use once the server has started and can successfully send messages.
   LogAction (LspM config) (WithSeverity LspServerLog) ->
   -- | Client input.
-  IO BS.ByteString ->
+  IO BS.StrictByteString ->
   -- | Function to provide output to.
-  (BSL.ByteString -> IO ()) ->
+  (BSL.LazyByteString -> IO ()) ->
   ServerDefinition config ->
   IO Int -- exit code
-runServerWith ioLogger logger clientIn clientOut serverDefinition = do
+runServerWith ioLogger lspLogger inwards outwards =
+  runServerWithConfig ServerConfig{prepareOutwards = prependHeader, parseInwards = parseHeaders, ..}
+
+-- ---------------------------------------------------------------------
+
+data ServerConfig config = ServerConfig
+  { ioLogger :: LogAction IO (WithSeverity LspServerLog)
+  -- ^ The logger to use outside the main body of the server where we can't assume the ability to send messages.
+  , lspLogger :: LogAction (LspM config) (WithSeverity LspServerLog)
+  -- ^ The logger to use once the server has started and can successfully send messages.
+  , inwards :: IO BS.StrictByteString
+  -- ^ Client input.
+  , outwards :: BSL.LazyByteString -> IO ()
+  -- ^ Function to provide output to.
+  , prepareOutwards :: BSL.LazyByteString -> BSL.LazyByteString
+  -- ^ how to prepare an outgoing response for sending. This can be used, to e.g. prepend the Content-Length header, c.f. 'prependHeader'
+  , parseInwards :: Attoparsec.Parser BS.StrictByteString
+  -- ^ how to parse the input. This can be used to consume the Content-Length and Content-Type headers, c.f. 'parseHeaders'
+  }
+
+runServerWithConfig ::
+  ServerConfig config ->
+  ServerDefinition config ->
+  IO Int
+runServerWithConfig ServerConfig{..} serverDefinition = do
   ioLogger <& Starting `WithSeverity` Info
 
   cout <- atomically newTChan :: IO (TChan J.Value)
-  _rhpid <- forkIO $ sendServer ioLogger cout clientOut
+  _rhpid <- forkIO $ sendServer ioLogger cout outwards prepareOutwards
 
   let sendMsg msg = atomically $ writeTChan cout $ J.toJSON msg
 
-  ioLoop ioLogger logger clientIn serverDefinition emptyVFS sendMsg
+  ioLoop ioLogger lspLogger inwards parseInwards serverDefinition emptyVFS sendMsg
 
   return 1
 
@@ -146,12 +195,13 @@ ioLoop ::
   forall config.
   LogAction IO (WithSeverity LspServerLog) ->
   LogAction (LspM config) (WithSeverity LspServerLog) ->
-  IO BS.ByteString ->
+  IO BS.StrictByteString ->
+  Attoparsec.Parser BS.StrictByteString ->
   ServerDefinition config ->
   VFS ->
   (FromServerMessage -> IO ()) ->
   IO ()
-ioLoop ioLogger logger clientIn serverDefinition vfs sendMsg = do
+ioLoop ioLogger logger clientIn parser serverDefinition vfs sendMsg = do
   minitialize <- parseOne ioLogger clientIn (parse parser "")
   case minitialize of
     Nothing -> pure ()
@@ -167,7 +217,7 @@ ioLoop ioLogger logger clientIn serverDefinition vfs sendMsg = do
   pioLogger = L.cmap (fmap LspProcessingLog) ioLogger
   pLogger = L.cmap (fmap LspProcessingLog) logger
 
-  loop :: Result BS.ByteString -> LspM config ()
+  loop :: Result BS.StrictByteString -> LspM config ()
   loop = go
    where
     go r = do
@@ -178,30 +228,11 @@ ioLoop ioLogger logger clientIn serverDefinition vfs sendMsg = do
           Processing.processMessage pLogger $ BSL.fromStrict msg
           go (parse parser remainder)
 
-  parser = do
-    try contentType <|> return ()
-    len <- contentLength
-    try contentType <|> return ()
-    _ <- string _ONE_CRLF
-    Attoparsec.take len
-
-  contentLength = do
-    _ <- string "Content-Length: "
-    len <- decimal
-    _ <- string _ONE_CRLF
-    return len
-
-  contentType = do
-    _ <- string "Content-Type: "
-    skipWhile (/= '\r')
-    _ <- string _ONE_CRLF
-    return ()
-
 parseOne ::
   MonadIO m =>
   LogAction m (WithSeverity LspServerLog) ->
-  IO BS.ByteString ->
-  Result BS.ByteString ->
+  IO BS.StrictByteString ->
+  Result BS.StrictByteString ->
   m (Maybe (BS.ByteString, BS.ByteString))
 parseOne logger clientIn = go
  where
@@ -223,30 +254,155 @@ parseOne logger clientIn = go
 
 -- ---------------------------------------------------------------------
 
+data WebsocketLog
+  = WebsocketShutDown
+  | WebsocketNewConnection
+  | WebsocketConnectionClosed
+  | WebsocketPing
+  | WebsocketStarted
+  | WebsocketIncomingRequest
+  | WebsocketOutgoingResponse
+  deriving stock (Show)
+
+instance Pretty WebsocketLog where
+  pretty l = case l of
+    WebsocketPing -> "Ping"
+    WebsocketStarted -> "Started Server, waiting for connections"
+    WebsocketShutDown -> "Shut down server"
+    WebsocketNewConnection -> "New connection established"
+    WebsocketIncomingRequest -> "Received request"
+    WebsocketConnectionClosed -> "Closed connection to client"
+    WebsocketOutgoingResponse -> "Sent response"
+
+-- | 'host' and 'port' of the websocket server to set up
+data WebsocketConfig = WebsocketConfig
+  { host :: !String
+  -- ^ the host of the websocket server, e.g. @"localhost"@
+  , port :: !Int
+  -- ^ the port of the websocket server, e.g. @8080@
+  }
+
+-- | Set up a websocket server, then call call the continuation (in our case this corresponds to the language server) after accepting a connection
+withWebsocket ::
+  -- | The logger
+  LogAction IO (WithSeverity LspServerLog) ->
+  -- | The configuration of the websocket server
+  WebsocketConfig ->
+  -- | invoke the lsp server, passing communication functions
+  (IO BS.StrictByteString -> (BSL.LazyByteString -> IO ()) -> IO r) ->
+  IO ()
+withWebsocket logger conf startLspServer = do
+  let wsLogger = L.cmap (fmap WebsocketLog) logger
+
+  WS.runServer (host conf) (port conf) $ \pending -> do
+    conn <- WS.acceptRequest pending
+    wsLogger <& WebsocketNewConnection `WithSeverity` Debug
+
+    outChan <- newChan
+    inChan <- newChan
+
+    let inwards = readChan inChan
+        outwards = writeChan outChan
+
+    WS.withPingThread conn 30 (wsLogger <& WebsocketPing `WithSeverity` Debug) $ do
+      withAsync (startLspServer inwards outwards) $ \_lspAsync ->
+        race_
+          ( forever $ do
+              msg <- readChan outChan
+              wsLogger <& WebsocketOutgoingResponse `WithSeverity` Debug
+              WS.sendTextData conn msg
+          )
+          ( forever $ do
+              msg <- WS.receiveData conn
+              wsLogger <& WebsocketIncomingRequest `WithSeverity` Debug
+              writeChan inChan msg
+          )
+          `finally` do
+            wsLogger <& WebsocketConnectionClosed `WithSeverity` Debug
+
+withWebsocketRunServer ::
+  -- | Configuration for the websocket
+  WebsocketConfig ->
+  -- | How to set up a new 'ServerDefinition' for a specific configuration. z
+  --   This is passed as CPS'd 'IO' to allow for setting (- and cleaning) up
+  --   a server per websocket connection
+  ((ServerDefinition config -> IO Int) -> IO Int) ->
+  -- | The 'IO' logger
+  LogAction IO (WithSeverity LspServerLog) ->
+  -- | The logger that logs in 'LspM' to the client
+  LogAction (LspM config) (WithSeverity LspServerLog) ->
+  IO ()
+withWebsocketRunServer wsConf withLspDefinition ioLogger lspLogger =
+  withWebsocket ioLogger wsConf $ \inwards outwards -> do
+    withLspDefinition $ \lspDefinition ->
+      runServerWithConfig
+        ServerConfig
+          { ioLogger
+          , lspLogger
+          , inwards
+          , outwards
+          , -- NOTE: if you run the language server on websockets, you do not
+            -- need to prepend headers to requests and responses, because
+            -- the chunking is already handled by the websocket, i.e. there's
+            -- no situation where the client or the server has to rely on input/
+            -- output chunking
+            prepareOutwards = id
+          , parseInwards = Attoparsec.takeByteString
+          }
+        lspDefinition
+
+-- ---------------------------------------------------------------------
+
 -- | Simple server to make sure all output is serialised
-sendServer :: LogAction IO (WithSeverity LspServerLog) -> TChan J.Value -> (BSL.ByteString -> IO ()) -> IO ()
-sendServer _logger msgChan clientOut = do
+sendServer :: LogAction IO (WithSeverity LspServerLog) -> TChan J.Value -> (BSL.LazyByteString -> IO ()) -> (BSL.LazyByteString -> BSL.LazyByteString) -> IO ()
+sendServer _logger msgChan clientOut prepareMessage = do
   forever $ do
     msg <- atomically $ readTChan msgChan
 
     -- We need to make sure we only send over the content of the message,
     -- and no other tags/wrapper stuff
     let str = J.encode msg
-
-    let out =
-          BSL.concat
-            [ TL.encodeUtf8 $ TL.pack $ "Content-Length: " ++ show (BSL.length str)
-            , BSL.fromStrict _TWO_CRLF
-            , str
-            ]
+    let out = prepareMessage str
 
     clientOut out
 
 -- TODO: figure out how to re-enable
 -- This can lead to infinite recursion in logging, see https://github.com/haskell/lsp/issues/447
 -- logger <& SendMsg (TL.decodeUtf8 str) `WithSeverity` Debug
 
-_ONE_CRLF :: BS.ByteString
+-- | prepend a Content-Length header to the given message
+prependHeader :: BSL.LazyByteString -> BSL.LazyByteString
+prependHeader str =
+  BSL.concat
+    [ TL.encodeUtf8 $ TL.pack $ "Content-Length: " ++ show (BSL.length str)
+    , BSL.fromStrict _TWO_CRLF
+    , str
+    ]
+
+{- | parse Content-Length and Content-Type headers and then consume
+  input with length of the Content-Length
+-}
+parseHeaders :: Attoparsec.Parser BS.StrictByteString
+parseHeaders = do
+  try contentType <|> return ()
+  len <- contentLength
+  try contentType <|> return ()
+  _ <- string _ONE_CRLF
+  Attoparsec.take len
+ where
+  contentLength = do
+    _ <- string "Content-Length: "
+    len <- decimal
+    _ <- string _ONE_CRLF
+    return len
+
+  contentType = do
+    _ <- string "Content-Type: "
+    skipWhile (/= '\r')
+    _ <- string _ONE_CRLF
+    return ()
+
+_ONE_CRLF :: BS.StrictByteString
 _ONE_CRLF = "\r\n"
-_TWO_CRLF :: BS.ByteString
+_TWO_CRLF :: BS.StrictByteString
 _TWO_CRLF = "\r\n\r\n"