GraphQLTransportWS: Always return an id upon errors (#1845)

* Always return an id upon errors * According to the server implementation in graphql-transport-ws always close the socket with 4403 on error * Moved over to using zio.Random for id generation
ghostdogpr · Aug 29, 2023 · 85d2352 · 85d2352
1 parent 2268b67
commit 85d2352
Showing 1 changed file with 23 additions and 8 deletions.
diff --git a/interop/tapir/src/main/scala/caliban/interop/tapir/ws/Protocol.scala b/interop/tapir/src/main/scala/caliban/interop/tapir/ws/Protocol.scala
@@ -8,7 +8,7 @@ import caliban.interop.tapir.TapirAdapter.CalibanPipe
 import caliban.interop.tapir.WebSocketHooks
 import zio.stm.TMap
 import zio.stream.{ UStream, ZStream }
-import zio.{ Duration, Promise, Queue, Ref, Schedule, UIO, URIO, ZIO }
+import zio.{ Duration, Promise, Queue, Random, Ref, Schedule, UIO, URIO, ZIO }
 
 sealed trait Protocol {
   def name: String
@@ -76,7 +76,7 @@ object Protocol {
                                case GraphQLWSInput(Ops.ConnectionInit, id, payload)  =>
                                  val before     = ZIO.whenCase((webSocketHooks.beforeInit, payload)) {
                                    case (Some(beforeInit), Some(payload)) =>
-                                     beforeInit(payload).catchAll(e => output.offer(Right(handler.error(id, e))))
+                                     beforeInit(payload).catchAll(_ => output.offer(Left(GraphQLWSClose(4403, "Forbidden"))))
                                  }
                                  val ackPayload = webSocketHooks.onAck.fold[URIO[R, Option[ResponseValue]]](ZIO.none)(_.option)
                                  val response   =
@@ -86,7 +86,7 @@ object Protocol {
                                    afterInit
                                      .catchAllCause(cause =>
                                        ZIO.foreachDiscard(cause.failureOption)(e =>
-                                         output.offer(Right(handler.error(id, e)))
+                                         generateId(id).flatMap(uuid => output.offer(Right(handler.error(uuid, e))))
                                        ) *> output.offer(Left(GraphQLWSClose(4401, "Unauthorized")))
                                      )
                                      .fork
@@ -95,14 +95,20 @@ object Protocol {
                                  before *> response *> ka *> after
                                case GraphQLWSInput(Ops.Pong, id, payload)            =>
                                  ZIO.whenCase(webSocketHooks.onPong -> payload) { case (Some(onPong), Some(payload)) =>
-                                   onPong(payload).catchAll(e => output.offer(Right(handler.error(id, e))))
+                                   onPong(payload).catchAll(e =>
+                                     generateId(id).flatMap(uuid => output.offer(Right(handler.error(uuid, e))))
+                                   )
                                  }
                                case GraphQLWSInput(Ops.Ping, id, payload)            =>
                                  def sendPong(p: Option[ResponseValue]) = output.offer(Right(GraphQLWSOutput(Ops.Pong, id, p)))
 
                                  webSocketHooks.onPing match {
                                    case Some(onPing) =>
-                                     onPing(payload).flatMap(sendPong).catchAll(e => output.offer(Right(handler.error(id, e))))
+                                     onPing(payload)
+                                       .flatMap(sendPong)
+                                       .catchAll(e =>
+                                         generateId(id).flatMap(uuid => output.offer(Right(handler.error(uuid, e))))
+                                       )
                                    case _            => sendPong(None)
                                  }
                                case GraphQLWSInput(Ops.Subscribe, Some(id), payload) =>
@@ -131,7 +137,8 @@ object Protocol {
                                          .unit
                                      )
 
-                                   case None => output.offer(Right(connectionError))
+                                   case None =>
+                                     generateId(None).flatMap(uuid => output.offer(Right(connectionError(uuid))))
                                  }
 
                                  ZIO.ifZIO(ack.get)(continue, output.offer(Left(GraphQLWSClose(4401, "Unauthorized"))))
@@ -140,18 +147,26 @@ object Protocol {
                                case GraphQLWSInput(unsupported, _, _)                =>
                                  output.offer(Left(GraphQLWSClose(4400, s"Unsupported operation: $unsupported")))
                              }.runDrain.interruptible
-                               .catchAll(_ => output.offer(Right(connectionError)))
+                               .catchAll(_ =>
+                                 generateId(None).flatMap(uuid => output.offer(Right(connectionError(Some(uuid.toString())))))
+                               )
                                .ensuring(subscriptions.untrackAll)
                                .provideEnvironment(env)
                                .forkScoped
                            ) *> ZStream.fromQueueWithShutdown(output)
                          }
       } yield pipe
 
-    private val connectionError: GraphQLWSOutput                               = GraphQLWSOutput(Ops.Error, None, None)
+    private def connectionError(id: Option[String]): GraphQLWSOutput           = GraphQLWSOutput(Ops.Error, id, None)
     private def connectionAck(payload: Option[ResponseValue]): GraphQLWSOutput =
       GraphQLWSOutput(Ops.ConnectionAck, None, payload)
 
+    private def generateId(id: Option[String]): ZIO[Any, Nothing, Option[String]] =
+      id match {
+        case Some(_) => ZIO.succeed(id)
+        case None    => Random.nextUUID.map(uuid => Some(uuid.toString))
+      }
+
     private def ping(keepAlive: Option[Duration]): UStream[Either[Nothing, GraphQLWSOutput]] =
       keepAlive match {
         case None           => ZStream.empty