Add support for setting up language servers to use websockets

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"