server: address recent graphql-ws related bugs

GITHUB_PR_NUMBER: 7730
GITHUB_PR_URL: https://github.com/hasura/graphql-engine/pull/7730

PR-URL: https://github.com/hasura/graphql-engine-mono/pull/2685
Co-authored-by: Sameer Kolhar <6604943+kolharsam@users.noreply.github.com>
GitOrigin-RevId: 55bafd4eb1576e95803350f3ba9c7920a21de037

.gitignore

@@ -1,6 +1,7 @@
 .hasura-dev-python-venv
 server/tests-py/.devsh_version
 npm-debug.log
+dist-newstyle/
 *.temp
 *.DS_Store
 .tern-project

server/src-lib/Hasura/GraphQL/Transport/WSServerApp.hs

@@ -9,7 +9,7 @@ import Control.Concurrent.Async.Lifted.Safe qualified as LA
 import Control.Concurrent.STM qualified as STM
 import Control.Exception.Lifted
 import Control.Monad.Trans.Control qualified as MC
-import Data.Aeson (toJSON)
+import Data.Aeson (object, toJSON, (.=))
 import Data.ByteString.Char8 qualified as B (pack)
 import Data.Environment qualified as Env
 import Data.Text (pack, unpack)
@@ -137,6 +137,7 @@ mkWSActions logger subProtocol =
     keepAliveAction
     getServerMsgType
     mkAcceptRequest
+    fmtErrorMessage
   where
     mkPostExecErrMessageAction wsConn opId execErr =
       sendMsg wsConn $ case subProtocol of
@@ -145,11 +146,11 @@ mkWSActions logger subProtocol =
 
     mkOnErrorMessageAction wsConn err mErrMsg = case subProtocol of
       Apollo -> sendMsg wsConn $ SMConnErr err
-      GraphQLWS -> sendCloseWithMsg logger wsConn (GenericError4400 $ (fromMaybe "" mErrMsg) <> (unpack . unConnErrMsg $ err)) Nothing
+      GraphQLWS -> sendCloseWithMsg logger wsConn (GenericError4400 $ (fromMaybe "" mErrMsg) <> (unpack . unConnErrMsg $ err)) Nothing Nothing
 
     mkConnectionCloseAction wsConn opId errMsg =
       when (subProtocol == GraphQLWS) $
-        sendCloseWithMsg logger wsConn (GenericError4400 errMsg) (Just . SMErr $ ErrorMsg opId $ toJSON (pack errMsg))
+        sendCloseWithMsg logger wsConn (GenericError4400 errMsg) (Just . SMErr $ ErrorMsg opId $ toJSON (pack errMsg)) (Just 1000)
 
     getServerMsgType = case subProtocol of
       Apollo -> SMData
@@ -164,3 +165,7 @@ mkWSActions logger subProtocol =
       WS.defaultAcceptRequest
         { WS.acceptSubprotocol = Just . B.pack . showSubProtocol $ subProtocol
         }
+
+    fmtErrorMessage errMsgs = case subProtocol of
+      Apollo -> object ["errors" .= errMsgs]
+      GraphQLWS -> toJSON errMsgs

server/src-lib/Hasura/GraphQL/Transport/WebSocket.hs

@@ -38,6 +38,7 @@ import Data.String
 import Data.Text qualified as T
 import Data.Text.Encoding qualified as TE
 import Data.Time.Clock qualified as TC
+import Data.Word (Word16)
 import GHC.AssertNF.CPP
 import Hasura.Backends.Postgres.Instances.Transport (runPGMutationTransaction)
 import Hasura.Base.Error
@@ -210,23 +211,36 @@ sendMsg :: (MonadIO m) => WSConn -> ServerMsg -> m ()
 sendMsg wsConn msg =
   liftIO $ WS.sendMsg wsConn $ WS.WSQueueResponse (encodeServerMsg msg) Nothing
 
+-- sendCloseWithMsg closes the websocket server with an error code that can be supplied as (Maybe Word16),
+-- if there is `Nothing`, the server will be closed with an error code derived from ServerErrorCode
 sendCloseWithMsg ::
   (MonadIO m) =>
   L.Logger L.Hasura ->
   WSConn ->
   ServerErrorCode ->
   Maybe ServerMsg ->
+  Maybe Word16 ->
   m ()
-sendCloseWithMsg logger wsConn errCode mErrServerMsg = do
+sendCloseWithMsg logger wsConn errCode mErrServerMsg mCode = do
   case mErrServerMsg of
     Just errServerMsg -> do
       sendMsg wsConn errServerMsg
     Nothing -> pure ()
   logWSEvent logger wsConn EClosed
-  liftIO $ WS.sendClose wsc errMsg
+  liftIO $ WS.sendCloseCode wsc errCloseCode errMsg
   where
     wsc = WS.getRawWebSocketConnection wsConn
     errMsg = encodeServerErrorMsg errCode
+    errCloseCode = fromMaybe (getErrCode errCode) mCode
+    getErrCode :: ServerErrorCode -> Word16
+    getErrCode err = case err of
+      ProtocolError1002 -> 1002
+      GenericError4400 _ -> 4400
+      Unauthorized4401 -> 4401
+      Forbidden4403 -> 4403
+      ConnectionInitTimeout4408 -> 4408
+      NonUniqueSubscription4409 _ -> 4409
+      TooManyRequests4429 -> 4429
 
 sendMsgWithMetadata ::
   (MonadIO m) =>
@@ -643,6 +657,7 @@ onStart env enabledLogTypes serverEnv wsConn (StartMsg opId q) onMessageActions
     sendDataMsg = WS._wsaGetDataMessageType onMessageActions
     closeConnAction = WS._wsaConnectionCloseAction onMessageActions
     postExecErrAction = WS._wsaPostExecErrMessageAction onMessageActions
+    fmtErrorMessage = WS._wsaErrorMsgFormat onMessageActions
 
     getExecStepActionWithActionInfo acc execStep = case execStep of
       E.ExecStepAction _ actionInfo _remoteJoins -> (actionInfo : acc)
@@ -757,7 +772,7 @@ onStart env enabledLogTypes serverEnv wsConn (StartMsg opId q) onMessageActions
       logOpEv (ODQueryErr qErr) (Just reqId) Nothing
       let err = case errRespTy of
             ERTLegacy -> errFn False qErr
-            ERTGraphqlCompliant -> J.object ["errors" J..= [errFn False qErr]]
+            ERTGraphqlCompliant -> fmtErrorMessage [errFn False qErr]
       sendMsg wsConn (SMErr $ ErrorMsg opId err)
 
     sendSuccResp ::

server/src-lib/Hasura/GraphQL/Transport/WebSocket/Protocol.hs

@@ -188,7 +188,7 @@ data ServerMsg
 data ServerErrorCode
   = ProtocolError1002
   | GenericError4400 !String
-  | UnAuthorized4401
+  | Unauthorized4401
   | Forbidden4403
   | ConnectionInitTimeout4408
   | NonUniqueSubscription4409 !OperationId
@@ -199,7 +199,7 @@ encodeServerErrorMsg :: ServerErrorCode -> BL.ByteString
 encodeServerErrorMsg ecode = encJToLBS . encJFromJValue $ case ecode of
   ProtocolError1002 -> packMsg "1002: Protocol Error"
   GenericError4400 msg -> packMsg $ "4400: " <> msg
-  UnAuthorized4401 -> packMsg "4401: Unauthorized"
+  Unauthorized4401 -> packMsg "4401: Unauthorized"
   Forbidden4403 -> packMsg "4403: Forbidden"
   ConnectionInitTimeout4408 -> packMsg "4408: Connection initialisation timeout"
   NonUniqueSubscription4409 opId -> packMsg $ "4409: Subscriber for " <> show opId <> " already exists"

server/src-lib/Hasura/GraphQL/Transport/WebSocket/Server.hs

@@ -236,7 +236,8 @@ data WSActions a = WSActions
     -- after the connection has been successfully established after `connection_init`
     _wsaKeepAliveAction :: !(WSKeepAliveMessageAction a),
     _wsaGetDataMessageType :: !(DataMsg -> ServerMsg),
-    _wsaAcceptRequest :: !WS.AcceptRequest
+    _wsaAcceptRequest :: !WS.AcceptRequest,
+    _wsaErrorMsgFormat :: !([J.Value] -> J.Value)
   }
 
 -- | to be used with `WSOnErrorMessageAction`

server/tests-py/validate.py

@@ -292,6 +292,8 @@ def validate_gql_ws_q(hge_ctx, conf, headers, retry=False, via_subscription=Fals
             assert resp['type'] in ['data', 'error', 'next'], resp
 
     if 'errors' in exp_http_response or 'error' in exp_http_response:
+        if gqlws:
+            resp['payload'] = {'errors':resp['payload']}
         assert resp['type'] in ['data', 'error', 'next'], resp
     else:
         assert resp['type'] == 'data' or resp['type'] == 'next', resp