Improve handling of close frames in web sockets

Author: adamw

build.sbt

@@ -8,7 +8,7 @@ val scala3_M2 = "3.0.0-M2"
 val scala3_M3 = "3.0.0-M3"
 val scala3 = List(scala3_M2, scala3_M3)
 
-val sttpModelVersion = "1.2.0-RC11"
+val sttpModelVersion = "1.2.0-RC12"
 
 val scalaTestVersion = "3.2.3"
 val scalaNativeTestInterfaceVersion = "0.4.0-M2"
@@ -74,7 +74,7 @@ lazy val projectAggregates: Seq[ProjectReference] = if (sys.env.isDefinedAt("STT
   println("[info] STTP_NATIVE *not* defined, *not* including sttp-native in the aggregate projects")
   scala2.flatMap(v => List[ProjectReference](core.js(v), ws.js(v))) ++
     scala2.flatMap(v => List[ProjectReference](core.jvm(v), ws.jvm(v), fs2.jvm(v), monix.jvm(v))) ++
-    scala3.flatMap(v => List[ProjectReference](core.jvm(v), ws.jvm(v), zio.jvm(v))) ++
+    scala3.flatMap(v => List[ProjectReference](core.jvm(v), ws.jvm(v))) ++ // TODO , zio.jvm(v)
     List[ProjectReference](
       akka.jvm(scala2_12),
       akka.jvm(scala2_13),

ws/src/main/scala/sttp/ws/WebSocket.scala

@@ -4,79 +4,113 @@ import sttp.model.Headers
 import sttp.monad.MonadError
 import sttp.monad.syntax._
 
-/** The `send` and `receive` methods may result in a failed effect, with either one of [[WebSocketException]]
-  * exceptions, or a backend-specific exception.
+/** The `send*` and `receive*` methods may result in a failed effect, with either one of [[WebSocketException]]
+  * exceptions, or a backend-specific exception. Specifically, they will fail with [[WebSocketClosed]] if the
+  * web socket is closed.
+  *
+  * See the `either` and `eitherClose` method to lift web socket closed events to the value level.
   */
 trait WebSocket[F[_]] {
 
-  /** After receiving a close frame, no further interactions with the web socket should happen. Subsequent invocations
-    * of `receive`, as well as `send`, will fail with the [[WebSocketClosed]] exception.
+  /** Receive the next frame from the web socket. This can be a data frame, or a control frame including
+    * [[WebSocketFrame.Close]]. After receiving a close frame, no further interactions with the web socket should
+    * happen.
+    *
+    * However, not all implementations expose the close frame, and web sockets might also get closed without the proper
+    * close frame exchange. In such cases, as well as when invoking `receive`/`send` after receiving a close frame,
+    * this effect will fail with the [[WebSocketClosed]] exception.
     */
   def receive(): F[WebSocketFrame]
   def send(f: WebSocketFrame, isContinuation: Boolean = false): F[Unit]
   def isOpen(): F[Boolean]
 
   /** Receive a single data frame, ignoring others. The frame might be a fragment.
+    * Will fail with [[WebSocketClosed]] if the web socket is closed, or if a close frame is received.
     * @param pongOnPing Should a [[WebSocketFrame.Pong]] be sent when a [[WebSocketFrame.Ping]] is received.
     */
-  def receiveDataFrame(pongOnPing: Boolean = true): F[Either[WebSocketFrame.Close, WebSocketFrame.Data[_]]] =
+  def receiveDataFrame(pongOnPing: Boolean = true): F[WebSocketFrame.Data[_]] =
     receive().flatMap {
-      case close: WebSocketFrame.Close => (Left(close): Either[WebSocketFrame.Close, WebSocketFrame.Data[_]]).unit
-      case d: WebSocketFrame.Data[_]   => (Right(d): Either[WebSocketFrame.Close, WebSocketFrame.Data[_]]).unit
+      case close: WebSocketFrame.Close => monad.error(WebSocketClosed(Some(close)))
+      case d: WebSocketFrame.Data[_]   => monad.unit(d)
       case WebSocketFrame.Ping(payload) if pongOnPing =>
         send(WebSocketFrame.Pong(payload)).flatMap(_ => receiveDataFrame(pongOnPing))
       case _ => receiveDataFrame(pongOnPing)
     }
 
   /** Receive a single text data frame, ignoring others. The frame might be a fragment. To receive whole messages,
     * use [[receiveText]].
+    * Will fail with [[WebSocketClosed]] if the web socket is closed, or if a close frame is received.
     * @param pongOnPing Should a [[WebSocketFrame.Pong]] be sent when a [[WebSocketFrame.Ping]] is received.
     */
-  def receiveTextFrame(pongOnPing: Boolean = true): F[Either[WebSocketFrame.Close, WebSocketFrame.Text]] =
+  def receiveTextFrame(pongOnPing: Boolean = true): F[WebSocketFrame.Text] =
     receiveDataFrame(pongOnPing).flatMap {
-      case Left(close)                   => (Left(close): Either[WebSocketFrame.Close, WebSocketFrame.Text]).unit
-      case Right(t: WebSocketFrame.Text) => (Right(t): Either[WebSocketFrame.Close, WebSocketFrame.Text]).unit
-      case _                             => receiveTextFrame(pongOnPing)
+      case t: WebSocketFrame.Text => t.unit
+      case _                      => receiveTextFrame(pongOnPing)
     }
 
   /** Receive a single binary data frame, ignoring others. The frame might be a fragment. To receive whole messages,
     * use [[receiveBinary]].
+    * Will fail with [[WebSocketClosed]] if the web socket is closed, or if a close frame is received.
     * @param pongOnPing Should a [[WebSocketFrame.Pong]] be sent when a [[WebSocketFrame.Ping]] is received.
     */
-  def receiveBinaryFrame(pongOnPing: Boolean = true): F[Either[WebSocketFrame.Close, WebSocketFrame.Binary]] =
+  def receiveBinaryFrame(pongOnPing: Boolean = true): F[WebSocketFrame.Binary] =
     receiveDataFrame(pongOnPing).flatMap {
-      case Left(close)                     => (Left(close): Either[WebSocketFrame.Close, WebSocketFrame.Binary]).unit
-      case Right(t: WebSocketFrame.Binary) => (Right(t): Either[WebSocketFrame.Close, WebSocketFrame.Binary]).unit
-      case _                               => receiveBinaryFrame(pongOnPing)
+      case t: WebSocketFrame.Binary => t.unit
+      case _                        => receiveBinaryFrame(pongOnPing)
     }
 
   /** Receive a single text message (which might come from multiple, fragmented frames).
     * Ignores non-text frames and returns combined results.
+    * Will fail with [[WebSocketClosed]] if the web socket is closed, or if a close frame is received.
     * @param pongOnPing Should a [[WebSocketFrame.Pong]] be sent when a [[WebSocketFrame.Ping]] is received.
     */
-  def receiveText(pongOnPing: Boolean = true): F[Either[WebSocketFrame.Close, String]] =
+  def receiveText(pongOnPing: Boolean = true): F[String] =
     receiveConcat(() => receiveTextFrame(pongOnPing), _ + _)
 
   /** Receive a single binary message (which might come from multiple, fragmented frames).
     * Ignores non-binary frames and returns combined results.
+    * Will fail with [[WebSocketClosed]] if the web socket is closed, or if a close frame is received.
     * @param pongOnPing Should a [[WebSocketFrame.Pong]] be sent when a [[WebSocketFrame.Ping]] is received.
     */
-  def receiveBinary(pongOnPing: Boolean): F[Either[WebSocketFrame.Close, Array[Byte]]] =
+  def receiveBinary(pongOnPing: Boolean): F[Array[Byte]] =
     receiveConcat(() => receiveBinaryFrame(pongOnPing), _ ++ _)
 
+  /** Extracts the received close frame (if available) as the left side of an either, or returns the original result
+    * on the right.
+    *
+    * Will fail with [[WebSocketClosed]] if the web socket is closed, but no close frame is available.
+    *
+    * @param f The effect describing web socket interactions.
+    */
+  def eitherClose[T](f: => F[T]): F[Either[WebSocketFrame.Close, T]] =
+    f.map(t => Right(t): Either[WebSocketFrame.Close, T]).handleError { case WebSocketClosed(Some(close)) =>
+      (Left(close): Either[WebSocketFrame.Close, T]).unit
+    }
+
+  /** Returns an effect computing a:
+    *
+    * - `Left` if the web socket is closed - optionally with the received close frame (if available).
+    * - `Right` with the original result otherwise.
+    *
+    * Will never fail with a [[WebSocketClosed]].
+    *
+    * @param f The effect describing web socket interactions.
+    */
+  def either[T](f: => F[T]): F[Either[Option[WebSocketFrame.Close], T]] =
+    f.map(t => Right(t): Either[Option[WebSocketFrame.Close], T]).handleError { case WebSocketClosed(close) =>
+      (Left(close): Either[Option[WebSocketFrame.Close], T]).unit
+    }
+
   private def receiveConcat[T, U <: WebSocketFrame.Data[T]](
-      receiveSingle: () => F[Either[WebSocketFrame.Close, U]],
+      receiveSingle: () => F[U],
       combine: (T, T) => T
-  ): F[Either[WebSocketFrame.Close, T]] = {
+  ): F[T] = {
     receiveSingle().flatMap {
-      case Left(close) => (Left(close): Either[WebSocketFrame.Close, T]).unit
-      case Right(data) if !data.finalFragment =>
-        receiveConcat(receiveSingle, combine).flatMap {
-          case Left(close) => (Left(close): Either[WebSocketFrame.Close, T]).unit
-          case Right(t)    => (Right(combine(data.payload, t)): Either[WebSocketFrame.Close, T]).unit
+      case data if !data.finalFragment =>
+        receiveConcat(receiveSingle, combine).map { t =>
+          combine(data.payload, t)
         }
-      case Right(data) /* if data.finalFragment */ =>
-        (Right(data.payload): Either[WebSocketFrame.Close, T]).unit
+      case data /* if data.finalFragment */ => data.payload.unit
     }
   }
 

ws/src/main/scala/sttp/ws/WebSocketException.scala

@@ -2,6 +2,8 @@ package sttp.ws
 
 abstract class WebSocketException(msg: String) extends Exception(msg)
 
-class WebSocketClosed() extends WebSocketException(null)
+/** @param frame The received closing frame, if available.
+  */
+case class WebSocketClosed(frame: Option[WebSocketFrame.Close]) extends WebSocketException(null)
 
-class WebSocketBufferFull(capacity: Int) extends WebSocketException(s"Buffered $capacity messages")
+case class WebSocketBufferFull(capacity: Int) extends WebSocketException(s"Buffered $capacity messages")

ws/src/main/scala/sttp/ws/testing/WebSocketStub.scala

@@ -71,6 +71,7 @@ class WebSocketStub[S](
     new WebSocket[F] {
       private var state: S = initialState
       private var _isOpen: Boolean = true
+      private var closeFrame: Option[WebSocketFrame.Close] = None
       private var responses = initialResponses
 
       override def monad: MonadError[F] = m
@@ -83,6 +84,7 @@ class WebSocketStub[S](
               responses.headOption match {
                 case Some(Success(close: WebSocketFrame.Close)) =>
                   _isOpen = false
+                  closeFrame = Some(close)
                   monad.unit(close)
                 case Some(Success(response)) =>
                   responses = responses.tail
@@ -94,7 +96,7 @@ class WebSocketStub[S](
                   monad.error(new IllegalStateException("Unexpected 'receive', no more prepared responses."))
               }
             } else {
-              monad.error(new WebSocketClosed())
+              monad.error(WebSocketClosed(closeFrame))
             }
           }
         })

ws/src/test/scala/sttp/ws/testing/WebSocketStubTests.scala

@@ -121,7 +121,7 @@ class WebSocketStubTests extends AnyFlatSpec with Matchers with ScalaFutures {
         try rt()
         catch {
           case e if h.isDefinedAt(e) => h.apply(e)()
-          case e                     => throw e
+          case e: Exception          => throw e
         }
     override def ensure[T](f: () => T, e: => () => Unit): () => T =
       () =>