improving handling of invalid client messages

src/Main/Exceptions.fs

@@ -0,0 +1,4 @@
+namespace GraphQLTransportWS
+
+type InvalidMessageException (explanation : string) =
+  inherit System.Exception(explanation)

src/Main/GraphQLWebsocketMiddleware.fs

@@ -50,7 +50,9 @@ type GraphQLWebSocketMiddleware<'Root>(next : RequestDelegate, applicationLifeti
       task { return "dummySerializedServerMessage" }
 
   let deserializeClientMessage (jsonOptions: JsonOptions) (msg: string) =
-    task { return ConnectionInit None }
+    task {
+      return JsonSerializer.Deserialize<RawMessage>(msg, jsonOptions.SerializerOptions)
+    }
 
   let isSocketOpen (theSocket : WebSocket) =
     not (theSocket.State = WebSocketState.Aborted) &&
@@ -61,7 +63,7 @@ type GraphQLWebSocketMiddleware<'Root>(next : RequestDelegate, applicationLifeti
     not (theSocket.State = WebSocketState.Aborted) &&
     not (theSocket.State = WebSocketState.Closed)
 
-  let receiveMessageViaSocket (cancellationToken : CancellationToken) (jsonOptions: JsonOptions) (executor : Executor<'Root>) (replacements : Map<string, obj>) (socket : WebSocket) =
+  let receiveMessageViaSocket (cancellationToken : CancellationToken) (jsonOptions: JsonOptions) (executor : Executor<'Root>) (replacements : Map<string, obj>) (socket : WebSocket) : Task<RopResult<ClientMessage, ClientMessageProtocolFailure> option> =
     task {
       let buffer = Array.zeroCreate 4096
       let completeMessage = new List<byte>()
@@ -84,8 +86,17 @@ type GraphQLWebSocketMiddleware<'Root>(next : RequestDelegate, applicationLifeti
       if String.IsNullOrWhiteSpace message then
         return None
       else
-        let! deserializedMsg = deserializeClientMessage jsonOptions message
-        return Some deserializedMsg
+        try
+          let! deserializedRawMsg = deserializeClientMessage jsonOptions message
+          return
+            deserializedRawMsg
+            |> MessageMapping.toClientMessage executor
+            |> Some
+        with :? JsonException as e ->
+          printfn "%s" (e.ToString())
+          return
+            (MessageMapping.invalidMsg <| "invalid json in client message"
+            |> Some)
     }
 
   let sendMessageViaSocket (cancellationToken : CancellationToken) (jsonOptions: JsonOptions) (socket : WebSocket) (message : ServerMessage) =
@@ -167,6 +178,15 @@ type GraphQLWebSocketMiddleware<'Root>(next : RequestDelegate, applicationLifeti
     let logMsgWithIdReceived id msgAsStr =
         printfn "%s (id: %s)" msgAsStr id
 
+    let tryToGracefullyCloseSocket theSocket =
+      task {
+        if theSocket |> canCloseSocket
+        then
+          do! theSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, "NormalClosure", new CancellationToken())
+        else
+          ()
+      }
+
     // <--------------
     // <-- Helpers --|
     // <--------------
@@ -183,15 +203,20 @@ type GraphQLWebSocketMiddleware<'Root>(next : RequestDelegate, applicationLifeti
                 printfn "Warn: websocket socket received empty message! (socket state = %A)" socket.State
             | Some msg ->
                 match msg with
-                | ConnectionInit p ->
+                | Failure failureMsgs ->
+                    "InvalidMessage" |> logMsgReceivedWithOptionalPayload None
+                    match failureMsgs |> List.head with
+                    | InvalidMessage (code, explanation) ->
+                      do! socket.CloseAsync(enum code, explanation, cancellationToken)
+                | Success (ConnectionInit p, _) ->
                     "ConnectionInit" |> logMsgReceivedWithOptionalPayload p
                     do! ConnectionAck |> safe_Send
-                | ClientPing p ->
+                | Success (ClientPing p, _) ->
                     "ClientPing" |> logMsgReceivedWithOptionalPayload p
                     do! ServerPong None |> safe_Send
-                | ClientPong p ->
+                | Success (ClientPong p, _) ->
                     "ClientPong" |> logMsgReceivedWithOptionalPayload p
-                | Subscribe (id, query) ->
+                | Success (Subscribe (id, query), _) ->
                     "Subscribe" |> logMsgWithIdReceived id
                     let! planExecutionResult =
                         Async.StartAsTask (
@@ -200,22 +225,19 @@ type GraphQLWebSocketMiddleware<'Root>(next : RequestDelegate, applicationLifeti
                         )
                     do! planExecutionResult
                         |> safe_ApplyPlanExecutionResult id
-                | ClientComplete id ->
+                | Success (ClientComplete id, _) ->
                     "ClientComplete" |> logMsgWithIdReceived id
                     id |> GraphQLSubscriptionsManagement.removeSubscription
                     do! Complete id |> safe_Send
-                | InvalidMessage explanation ->
-                    "InvalidMessage" |> logMsgReceivedWithOptionalPayload None
-                    do! socket.CloseAsync(enum CustomWebSocketStatus.invalidMessage, explanation, cancellationToken)
         printfn "Leaving graphql-ws connection loop..."
-        if socket |> canCloseSocket
-        then
-          do! socket.CloseAsync(WebSocketCloseStatus.NormalClosure, "NormalClosure", new CancellationToken())
-        else
-          ()
-      with // TODO: MAKE A PROPER GRAPHQL ERROR HANDLING!
+        do! socket |> tryToGracefullyCloseSocket
+      with
       | ex ->
         printfn "Unexpected exception \"%s\" in GraphQLWebsocketMiddleware (handleMessages). More:\n%s" (ex.GetType().Name) (ex.ToString())
+        // at this point, only something really weird must have happened.
+        // In order to avoid faulty state scenarios and unimagined damages,
+        // just close the socket without further ado.
+        do! socket |> tryToGracefullyCloseSocket
     }
 
   let waitForConnectionInit (jsonOptions : JsonOptions) (schemaExecutor : Executor<'Root>) (replacements : Map<string, obj>) (connectionInitTimeoutInMs : int) (socket : WebSocket) : Task<RopResult<unit, string>> =

src/Main/Main.fsproj

@@ -12,6 +12,7 @@
   <ItemGroup>
     <Compile Include="Rop/Rop.fs" />
     <Compile Include="Rop/RopAsync.fs" />
+    <Compile Include="Exceptions.fs" />
     <Compile Include="Messages.fs" />
     <Compile Include="MessageMapping.fs" />
     <Compile Include="GraphQLWebsocketMiddlewareOptions.fs" />

src/Main/MessageMapping.fs

@@ -2,21 +2,32 @@ namespace GraphQLTransportWS
 
 module MessageMapping =
   open FSharp.Data.GraphQL
+  open Rop
 
-  let requireId (raw : RawMessage) : string =
+  /// From the spec: "Receiving a message of a type or format which is not specified in this document will result in an immediate socket closure with the event 4400: <error-message>.
+  /// The <error-message> can be vaguely descriptive on why the received message is invalid."
+  let invalidMsg (explanation : string) =
+    InvalidMessage (4400, explanation)
+    |> fail
+
+  let private requireId (raw : RawMessage) : RopResult<string, ClientMessageProtocolFailure> =
     match raw.Id with
-    | Some s -> s
-    | None -> failwith "property \"id\" is required but was not there"
+    | Some s -> succeed s
+    | None -> invalidMsg <| "property \"id\" is required for this message but was not present."
 
-  let requirePayloadToBeAnOptionalString (payload : RawPayload option) : string option =
+  let private requirePayloadToBeAnOptionalString (payload : RawPayload option) : RopResult<string option, ClientMessageProtocolFailure> =
     match payload with
     | Some p  ->
       match p with
-      | StringPayload strPayload -> Some strPayload
-      | _ -> failwith "payload was expected to be a string, but it wasn't"
-    | None -> None
+      | StringPayload strPayload ->
+        Some strPayload
+        |> succeed
+      | SubscribePayload _ ->
+        invalidMsg <| "for this message, payload was expected to be an optional string, but it was a \"subscribe\" payload instead."
+    | None ->
+      succeed None
 
-  let requireSubscribePayload (executor : Executor<'a>) (payload : RawPayload option) : GraphQLQuery =
+  let private requireSubscribePayload (executor : Executor<'a>) (payload : RawPayload option) : RopResult<GraphQLQuery, ClientMessageProtocolFailure> =
     match payload with
     | Some p ->
       match p with
@@ -25,29 +36,44 @@ module MessageMapping =
         | Some query ->
           { ExecutionPlan = executor.CreateExecutionPlan(query)
             Variables = Map.empty }
+          |> succeed
         | None ->
-          failwith "there was no query in subscribe message!"
+          invalidMsg <| sprintf "there was no query in the client's subscribe message!"
       | _ ->
-        failwith "payload was expected to be a subscribe payload object, but it wasn't."
+        invalidMsg <| "for this message, payload was expected to be a \"subscribe\" payload object, but it wasn't."
     | None ->
-      failwith "payload is required for this message, but none was available"
+      invalidMsg <| "payload is required for this message, but none was present."
 
-  let toClientMessage (executor : Executor<'a>) (raw : RawMessage) : ClientMessage =
+  let toClientMessage (executor : Executor<'a>) (raw : RawMessage) : RopResult<ClientMessage, ClientMessageProtocolFailure> =
     match raw.Type with
     | None ->
-      failwithf "property \"type\" was not found in the client message"
+      invalidMsg <| sprintf "message type was not specified by client."
     | Some "connection_init" ->
-      ConnectionInit (raw.Payload |> requirePayloadToBeAnOptionalString)
+      raw.Payload
+      |> requirePayloadToBeAnOptionalString
+      |> mapR ConnectionInit
     | Some "ping" ->
-      ClientPing (raw.Payload |> requirePayloadToBeAnOptionalString)
+      raw.Payload
+      |> requirePayloadToBeAnOptionalString
+      |> mapR ClientPing
     | Some "pong" ->
-      ClientPong (raw.Payload |> requirePayloadToBeAnOptionalString)
+      raw.Payload
+      |> requirePayloadToBeAnOptionalString
+      |> mapR ClientPong
     | Some "complete" ->
-      ClientComplete (raw |> requireId)
+      raw
+      |> requireId
+      |> mapR ClientComplete
     | Some "subscribe" ->
-      let id = raw |> requireId
-      let payload = raw.Payload |> requireSubscribePayload executor
-      Subscribe (id, payload)
+      raw
+      |> requireId
+      |> bindR
+        (fun id ->
+          raw.Payload
+          |> requireSubscribePayload executor
+          |> mapR (fun payload -> (id, payload))
+        )
+      |> mapR Subscribe
     | Some other ->
-      failwithf "type \"%s\" is not supported as a client message type" other
+      invalidMsg <| sprintf "invalid type \"%s\" specified by client." other
 

src/Main/Messages.fs

@@ -28,7 +28,9 @@ type ClientMessage =
     | ClientPong of payload: string option
     | Subscribe of id: string * query: GraphQLQuery
     | ClientComplete of id: string
-    | InvalidMessage of explanation: string
+
+type ClientMessageProtocolFailure =
+    | InvalidMessage of code: int * explanation: string
 
 type ServerMessage =
     | ConnectionAck

src/Main/RawMessageConverter.fs

@@ -8,6 +8,9 @@ open System.Text.Json.Serialization
 type RawMessageConverter() =
   inherit JsonConverter<RawMessage>()
 
+  let raiseInvalidMsg explanation =
+    raise <| InvalidMessageException explanation
+
   let getOptionalString (reader : byref<Utf8JsonReader>) =
     if reader.TokenType.Equals(JsonTokenType.Null) then
       None
@@ -18,7 +21,7 @@ type RawMessageConverter() =
     if reader.Read() then
       getOptionalString(&reader)
     else
-      failwithf "was expecting a value for property \"%s\"" propertyName
+      raiseInvalidMsg <| sprintf "was expecting a value for property \"%s\"" propertyName
 
   let readSubscribePayload (reader : byref<Utf8JsonReader>) : RawSubscribePayload =
     let mutable operationName : string option = None
@@ -36,7 +39,7 @@ type RawMessageConverter() =
       | "extensions" ->
         extensions <- readPropertyValueAsAString "extensions" &reader
       | other ->
-        failwithf "unexpected property \"%s\" in payload object" other
+        raiseInvalidMsg <| sprintf "unexpected property \"%s\" in payload object" other
     { OperationName = operationName
       Query = query
       Variables = variables
@@ -51,11 +54,11 @@ type RawMessageConverter() =
         SubscribePayload (readSubscribePayload &reader)
         |> Some
       elif reader.TokenType.Equals(JsonTokenType.Null) then
-        failwith "was expecting a value for property \"payload\""
+        raiseInvalidMsg <| "was expecting a value for property \"payload\""
       else
-        failwith "Not implemented yet. Uh-oh, this is a bug."
+        raiseInvalidMsg <| sprintf "payload is a \"%A\", which is not supported" reader.TokenType
     else
-      failwith "was expecting a value for property \"payload\""
+      raiseInvalidMsg <| "was expecting a value for property \"payload\""
 
   override __.Read(reader : byref<Utf8JsonReader>, typeToConvert: Type, options: JsonSerializerOptions) : RawMessage =
     if not (reader.TokenType.Equals(JsonTokenType.StartObject))
@@ -73,7 +76,7 @@ type RawMessageConverter() =
         | "payload" ->
           payload <- readPayload &reader
         | other ->
-          failwithf "unknown property \"%s\"" other
+          raiseInvalidMsg <| sprintf "unknown property \"%s\"" other
       { Id = id
         Type = theType
         Payload = payload }