Add selective WebSocket message compression

docs/src/main/paradox/routing-dsl/directives/websocket-directives/handleWebSocketMessages.md

@@ -13,6 +13,9 @@
 The directive first checks if the request was a valid WebSocket handshake request and if yes, it completes the request
 with the passed handler. Otherwise, the request is rejected with an @apidoc[ExpectedWebSocketRequestRejection$].
 
+The overload that accepts a `shouldCompress` filter can select compression separately for each outbound message after
+`permessage-deflate` is negotiated.
+
 WebSocket subprotocols offered in the `Sec-WebSocket-Protocol` header of the request are ignored. If you want to
 support several protocols use the @ref[handleWebSocketMessagesForProtocol](handleWebSocketMessagesForProtocol.md) directive, instead.
 

docs/src/main/paradox/routing-dsl/directives/websocket-directives/handleWebSocketMessagesForOptionalProtocol.md

@@ -18,6 +18,9 @@ If the `subprotocol` parameter is @scala[None]@java[@javadoc:[empty](java.util.O
 announced in the WebSocket request) @scala[contains `protocol`]@java[matches the contained subprotocol]. If the client did not offer the protocol in question
 the request is rejected with an @apidoc[UnsupportedWebSocketSubprotocolRejection].
 
+The overload that accepts a `shouldCompress` filter can select compression separately for each outbound message after
+`permessage-deflate` is negotiated.
+
 To support several subprotocols you may chain several `handleWebSocketMessagesForOptionalProtocol` routes.
 
 The `handleWebSocketMessagesForOptionalProtocol` directive is used as a building block for @ref[WebSocket Directives](index.md) to handle websocket messages.

docs/src/main/paradox/routing-dsl/directives/websocket-directives/handleWebSocketMessagesForProtocol.md

@@ -18,6 +18,9 @@ The directive first checks if the request was a valid WebSocket handshake reques
 subprotocol name. If yes, the directive completes the request with the passed handler. Otherwise, the request is
 either rejected with an @apidoc[ExpectedWebSocketRequestRejection$] or an @apidoc[UnsupportedWebSocketSubprotocolRejection].
 
+The overload that accepts a `shouldCompress` filter can select compression separately for each outbound message after
+`permessage-deflate` is negotiated.
+
 To support several subprotocols, for example at the same path, several instances of `handleWebSocketMessagesForProtocol` can
 be chained using `~` as you can see in the below example.
 

docs/src/main/paradox/server-side/websocket-support.md

@@ -171,6 +171,20 @@ The server exposes additional settings for the negotiated extension under
 If compression is enabled globally, a route can still decline compression for a single accepted WebSocket by using the
 `handleMessages` or `handleMessagesWith` overload with `compressionEnabled = false`.
 
+After `permessage-deflate` is negotiated, an application can also select which server-to-client messages are compressed.
+The compression filter is evaluated once for each outbound text or binary message. Returning `true` compresses that
+message; returning `false` sends it uncompressed. The same decision applies to every fragment of a streamed message.
+
+Scala
+:  @@snip [WebSocketExampleSpec.scala](/docs/src/test/scala/docs/http/scaladsl/server/WebSocketExampleSpec.scala) { #websocket-selective-compression }
+
+Java
+:  @@snip [WebSocketCoreExample.java](/docs/src/test/java/docs/http/javadsl/server/WebSocketCoreExample.java) { #websocket-selective-compression }
+
+The filter only controls outbound messages. It does not affect extension negotiation, and the client can still send both
+compressed and uncompressed messages. The filter is not evaluated if `permessage-deflate` was not negotiated. To disable
+compression in both directions for a security-sensitive endpoint, use the `compressionEnabled = false` overload instead.
+
 @@@ note
 The `server_no_context_takeover` and `client_no_context_takeover` extension parameters affect whether compression
 dictionaries are retained across messages. Retaining context generally improves compression ratio, while disabling

docs/src/test/java/docs/http/javadsl/server/WebSocketCoreExample.java

@@ -21,6 +21,7 @@
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Predicate;
 import org.apache.pekko.NotUsed;
 import org.apache.pekko.actor.ActorSystem;
 import org.apache.pekko.http.javadsl.ConnectionContext;
@@ -32,6 +33,7 @@
 import org.apache.pekko.http.javadsl.model.ws.Message;
 import org.apache.pekko.http.javadsl.model.ws.TextMessage;
 import org.apache.pekko.http.javadsl.model.ws.WebSocketRequest;
+import org.apache.pekko.http.javadsl.model.ws.WebSocketUpgrade;
 import org.apache.pekko.http.javadsl.settings.ClientConnectionSettings;
 import org.apache.pekko.http.javadsl.settings.ServerSettings;
 import org.apache.pekko.http.javadsl.settings.WebSocketSettings;
@@ -123,6 +125,16 @@ public static TextMessage handleTextMessage(TextMessage msg) {
 
   // #websocket-handler
 
+  static HttpResponse selectiveCompression(
+      WebSocketUpgrade upgrade, Flow<Message, Message, NotUsed> handler) {
+    // #websocket-selective-compression
+    Predicate<Message> shouldCompress = Message::isText;
+
+    HttpResponse response = upgrade.handleMessagesWith(handler, shouldCompress);
+    // #websocket-selective-compression
+    return response;
+  }
+
   {
     ActorSystem system = null;
     Flow<HttpRequest, HttpResponse, NotUsed> handler = null;

docs/src/test/scala/docs/http/scaladsl/server/WebSocketExampleSpec.scala

@@ -125,6 +125,20 @@ class WebSocketExampleSpec extends AnyWordSpec with Matchers with CompileOnlySpe
       .onComplete(_ => system.terminate()) // and shutdown when done
   }
 
+  "selective-compression-example" in compileOnlySpec {
+    import pekko.http.scaladsl.model.ws.{ Message, WebSocketUpgrade }
+    import pekko.stream.scaladsl.Flow
+
+    val upgrade: WebSocketUpgrade = null
+    val handler: Flow[Message, Message, Any] = null
+
+    // #websocket-selective-compression
+    val shouldCompress: Message => Boolean = _.isText
+
+    val response = upgrade.handleMessages(handler, shouldCompress)
+    // #websocket-selective-compression
+  }
+
   "ping-server-example" in compileOnlySpec {
     implicit val system: ActorSystem = null
     val route = null

http-core/src/main/scala/org/apache/pekko/http/impl/engine/ws/Handshake.scala

@@ -42,6 +42,9 @@ private[http] object Handshake {
 
   object Server {
 
+    private val CompressEveryMessage: Message => Boolean = _ => true
+    private val CompressEveryFrame: FrameStart => Boolean = _ => true
+
     /**
      *  Validates a client WebSocket handshake. Returns either `OptionVal.Some(UpgradeToWebSocketLowLevel)` or
      *  `OptionVal.None`
@@ -140,12 +143,22 @@ private[http] object Handshake {
             def handle(
                 handler: Either[Graph[FlowShape[FrameEvent, FrameEvent], Any], Graph[FlowShape[Message, Message], Any]],
                 subprotocol: Option[String],
-                compressionEnabled: Boolean): HttpResponse = {
+                compressionEnabled: Boolean,
+                shouldCompressMessage: Message => Boolean = CompressEveryMessage,
+                shouldCompressFrame: FrameStart => Boolean = CompressEveryFrame): HttpResponse = {
               require(
                 subprotocol.forall(chosen => clientSupportedSubprotocols.contains(chosen)),
                 s"Tried to choose invalid subprotocol '$subprotocol' which wasn't offered by the client: [${requestedProtocols.mkString(", ")}]")
               val acceptedPerMessageDeflate = if (compressionEnabled) perMessageDeflate else None
-              buildResponse(key.get, handler, subprotocol, acceptedPerMessageDeflate, settings, log)
+              buildResponse(
+                key.get,
+                handler,
+                subprotocol,
+                acceptedPerMessageDeflate,
+                settings,
+                log,
+                shouldCompressMessage,
+                shouldCompressFrame)
             }
 
             def handleFrames(
@@ -158,6 +171,13 @@ private[http] object Handshake {
                 compressionEnabled: Boolean): HttpResponse =
               handle(Left(handlerFlow), subprotocol, compressionEnabled)
 
+            override private[http] def handleFrames(
+                handlerFlow: Graph[FlowShape[FrameEvent, FrameEvent], Any],
+                subprotocol: Option[String],
+                compressionEnabled: Boolean,
+                shouldCompress: FrameStart => Boolean): HttpResponse =
+              handle(Left(handlerFlow), subprotocol, compressionEnabled, shouldCompressFrame = shouldCompress)
+
             override def handleMessages(handlerFlow: Graph[FlowShape[Message, Message], Any],
                 subprotocol: Option[String] = None): HttpResponse =
               handle(Right(handlerFlow), subprotocol, compressionEnabled = true)
@@ -167,6 +187,12 @@ private[http] object Handshake {
                 subprotocol: Option[String],
                 compressionEnabled: Boolean): HttpResponse =
               handle(Right(handlerFlow), subprotocol, compressionEnabled)
+
+            override def handleMessages(
+                handlerFlow: Graph[FlowShape[Message, Message], Any],
+                subprotocol: Option[String],
+                shouldCompress: Message => Boolean): HttpResponse =
+              handle(Right(handlerFlow), subprotocol, compressionEnabled = true, shouldCompressMessage = shouldCompress)
           }
           OptionVal.Some(header)
         } else OptionVal.None
@@ -197,12 +223,33 @@ private[http] object Handshake {
         subprotocol: Option[String],
         perMessageDeflate: Option[PerMessageDeflate.Negotiated],
         settings: WebSocketSettings,
-        log: LoggingAdapter): HttpResponse = {
+        log: LoggingAdapter): HttpResponse =
+      buildResponse(
+        key,
+        handler,
+        subprotocol,
+        perMessageDeflate,
+        settings,
+        log,
+        CompressEveryMessage,
+        CompressEveryFrame)
+
+    private def buildResponse(
+        key: `Sec-WebSocket-Key`,
+        handler: Either[Graph[FlowShape[FrameEvent, FrameEvent], Any], Graph[FlowShape[Message, Message], Any]],
+        subprotocol: Option[String],
+        perMessageDeflate: Option[PerMessageDeflate.Negotiated],
+        settings: WebSocketSettings,
+        log: LoggingAdapter,
+        shouldCompressMessage: Message => Boolean,
+        shouldCompressFrame: FrameStart => Boolean): HttpResponse = {
       val frameHandler = handler match {
         case Left(frameHandler) =>
-          perMessageDeflate.map(_.frameEventBidiFlow(settings.randomFactory).join(frameHandler)).getOrElse(frameHandler)
+          perMessageDeflate
+            .map(_.frameEventBidiFlow(settings.randomFactory, shouldCompressFrame).join(frameHandler))
+            .getOrElse(frameHandler)
         case Right(messageHandler) =>
-          WebSocket.stack(serverSide = true, settings, perMessageDeflate = perMessageDeflate, log = log)
+          WebSocket.stack(true, settings, perMessageDeflate, log, shouldCompressMessage)
             .join(messageHandler)
       }
 

http-core/src/main/scala/org/apache/pekko/http/impl/engine/ws/MessageToFrameRenderer.scala

@@ -30,23 +30,34 @@ import pekko.http.scaladsl.model.ws._
  */
 @InternalApi
 private[http] object MessageToFrameRenderer {
-  def create(serverSide: Boolean): Flow[Message, FrameStart, NotUsed] = {
-    def strictFrames(opcode: Opcode, data: ByteString): Source[FrameStart, ?] =
+  def create(serverSide: Boolean): Flow[Message, FrameStart, NotUsed] =
+    create(serverSide, None)
+
+  def create(serverSide: Boolean, shouldCompress: Message => Boolean): Flow[Message, FrameStart, NotUsed] =
+    create(serverSide, Some(shouldCompress))
+
+  private def create(
+      serverSide: Boolean,
+      shouldCompress: Option[Message => Boolean]): Flow[Message, FrameStart, NotUsed] = {
+    def strictFrames(opcode: Opcode, data: ByteString, compress: Boolean): Source[FrameStart, ?] =
       // FIXME: fragment?
-      Source.single(FrameEvent.fullFrame(opcode, None, data, fin = true))
+      Source.single(FrameEvent.fullFrame(opcode, None, data, fin = true, rsv1 = compress))
 
-    def streamedFrames[M](opcode: Opcode, data: Source[ByteString, M]): Source[FrameStart, Any] =
+    def streamedFrames[M](opcode: Opcode, data: Source[ByteString, M], compress: Boolean): Source[FrameStart, Any] =
       data.statefulMap(() => true)((isFirst, data) => {
           val frameOpcode = if (isFirst) opcode else Opcode.Continuation
-          (false, FrameEvent.fullFrame(frameOpcode, None, data, fin = false))
+          (false, FrameEvent.fullFrame(frameOpcode, None, data, fin = false, rsv1 = isFirst && compress))
         }, _ => None) ++ Source.single(FrameEvent.emptyLastContinuationFrame)
 
     Flow[Message]
-      .flatMapConcat {
-        case BinaryMessage.Strict(data) => strictFrames(Opcode.Binary, data)
-        case bm: BinaryMessage          => streamedFrames(Opcode.Binary, bm.dataStream)
-        case TextMessage.Strict(text)   => strictFrames(Opcode.Text, ByteString(text, StandardCharsets.UTF_8))
-        case tm: TextMessage            => streamedFrames(Opcode.Text, tm.textStream.via(Utf8Encoder))
+      .flatMapConcat { message =>
+        val compress = shouldCompress.exists(_(message))
+        message match {
+          case BinaryMessage.Strict(data) => strictFrames(Opcode.Binary, data, compress)
+          case bm: BinaryMessage          => streamedFrames(Opcode.Binary, bm.dataStream, compress)
+          case TextMessage.Strict(text)   => strictFrames(Opcode.Text, ByteString(text, StandardCharsets.UTF_8), compress)
+          case tm: TextMessage            => streamedFrames(Opcode.Text, tm.textStream.via(Utf8Encoder), compress)
+        }
       }
   }
 }

http-core/src/main/scala/org/apache/pekko/http/impl/engine/ws/PerMessageDeflate.scala

@@ -70,8 +70,16 @@ private[http] object PerMessageDeflate {
     def bidiFlow: BidiFlow[FrameEventOrError, FrameEventOrError, FrameEvent, FrameEvent, NotUsed] =
       BidiFlow.fromFlows(inflaterFlow, deflaterFlow)
 
+    def messageBidiFlow: BidiFlow[FrameEventOrError, FrameEventOrError, FrameEvent, FrameEvent, NotUsed] =
+      BidiFlow.fromFlows(inflaterFlow, selectiveDeflaterFlow(_.header.rsv1, rsv1IndicatesCompression = true))
+
     def frameEventBidiFlow(
         maskRandom: () => Random): BidiFlow[FrameEvent, FrameEvent, FrameEvent, FrameEvent, NotUsed] =
+      frameEventBidiFlow(maskRandom, _ => true)
+
+    def frameEventBidiFlow(
+        maskRandom: () => Random,
+        shouldCompress: FrameStart => Boolean): BidiFlow[FrameEvent, FrameEvent, FrameEvent, FrameEvent, NotUsed] =
       BidiFlow.fromFlows(
         Flow[FrameEvent]
           .via(Masking.unmaskIf(condition = true))
@@ -81,13 +89,26 @@ private[http] object PerMessageDeflate {
             case FrameError(ex)    => throw ex
           }
           .via(Masking.maskIf(condition = true, maskRandom)),
-        deflaterFlow)
+        selectiveDeflaterFlow(shouldCompress, rsv1IndicatesCompression = false))
 
     private def inflaterFlow: Flow[FrameEventOrError, FrameEventOrError, NotUsed] =
       createInflaterFlow(clientNoContextTakeover, settings, DefaultCompressionFactory)
 
     private def deflaterFlow: Flow[FrameEvent, FrameEvent, NotUsed] =
       createDeflaterFlow(serverNoContextTakeover, settings, DefaultCompressionFactory)
+
+    private def selectiveDeflaterFlow(
+        shouldCompress: FrameStart => Boolean,
+        rsv1IndicatesCompression: Boolean): Flow[FrameEvent, FrameEvent, NotUsed] =
+      Flow.fromGraph(new LifecycleMapConcatStage(
+        "PerMessageDeflate.deflater",
+        () =>
+          new DeflaterFlow(
+            serverNoContextTakeover,
+            settings,
+            DefaultCompressionFactory,
+            shouldCompress,
+            rsv1IndicatesCompression)))
   }
 
   private[ws] def createInflaterFlow(
@@ -104,7 +125,13 @@ private[http] object PerMessageDeflate {
       compressionFactory: CompressionFactory): Flow[FrameEvent, FrameEvent, NotUsed] =
     Flow.fromGraph(new LifecycleMapConcatStage(
       "PerMessageDeflate.deflater",
-      () => new DeflaterFlow(noContextTakeover, settings, compressionFactory)))
+      () =>
+        new DeflaterFlow(
+          noContextTakeover,
+          settings,
+          compressionFactory,
+          _ => true,
+          rsv1IndicatesCompression = false)))
 
   def negotiate(
       requested: immutable.Seq[WebSocketExtension],
@@ -242,19 +269,22 @@ private[http] object PerMessageDeflate {
   private final class DeflaterFlow(
       noContextTakeover: Boolean,
       settings: WebSocketCompressionSettingsImpl,
-      compressionFactory: CompressionFactory)
+      compressionFactory: CompressionFactory,
+      shouldCompress: FrameStart => Boolean,
+      rsv1IndicatesCompression: Boolean)
       extends LifecycleMapConcat[FrameEvent, FrameEvent] {
     private var deflater = compressionFactory.newDeflater(settings.compressionLevel)
     private var frame: Option[UncompressedFrame] = None
     private var messageInProgress = false
+    private var compressFragmentedMessage = false
     private var bypassFrameInProgress = false
     private val buffer = new Array[Byte](8192)
 
     override def apply(event: FrameEvent): immutable.Iterable[FrameEvent] = event match {
       case FrameStart(header, _)
           if (header.opcode == Protocol.Opcode.Text ||
           header.opcode == Protocol.Opcode.Binary) &&
-          (header.rsv1 || header.rsv2 || header.rsv3) =>
+          (header.rsv2 || header.rsv3 || (header.rsv1 && !rsv1IndicatesCompression)) =>
         throw new ProtocolException("Unexpected reserved bit for outbound WebSocket message")
       case FrameStart(header, _)
           if header.opcode == Protocol.Opcode.Continuation &&
@@ -265,18 +295,32 @@ private[http] object PerMessageDeflate {
           header.opcode == Protocol.Opcode.Binary =>
         if (messageInProgress || frame.isDefined)
           throw new ProtocolException("Unexpected data frame while fragmented message is open")
+        val compress = if (rsv1IndicatesCompression) header.rsv1 else shouldCompress(start)
         messageInProgress = !header.fin
-        frame = Some(UncompressedFrame(header.copy(length = 0, rsv1 = true), data, removeTail = header.fin))
-        if (start.lastPart) finishFrame() else Nil
+        compressFragmentedMessage = compress && messageInProgress
+        if (compress) {
+          frame = Some(UncompressedFrame(header.copy(length = 0, rsv1 = true), data, removeTail = header.fin))
+          if (start.lastPart) finishFrame() else Nil
+        } else {
+          bypassFrameInProgress = !start.lastPart
+          start :: Nil
+        }
       case start @ FrameStart(header, _) if bypassFrameInProgress =>
         throw new ProtocolException(s"Unexpected frame ${header.opcode} while frame data is open")
       case start @ FrameStart(header, _) if (frame.isDefined || messageInProgress) && header.opcode.isControl =>
         bypassFrameInProgress = !start.lastPart
         start :: Nil
       case start @ FrameStart(header, data) if messageInProgress && header.opcode == Protocol.Opcode.Continuation =>
+        val compress = compressFragmentedMessage
         messageInProgress = !header.fin
-        frame = Some(UncompressedFrame(header.copy(length = 0), data, removeTail = header.fin))
-        if (start.lastPart) finishFrame() else Nil
+        if (!messageInProgress) compressFragmentedMessage = false
+        if (compress) {
+          frame = Some(UncompressedFrame(header.copy(length = 0), data, removeTail = header.fin))
+          if (start.lastPart) finishFrame() else Nil
+        } else {
+          bypassFrameInProgress = !start.lastPart
+          start :: Nil
+        }
       case start @ FrameStart(header, _) if frame.isDefined || messageInProgress =>
         throw new ProtocolException(s"Unexpected frame ${header.opcode} while fragmented message is open")
       case data: FrameData if bypassFrameInProgress =>

http-core/src/main/scala/org/apache/pekko/http/impl/engine/ws/UpgradeToWebSocketLowLevel.scala

@@ -49,4 +49,16 @@ private[http] abstract class UpgradeToWebSocketLowLevel extends InternalCustomHe
       subprotocol: Option[String],
       compressionEnabled: Boolean): HttpResponse =
     handleFrames(handlerFlow, subprotocol)
+
+  /**
+   * The `shouldCompress` function is evaluated once for every outbound text or binary message when
+   * `permessage-deflate` was negotiated. The decision for the initial frame is retained for every continuation frame.
+   */
+  @InternalApi
+  private[http] def handleFrames(
+      handlerFlow: Graph[FlowShape[FrameEvent, FrameEvent], Any],
+      subprotocol: Option[String],
+      compressionEnabled: Boolean,
+      shouldCompress: FrameStart => Boolean): HttpResponse =
+    handleFrames(handlerFlow, subprotocol, compressionEnabled)
 }