Ember Web Socket Client

.github/workflows/ci.yml

@@ -142,7 +142,7 @@ jobs:
   publish:
     name: Publish Artifacts
     needs: [build]
-    if: github.event_name != 'pull_request' && (startsWith(github.ref, 'refs/tags/v') || github.ref == 'refs/heads/series/0.23')
+    if: github.event_name != 'pull_request' && (startsWith(github.ref, 'refs/tags/v') || github.ref == 'refs/heads/series/0.23' || github.ref == 'refs/heads/ember-wsclient')
     strategy:
       matrix:
         os: [ubuntu-latest]

build.sbt

@@ -13,6 +13,8 @@ ThisBuild / developers += tlGitHubDev("rossabaker", "Ross A. Baker")
 ThisBuild / tlCiReleaseBranches := Seq("series/0.23")
 ThisBuild / tlSitePublishBranch := Some("series/0.23")
 
+ThisBuild / tlCiReleaseBranches += "ember-wsclient"
+
 ThisBuild / scalafixAll / skip := tlIsScala3.value
 ThisBuild / ScalafixConfig / skip := tlIsScala3.value
 ThisBuild / Test / scalafixConfig := Some(file(".scalafix.test.conf"))

ember-client/shared/src/main/scala/org/http4s/ember/client/EmberClientBuilder.scala

@@ -30,7 +30,10 @@ import org.http4s.Response
 import org.http4s.client._
 import org.http4s.client.middleware.Retry
 import org.http4s.client.middleware.RetryPolicy
+import org.http4s.client.websocket.WSClient
 import org.http4s.ember.client.internal.ClientHelpers
+import org.http4s.ember.client.internal.EmberWSClient
+import org.http4s.ember.client.internal.WebSocketKey
 import org.http4s.ember.core.h2.H2Client
 import org.http4s.ember.core.h2.H2Frame
 import org.http4s.ember.core.h2.H2Frame.Settings.ConnectionSettings.default
@@ -246,7 +249,7 @@ final class EmberClientBuilder[F[_]: Async: Network] private (
       )
       .whenA(timeout.isFinite && timeout >= idleConnectionTime)
 
-  def build: Resource[F, Client[F]] =
+  private def buildHelper(ws: Boolean = false): Resource[F, Client[F]] =
     for {
       _ <- Resource.eval(verifyTimeoutRelations)
       sg <- Resource.pure(sgOpt.getOrElse(Network[F]))
@@ -334,7 +337,14 @@ final class EmberClientBuilder[F[_]: Async: Network] private (
               case _ => Applicative[F].unit
             }
           }
-        } yield responseResource._1
+          _ <- Resource.eval(managed.canBeReused.set(Reusable.DontReuse))
+        } yield
+          if (ws)
+            responseResource._1.withAttribute(
+              WebSocketKey.webSocketConnection[F],
+              managed.value.keySocket.socket,
+            )
+          else responseResource._1
 
       def unixSocketClient(
           request: Request[F],
@@ -395,6 +405,14 @@ final class EmberClientBuilder[F[_]: Async: Network] private (
         new EmberClient(h2Client, pool)
       }
     }
+
+  def build: Resource[F, Client[F]] = buildHelper()
+
+  def buildWebSocket: Resource[F, (Client[F], WSClient[F])] =
+    for {
+      httpClient <- buildHelper(ws = true)
+      wsClient <- Resource.eval(EmberWSClient[F](httpClient))
+    } yield (httpClient, wsClient)
 }
 
 object EmberClientBuilder extends EmberClientBuilderCompanionPlatform {

ember-client/shared/src/main/scala/org/http4s/ember/client/internal/EmberWSClient.scala

@@ -0,0 +1,126 @@
+/*
+ * Copyright 2019 http4s.org
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.http4s.ember.client.internal
+
+import cats.MonadThrow
+import cats.effect.Async
+import cats.effect.implicits._
+import cats.effect.kernel.Resource
+import cats.effect.std.Queue
+import cats.effect.std.SecureRandom
+import cats.syntax.all._
+import fs2.concurrent.Channel
+import org.http4s.Request
+import org.http4s._
+import org.http4s.client.Client
+import org.http4s.client.websocket.WSClient
+import org.http4s.client.websocket.WSConnection
+import org.http4s.client.websocket.WSFrame
+import org.http4s.ember.client.internal.WebSocketHelpers._
+import org.http4s.ember.core.WebSocketHelpers._
+import org.http4s.ember.core.h2.H2Keys.WebSocketUpgradeIdentifier
+import org.http4s.headers.`Sec-WebSocket-Key`
+import org.http4s.websocket.WebSocketFrame
+import scodec.bits.ByteVector
+
+import java.util.Base64
+
+private[client] object EmberWSClient {
+  def apply[F[_]](
+      emberClient: Client[F]
+  )(implicit F: Async[F]): F[WSClient[F]] =
+    SecureRandom.javaSecuritySecureRandom[F].map { random =>
+      WSClient[F](respondToPings = false) { wsRequest =>
+        for {
+          randomByteArray <- Resource.eval(random.nextBytes(16))
+
+          uriScheme = wsRequest.uri.scheme.map(scheme =>
+            scheme.value match {
+              case "wss" => Uri.Scheme.https
+              case "ws" => Uri.Scheme.http
+              case _ => scheme
+            }
+          )
+
+          httpWSRequest = Request[F]()
+            .withUri(wsRequest.uri.copy(uriScheme))
+            .withHeaders(
+              Headers(
+                upgradeWebSocket,
+                connectionUpgrade,
+                supportedWebSocketVersionHeader,
+                new `Sec-WebSocket-Key`(ByteVector(Base64.getEncoder().encode(randomByteArray))),
+              )
+            )
+            .withMethod(Method.GET)
+            .withAttribute(WebSocketUpgradeIdentifier, ())
+
+          socketOption <- getSocket(emberClient, httpWSRequest)
+          socket <- socketOption.liftTo[F](new RuntimeException("Not an Ember client")).toResource
+
+          closeFrameDeffered <- F.deferred[WebSocketFrame.Close].toResource
+
+          clientReceiveQueue <- Queue.bounded[F, WebSocketFrame](100).toResource
+          clientSendChannel <- Channel.bounded[F, WebSocketFrame](100).toResource
+
+          _ <- socket.reads
+            .through(decodeFrames(true))
+            .foreach {
+              case f @ WebSocketFrame.Close(_) =>
+                closeFrameDeffered.complete(f).ifM(clientReceiveQueue.offer(f), F.unit)
+              case f =>
+                closeFrameDeffered.tryGet.flatMap { x =>
+                  if (x.isDefined) F.unit else clientReceiveQueue.offer(f)
+                }
+            }
+            .compile
+            .drain
+            .background
+
+          sendingFinished <- clientSendChannel.stream
+            .foreach(f => frameToBytes(f, true).traverse_(c => socket.write(c)))
+            .compile
+            .drain
+            .background
+
+          _ <- Resource.onFinalize {
+            MonadThrow[F]
+              .fromEither(WebSocketFrame.Close(1000, "Connection automatically closed"))
+              .flatMap(clientSendChannel.closeWithElement(_)) *> sendingFinished.void
+          }
+        } yield new WSConnection[F] {
+          def receive: F[Option[WSFrame]] = clientReceiveQueue.take.flatMap {
+            case f @ WebSocketFrame.Close(_) =>
+              closeChannelWithCloseFrame(clientSendChannel).as(toWSFrame(f).some)
+            case f =>
+              toWSFrame(f).some.pure[F]
+          }
+          def send(wsf: WSFrame): F[Unit] =
+            toWebSocketFrame(wsf).flatMap {
+              case WebSocketFrame.Close(_) =>
+                closeChannelWithCloseFrame(clientSendChannel)
+              case f =>
+                clientSendChannel.send(f).void
+            }
+          def sendMany[G[_], A <: WSFrame](wsfs: G[A])(implicit
+              evidence$1: cats.Foldable[G]
+          ): F[Unit] = wsfs.traverse_(send(_))
+          def subprotocol: Option[String] = ???
+        }
+      }
+    }
+}

ember-client/shared/src/main/scala/org/http4s/ember/client/internal/WebSocketHelpers.scala

@@ -0,0 +1,172 @@
+/*
+ * Copyright 2019 http4s.org
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.http4s.ember.client.internal
+
+import cats.Applicative
+import cats.MonadThrow
+import cats.data.EitherT
+import cats.data.NonEmptyList
+import cats.effect.Concurrent
+import cats.effect.MonadCancel
+import cats.effect.Resource
+import cats.syntax.all._
+import fs2.concurrent.Channel
+import fs2.io.net.Socket
+import org.http4s.Request
+import org.http4s.Status
+import org.http4s._
+import org.http4s.client.Client
+import org.http4s.client.websocket.WSFrame
+import org.http4s.crypto.Hash
+import org.http4s.crypto.HashAlgorithm
+import org.http4s.headers._
+import org.http4s.websocket.Rfc6455
+import org.http4s.websocket.WebSocketFrame
+import org.typelevel.ci._
+import scodec.bits.ByteVector
+
+private[internal] object WebSocketHelpers {
+
+  val supportedWebSocketVersion = 13L
+
+  val supportedWebSocketVersionHeader: `Sec-WebSocket-Version` =
+    `Sec-WebSocket-Version`.unsafeFromLong(
+      supportedWebSocketVersion
+    )
+  val upgradeCi: CIString = ci"upgrade"
+  val webSocketProtocol: Protocol = Protocol(ci"websocket", None)
+  val connectionUpgrade: Connection = Connection(NonEmptyList.of(upgradeCi))
+  val upgradeWebSocket: Upgrade = Upgrade(webSocketProtocol)
+
+  def closeChannelWithCloseFrame[F[_]: MonadThrow](
+      clientSendChannel: Channel[F, WebSocketFrame]
+  ): F[Unit] =
+    for {
+      closeFrame <-
+        MonadThrow[F]
+          .fromEither(WebSocketFrame.Close(1000, "Connection automatically closed"))
+
+      _ <- EitherT(clientSendChannel.closeWithElement(closeFrame))
+        .getOrRaise(new RuntimeException("Connection already closed"))
+    } yield ()
+
+  def getSocket[F[_]](client: Client[F], request: Request[F])(implicit
+      F: MonadCancel[F, Throwable]
+  ): Resource[F, Option[Socket[F]]] = {
+    val webSocketKey = WebSocketKey.webSocketConnection[F]
+    client
+      .run(request)
+      .evalMap { res =>
+        for {
+          secWebSocketKeyString <- request.headers
+            .get[`Sec-WebSocket-Key`]
+            .liftTo[F](new RuntimeException("Sec-WebSocket-Key header not found"))
+            .map(_.hashString)
+          isValid <- validateServerHandshake(res, secWebSocketKeyString)
+        } yield isValid.toOption *> res.attributes.lookup(webSocketKey)
+      }
+  }
+
+  def toWebSocketFrame[F[_]: Concurrent](wsFrame: WSFrame): F[WebSocketFrame] =
+    wsFrame match {
+      case WSFrame.Close(code, reason) =>
+        MonadThrow[F].fromEither(WebSocketFrame.Close(code, reason))
+      case WSFrame.Ping(data) => Applicative[F].pure(WebSocketFrame.Ping(data))
+      case WSFrame.Pong(data) => Applicative[F].pure(WebSocketFrame.Pong(data))
+      case WSFrame.Text(data, last) => Applicative[F].pure(WebSocketFrame.Text(data, last))
+      case WSFrame.Binary(data, last) => Applicative[F].pure(WebSocketFrame.Binary(data, last))
+    }
+
+  def toWSFrame(wsf: WebSocketFrame): WSFrame =
+    (wsf: @unchecked) match {
+      case c: WebSocketFrame.Close => WSFrame.Close(c.closeCode, c.reason)
+      case WebSocketFrame.Ping(data) => WSFrame.Ping(data)
+      case WebSocketFrame.Pong(data) => WSFrame.Pong(data)
+      case WebSocketFrame.Text(data, last) => WSFrame.Text(data, last)
+      case WebSocketFrame.Binary(data, last) => WSFrame.Binary(data, last)
+    }
+
+  /** Validate the opening handshake response from the server
+    * https://datatracker.ietf.org/doc/html/rfc6455#page-6
+    */
+  def validateServerHandshake[F[_]](
+      response: Response[F],
+      secWebSocketKey: String,
+  )(implicit F: MonadThrow[F]): F[Either[ServerHandshakeError, Unit]] =
+    for {
+      secWebSocketAccept <- serverHandshake(response).pure[F]
+      correctSecWebSocketAccept <- clientHandshake(secWebSocketKey)
+      validated = secWebSocketAccept.flatMap(s =>
+        if (s == correctSecWebSocketAccept) Either.unit else Left(InvalidSecWebSocketAccept)
+      )
+    } yield validated
+
+  private[this] val magic = ByteVector.view(Rfc6455.handshakeMagicBytes)
+
+  def clientHandshake[F[_]](
+      value: String
+  )(implicit F: MonadThrow[F]): F[ByteVector] = for {
+    value <- ByteVector.encodeAscii(value).liftTo[F]
+    digest <- Hash[F].digest(HashAlgorithm.SHA1, value ++ magic)
+  } yield digest
+
+  private def serverHandshake[F[_]](res: Response[F]): Either[ServerHandshakeError, ByteVector] = {
+    val status = res.status match {
+      case Status.SwitchingProtocols => Either.unit
+      case _ => Left(InvalidStatus)
+    }
+
+    val connection = res.headers.get[Connection] match {
+      case Some(header) if header.hasUpgrade => Either.unit
+      case _ => Left(UpgradeRequired)
+    }
+
+    val upgrade = res.headers.get[Upgrade] match {
+      case Some(header) if header.values.contains_(webSocketProtocol) => Either.unit
+      case _ => Left(UpgradeRequired)
+    }
+
+    val secWebSocketAcceptKey = res.headers.get[`Sec-WebSocket-Accept`] match {
+      case Some(header) => Right(header.hashedKey)
+      case None => Left(SecWebSocketAcceptNotFound)
+    }
+
+    (status, connection, upgrade, secWebSocketAcceptKey).mapN {
+      case (_, _, _, secWebSocketAcceptKey) =>
+        secWebSocketAcceptKey
+    }
+  }
+
+  sealed abstract class ServerHandshakeError(val status: Status, val message: String)
+  case object InvalidStatus
+      extends ServerHandshakeError(
+        Status.BadRequest,
+        "Not found HTTP Status 101 Switching Protocol.",
+      )
+  case object UpgradeRequired
+      extends ServerHandshakeError(
+        Status.UpgradeRequired,
+        "Upgrade required for WebSocket communication.",
+      )
+  case object SecWebSocketAcceptNotFound
+      extends ServerHandshakeError(Status.BadRequest, "Sec-WebSocket-Accept header not present.")
+  case object InvalidSecWebSocketAccept
+      extends ServerHandshakeError(
+        Status.BadRequest,
+        "Sec-WebSocket-Accept does not correspond to the Sec-WebSocket-Key",
+      )
+}

ember-client/shared/src/main/scala/org/http4s/ember/client/internal/WebSocketKey.scala

@@ -0,0 +1,28 @@
+/*
+ * Copyright 2019 http4s.org
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.http4s.ember.client.internal
+
+import cats.effect.SyncIO
+import fs2.io.net.Socket
+import org.typelevel.vault._
+
+private[client] object WebSocketKey {
+
+  private[this] val wsConnectionInternal: Key[Any] = Key.newKey[SyncIO, Any].unsafeRunSync()
+  def webSocketConnection[F[_]]: Key[Socket[F]] =
+    wsConnectionInternal.asInstanceOf[Key[Socket[F]]]
+}