diff --git a/docs/src/main/paradox/routing-dsl/directives/websocket-directives/handleWebSocketMessages.md b/docs/src/main/paradox/routing-dsl/directives/websocket-directives/handleWebSocketMessages.md index 0d30e5a1c3..fb675362f2 100644 --- a/docs/src/main/paradox/routing-dsl/directives/websocket-directives/handleWebSocketMessages.md +++ b/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. diff --git a/docs/src/main/paradox/routing-dsl/directives/websocket-directives/handleWebSocketMessagesForOptionalProtocol.md b/docs/src/main/paradox/routing-dsl/directives/websocket-directives/handleWebSocketMessagesForOptionalProtocol.md index 9c978800c2..40f4222edb 100644 --- a/docs/src/main/paradox/routing-dsl/directives/websocket-directives/handleWebSocketMessagesForOptionalProtocol.md +++ b/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. diff --git a/docs/src/main/paradox/routing-dsl/directives/websocket-directives/handleWebSocketMessagesForProtocol.md b/docs/src/main/paradox/routing-dsl/directives/websocket-directives/handleWebSocketMessagesForProtocol.md index 9a3184a54e..0664a1fbaf 100644 --- a/docs/src/main/paradox/routing-dsl/directives/websocket-directives/handleWebSocketMessagesForProtocol.md +++ b/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. diff --git a/docs/src/main/paradox/server-side/websocket-support.md b/docs/src/main/paradox/server-side/websocket-support.md index f0f8326a6a..6a2256a684 100644 --- a/docs/src/main/paradox/server-side/websocket-support.md +++ b/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 diff --git a/docs/src/test/java/docs/http/javadsl/server/WebSocketCoreExample.java b/docs/src/test/java/docs/http/javadsl/server/WebSocketCoreExample.java index e28988938d..b09f2ad792 100644 --- a/docs/src/test/java/docs/http/javadsl/server/WebSocketCoreExample.java +++ b/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 handler) { + // #websocket-selective-compression + Predicate shouldCompress = Message::isText; + + HttpResponse response = upgrade.handleMessagesWith(handler, shouldCompress); + // #websocket-selective-compression + return response; + } + { ActorSystem system = null; Flow handler = null; diff --git a/docs/src/test/scala/docs/http/scaladsl/server/WebSocketExampleSpec.scala b/docs/src/test/scala/docs/http/scaladsl/server/WebSocketExampleSpec.scala index a446a75220..9cb7b63905 100644 --- a/docs/src/test/scala/docs/http/scaladsl/server/WebSocketExampleSpec.scala +++ b/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 diff --git a/http-core/src/main/scala/org/apache/pekko/http/impl/engine/ws/Handshake.scala b/http-core/src/main/scala/org/apache/pekko/http/impl/engine/ws/Handshake.scala index b9b293c9ec..77a2096fe1 100644 --- a/http-core/src/main/scala/org/apache/pekko/http/impl/engine/ws/Handshake.scala +++ b/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) } diff --git a/http-core/src/main/scala/org/apache/pekko/http/impl/engine/ws/MessageToFrameRenderer.scala b/http-core/src/main/scala/org/apache/pekko/http/impl/engine/ws/MessageToFrameRenderer.scala index 1c5df2dcec..0ae9389ce4 100644 --- a/http-core/src/main/scala/org/apache/pekko/http/impl/engine/ws/MessageToFrameRenderer.scala +++ b/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) + } } } } diff --git a/http-core/src/main/scala/org/apache/pekko/http/impl/engine/ws/PerMessageDeflate.scala b/http-core/src/main/scala/org/apache/pekko/http/impl/engine/ws/PerMessageDeflate.scala index 7bf669aa9d..771da9027c 100644 --- a/http-core/src/main/scala/org/apache/pekko/http/impl/engine/ws/PerMessageDeflate.scala +++ b/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,11 +269,14 @@ 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) @@ -254,7 +284,7 @@ private[http] object PerMessageDeflate { 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 => diff --git a/http-core/src/main/scala/org/apache/pekko/http/impl/engine/ws/UpgradeToWebSocketLowLevel.scala b/http-core/src/main/scala/org/apache/pekko/http/impl/engine/ws/UpgradeToWebSocketLowLevel.scala index ca5fbd1653..d2274f4d20 100644 --- a/http-core/src/main/scala/org/apache/pekko/http/impl/engine/ws/UpgradeToWebSocketLowLevel.scala +++ b/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) } diff --git a/http-core/src/main/scala/org/apache/pekko/http/impl/engine/ws/WebSocket.scala b/http-core/src/main/scala/org/apache/pekko/http/impl/engine/ws/WebSocket.scala index 5d1d37656a..3e6fdcb7a7 100644 --- a/http-core/src/main/scala/org/apache/pekko/http/impl/engine/ws/WebSocket.scala +++ b/http-core/src/main/scala/org/apache/pekko/http/impl/engine/ws/WebSocket.scala @@ -50,12 +50,29 @@ private[http] object WebSocket { perMessageDeflate: Option[PerMessageDeflate.Negotiated] = None, closeTimeout: FiniteDuration = 3.seconds, // TODO put close timeout into the settings? log: LoggingAdapter): BidiFlow[FrameEvent, Message, Message, FrameEvent, NotUsed] = + stack(serverSide, websocketSettings, perMessageDeflate, closeTimeout, log, _ => true) + + def stack( + serverSide: Boolean, + websocketSettings: WebSocketSettings, + perMessageDeflate: Option[PerMessageDeflate.Negotiated], + log: LoggingAdapter, + shouldCompress: Message => Boolean): BidiFlow[FrameEvent, Message, Message, FrameEvent, NotUsed] = + stack(serverSide, websocketSettings, perMessageDeflate, 3.seconds, log, shouldCompress) + + def stack( + serverSide: Boolean, + websocketSettings: WebSocketSettings, + perMessageDeflate: Option[PerMessageDeflate.Negotiated], + closeTimeout: FiniteDuration, + log: LoggingAdapter, + shouldCompress: Message => Boolean): BidiFlow[FrameEvent, Message, Message, FrameEvent, NotUsed] = masking(serverSide, websocketSettings.randomFactory).atop( FrameLogger.logFramesIfEnabled(websocketSettings.logFrames)).atop( - perMessageDeflate.map(_.bidiFlow).getOrElse(BidiFlow.identity)).atop( + perMessageDeflate.map(_.messageBidiFlow).getOrElse(BidiFlow.identity)).atop( frameHandling(serverSide, closeTimeout, log)).atop( periodicKeepAlive(websocketSettings)).atop( - messageAPI(serverSide, closeTimeout)) + messageAPI(serverSide, closeTimeout, perMessageDeflate.map(_ => shouldCompress))) /** The lowest layer that implements the binary protocol */ def framing: BidiFlow[ByteString, FrameEvent, FrameEvent, ByteString, NotUsed] = @@ -154,7 +171,14 @@ private[http] object WebSocket { */ def messageAPI( serverSide: Boolean, - closeTimeout: FiniteDuration): BidiFlow[FrameHandler.Output, Message, Message, FrameOutHandler.Input, NotUsed] = { + closeTimeout: FiniteDuration): BidiFlow[FrameHandler.Output, Message, Message, FrameOutHandler.Input, NotUsed] = + messageAPI(serverSide, closeTimeout, None) + + private def messageAPI( + serverSide: Boolean, + closeTimeout: FiniteDuration, + shouldCompress: Option[Message => Boolean]) + : BidiFlow[FrameHandler.Output, Message, Message, FrameOutHandler.Input, NotUsed] = { /* Collects user-level API messages from MessageDataParts */ val collectMessage: Flow[MessageDataPart, Message, NotUsed] = Flow[MessageDataPart] @@ -188,7 +212,9 @@ private[http] object WebSocket { .named("ws-prepare-messages") def renderMessages: Flow[Message, FrameStart, NotUsed] = - MessageToFrameRenderer.create(serverSide) + shouldCompress + .map(MessageToFrameRenderer.create(serverSide, _)) + .getOrElse(MessageToFrameRenderer.create(serverSide)) .named("ws-render-messages") BidiFlow.fromGraph(GraphDSL.create() { implicit b => diff --git a/http-core/src/main/scala/org/apache/pekko/http/javadsl/model/ws/WebSocketUpgrade.scala b/http-core/src/main/scala/org/apache/pekko/http/javadsl/model/ws/WebSocketUpgrade.scala index ceb70296ac..04694d8df1 100644 --- a/http-core/src/main/scala/org/apache/pekko/http/javadsl/model/ws/WebSocketUpgrade.scala +++ b/http-core/src/main/scala/org/apache/pekko/http/javadsl/model/ws/WebSocketUpgrade.scala @@ -14,6 +14,7 @@ package org.apache.pekko.http.javadsl.model.ws import java.lang.{ Iterable => JIterable } +import java.util.function.Predicate import org.apache.pekko import pekko.http.javadsl.model.HttpResponse @@ -52,6 +53,20 @@ trait WebSocketUpgrade { def handleMessagesWith( handlerFlow: Graph[FlowShape[Message, Message], ? <: Any], compressionEnabled: Boolean): HttpResponse + /** + * Returns a response that can be used to answer a WebSocket handshake request. The connection will afterwards + * use the given handlerFlow to handle WebSocket messages from the client. + * + * If {@code permessage-deflate} is negotiated, {@code shouldCompress} is evaluated once for each outbound text or + * binary message. Returning {@code true} compresses that message and returning {@code false} sends it uncompressed. + * The filter does not affect inbound messages or whether compression is negotiated. + * + * @since 2.0.0 + */ + def handleMessagesWith( + handlerFlow: Graph[FlowShape[Message, Message], ? <: Any], shouldCompress: Predicate[Message]): HttpResponse = + handleMessagesWith(handlerFlow) + /** * Returns a response that can be used to answer a WebSocket handshake request. The connection will afterwards * use the given handlerFlow to handle WebSocket messages from the client. The given subprotocol must be one @@ -73,6 +88,23 @@ trait WebSocketUpgrade { subprotocol: String, compressionEnabled: Boolean): HttpResponse + /** + * Returns a response that can be used to answer a WebSocket handshake request. The connection will afterwards + * use the given handlerFlow to handle WebSocket messages from the client. The given subprotocol must be one + * of the ones offered by the client. + * + * If {@code permessage-deflate} is negotiated, {@code shouldCompress} is evaluated once for each outbound text or + * binary message. Returning {@code true} compresses that message and returning {@code false} sends it uncompressed. + * The filter does not affect inbound messages or whether compression is negotiated. + * + * @since 2.0.0 + */ + def handleMessagesWith( + handlerFlow: Graph[FlowShape[Message, Message], ? <: Any], + subprotocol: String, + shouldCompress: Predicate[Message]): HttpResponse = + handleMessagesWith(handlerFlow, subprotocol) + /** * Returns a response that can be used to answer a WebSocket handshake request. The connection will afterwards * use the given inSink to handle WebSocket messages from the client and the given outSource to send messages to the client. @@ -91,6 +123,23 @@ trait WebSocketUpgrade { def handleMessagesWith(inSink: Graph[SinkShape[Message], ? <: Any], outSource: Graph[SourceShape[Message], ? <: Any], compressionEnabled: Boolean): HttpResponse + /** + * Returns a response that can be used to answer a WebSocket handshake request. The connection will afterwards + * use the given inSink to handle WebSocket messages from the client and the given outSource to send messages to the + * client. + * + * If {@code permessage-deflate} is negotiated, {@code shouldCompress} is evaluated once for each outbound text or + * binary message. Returning {@code true} compresses that message and returning {@code false} sends it uncompressed. + * The filter does not affect inbound messages or whether compression is negotiated. + * + * @since 2.0.0 + */ + def handleMessagesWith( + inSink: Graph[SinkShape[Message], ? <: Any], + outSource: Graph[SourceShape[Message], ? <: Any], + shouldCompress: Predicate[Message]): HttpResponse = + handleMessagesWith(inSink, outSource) + /** * Returns a response that can be used to answer a WebSocket handshake request. The connection will afterwards * use the given inSink to handle WebSocket messages from the client and the given outSource to send messages to the client. @@ -111,4 +160,22 @@ trait WebSocketUpgrade { */ def handleMessagesWith(inSink: Graph[SinkShape[Message], ? <: Any], outSource: Graph[SourceShape[Message], ? <: Any], subprotocol: String, compressionEnabled: Boolean): HttpResponse + + /** + * Returns a response that can be used to answer a WebSocket handshake request. The connection will afterwards + * use the given inSink to handle WebSocket messages from the client and the given outSource to send messages to the + * client. The given subprotocol must be one of the ones offered by the client. + * + * If {@code permessage-deflate} is negotiated, {@code shouldCompress} is evaluated once for each outbound text or + * binary message. Returning {@code true} compresses that message and returning {@code false} sends it uncompressed. + * The filter does not affect inbound messages or whether compression is negotiated. + * + * @since 2.0.0 + */ + def handleMessagesWith( + inSink: Graph[SinkShape[Message], ? <: Any], + outSource: Graph[SourceShape[Message], ? <: Any], + subprotocol: String, + shouldCompress: Predicate[Message]): HttpResponse = + handleMessagesWith(inSink, outSource, subprotocol) } diff --git a/http-core/src/main/scala/org/apache/pekko/http/scaladsl/model/ws/WebSocketUpgrade.scala b/http-core/src/main/scala/org/apache/pekko/http/scaladsl/model/ws/WebSocketUpgrade.scala index 34de701c24..2f60e6c9fd 100644 --- a/http-core/src/main/scala/org/apache/pekko/http/scaladsl/model/ws/WebSocketUpgrade.scala +++ b/http-core/src/main/scala/org/apache/pekko/http/scaladsl/model/ws/WebSocketUpgrade.scala @@ -14,6 +14,7 @@ package org.apache.pekko.http.scaladsl.model.ws import java.lang.Iterable +import java.util.function.{ Predicate => JPredicate } import scala.collection.immutable import org.apache.pekko import pekko.NotUsed @@ -71,6 +72,36 @@ trait WebSocketUpgrade extends jm.ws.WebSocketUpgrade { compressionEnabled: Boolean): HttpResponse = handleMessages(handlerFlow, subprotocol) + /** + * The high-level interface to create a WebSocket server based on "messages". + * + * If `permessage-deflate` is negotiated, `shouldCompress` is evaluated once for each outbound text or binary + * message. Returning `true` compresses that message and returning `false` sends it uncompressed. The filter does not + * affect inbound messages or whether compression is negotiated. + * + * @since 2.0.0 + */ + def handleMessages( + handlerFlow: Graph[FlowShape[Message, Message], Any], + shouldCompress: Message => Boolean): HttpResponse = + handleMessages(handlerFlow, None, shouldCompress) + + /** + * The high-level interface to create a WebSocket server based on "messages". + * + * Optionally, a subprotocol out of the ones requested by the client can be chosen. If `permessage-deflate` is + * negotiated, `shouldCompress` is evaluated once for each outbound text or binary message. Returning `true` + * compresses that message and returning `false` sends it uncompressed. The filter does not affect inbound messages + * or whether compression is negotiated. + * + * @since 2.0.0 + */ + def handleMessages( + handlerFlow: Graph[FlowShape[Message, Message], Any], + subprotocol: Option[String], + shouldCompress: Message => Boolean): HttpResponse = + handleMessages(handlerFlow, subprotocol) + /** * The high-level interface to create a WebSocket server based on "messages". * @@ -108,6 +139,38 @@ trait WebSocketUpgrade extends jm.ws.WebSocketUpgrade { compressionEnabled: Boolean): HttpResponse = handleMessages(scaladsl.Flow.fromSinkAndSource(inSink, outSource), subprotocol, compressionEnabled) + /** + * The high-level interface to create a WebSocket server based on "messages". + * + * If `permessage-deflate` is negotiated, `shouldCompress` is evaluated once for each outbound text or binary + * message. Returning `true` compresses that message and returning `false` sends it uncompressed. The filter does not + * affect inbound messages or whether compression is negotiated. + * + * @since 2.0.0 + */ + def handleMessagesWithSinkSource( + inSink: Graph[SinkShape[Message], Any], + outSource: Graph[SourceShape[Message], Any], + shouldCompress: Message => Boolean): HttpResponse = + handleMessages(scaladsl.Flow.fromSinkAndSource(inSink, outSource), None, shouldCompress) + + /** + * The high-level interface to create a WebSocket server based on "messages". + * + * Optionally, a subprotocol out of the ones requested by the client can be chosen. If `permessage-deflate` is + * negotiated, `shouldCompress` is evaluated once for each outbound text or binary message. Returning `true` + * compresses that message and returning `false` sends it uncompressed. The filter does not affect inbound messages + * or whether compression is negotiated. + * + * @since 2.0.0 + */ + def handleMessagesWithSinkSource( + inSink: Graph[SinkShape[Message], Any], + outSource: Graph[SourceShape[Message], Any], + subprotocol: Option[String], + shouldCompress: Message => Boolean): HttpResponse = + handleMessages(scaladsl.Flow.fromSinkAndSource(inSink, outSource), subprotocol, shouldCompress) + import scala.jdk.CollectionConverters._ /** @@ -131,6 +194,16 @@ trait WebSocketUpgrade extends jm.ws.WebSocketUpgrade { compressionEnabled: Boolean): HttpResponse = handleMessages(JavaMapping.toScala(handlerFlow), None, compressionEnabled) + /** + * Java API + * + * @since 2.0.0 + */ + override def handleMessagesWith( + handlerFlow: Graph[FlowShape[jm.ws.Message, jm.ws.Message], ? <: Any], + shouldCompress: JPredicate[jm.ws.Message]): HttpResponse = + handleMessages(JavaMapping.toScala(handlerFlow), message => shouldCompress.test(message)) + /** * Java API */ @@ -149,6 +222,17 @@ trait WebSocketUpgrade extends jm.ws.WebSocketUpgrade { compressionEnabled: Boolean): HttpResponse = handleMessages(JavaMapping.toScala(handlerFlow), subprotocol = Some(subprotocol), compressionEnabled) + /** + * Java API + * + * @since 2.0.0 + */ + override def handleMessagesWith( + handlerFlow: Graph[FlowShape[jm.ws.Message, jm.ws.Message], ? <: Any], + subprotocol: String, + shouldCompress: JPredicate[jm.ws.Message]): HttpResponse = + handleMessages(JavaMapping.toScala(handlerFlow), Some(subprotocol), message => shouldCompress.test(message)) + /** * Java API */ @@ -167,6 +251,17 @@ trait WebSocketUpgrade extends jm.ws.WebSocketUpgrade { compressionEnabled: Boolean): HttpResponse = handleMessages(createScalaFlow(inSink, outSource), None, compressionEnabled) + /** + * Java API + * + * @since 2.0.0 + */ + override def handleMessagesWith( + inSink: Graph[SinkShape[jm.ws.Message], ? <: Any], + outSource: Graph[SourceShape[jm.ws.Message], ? <: Any], + shouldCompress: JPredicate[jm.ws.Message]): HttpResponse = + handleMessages(createScalaFlow(inSink, outSource), message => shouldCompress.test(message)) + /** * Java API * @@ -190,6 +285,18 @@ trait WebSocketUpgrade extends jm.ws.WebSocketUpgrade { compressionEnabled: Boolean): HttpResponse = handleMessages(createScalaFlow(inSink, outSource), subprotocol = Some(subprotocol), compressionEnabled) + /** + * Java API + * + * @since 2.0.0 + */ + override def handleMessagesWith( + inSink: Graph[SinkShape[jm.ws.Message], ? <: Any], + outSource: Graph[SourceShape[jm.ws.Message], ? <: Any], + subprotocol: String, + shouldCompress: JPredicate[jm.ws.Message]): HttpResponse = + handleMessages(createScalaFlow(inSink, outSource), Some(subprotocol), message => shouldCompress.test(message)) + private def createScalaFlow(inSink: Graph[SinkShape[jm.ws.Message], ? <: Any], outSource: Graph[SourceShape[jm.ws.Message], ? <: Any]): Graph[FlowShape[Message, Message], NotUsed] = JavaMapping.toScala(scaladsl.Flow.fromSinkAndSourceMat(inSink, outSource)(scaladsl.Keep.none): Graph[FlowShape[ diff --git a/http-core/src/test/scala/org/apache/pekko/http/impl/engine/ws/WebSocketServerSpec.scala b/http-core/src/test/scala/org/apache/pekko/http/impl/engine/ws/WebSocketServerSpec.scala index bffc62f9ad..ec29e5c068 100644 --- a/http-core/src/test/scala/org/apache/pekko/http/impl/engine/ws/WebSocketServerSpec.scala +++ b/http-core/src/test/scala/org/apache/pekko/http/impl/engine/ws/WebSocketServerSpec.scala @@ -17,11 +17,13 @@ import java.io.ByteArrayOutputStream import java.util.concurrent.CountDownLatch import java.util.concurrent.TimeUnit import java.util.concurrent.atomic.AtomicInteger +import java.util.function.Predicate import java.util.zip.Deflater import java.util.zip.Inflater import org.apache.pekko import pekko.actor.ActorSystem +import pekko.http.javadsl.model.ws.{ Message => JavaMessage, TextMessage => JavaTextMessage } import pekko.http.scaladsl.model.ws._ import pekko.http.scaladsl.model.AttributeKeys.webSocketUpgrade import pekko.stream.Materializer @@ -536,6 +538,131 @@ class WebSocketServerSpec extends PekkoSpecWithMaterializer("pekko.http.server.w } } + "select compression separately for each outbound message" in Utils.assertAllStagesStopped { + new TestSetup { + sendWebSocketRequest("Sec-WebSocket-Extensions: permessage-deflate\r\n") + + val request = expectRequest() + val upgrade = request.attribute(webSocketUpgrade) + val evaluated = new AtomicInteger() + val firstMessage = "compressed first server message" + val secondMessage = "plain" + val thirdMessage = "compressed third server message" + val response = upgrade.get.handleMessages( + Flow.fromSinkAndSource( + Sink.ignore, + Source(List( + TextMessage.Strict(firstMessage), + TextMessage.Strict(secondMessage), + TextMessage.Strict(thirdMessage)))), + None, + message => { + evaluated.incrementAndGet() + message.asInstanceOf[TextMessage.Strict].text.startsWith("compressed") + }) + responses.sendNext(response) + + expectResponseWithWipedDate( + """HTTP/1.1 101 Switching Protocols + |Upgrade: websocket + |Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo= + |Sec-WebSocket-Extensions: permessage-deflate + |Server: pekko-http/test + |Date: XXXX + |Connection: upgrade + | + |""") + + val firstPayload = expectCompressedFrame(Protocol.Opcode.Text, fin = true, rsv1 = true) + expectWSFrame(Protocol.Opcode.Text, ByteString(secondMessage), fin = true) + val thirdPayload = expectCompressedFrame(Protocol.Opcode.Text, fin = true, rsv1 = true) + inflatePerMessagesWithContext(firstPayload, thirdPayload).map(_.utf8String) shouldEqual + Seq(firstMessage, thirdMessage) + evaluated.get() shouldEqual 3 + expectWSCloseFrame(Protocol.CloseCodes.Regular) + + sendWSCloseFrame(Protocol.CloseCodes.Regular, mask = true) + closeNetworkInput() + expectNetworkClose() + } + } + + "select compression separately through the Java API" in Utils.assertAllStagesStopped { + new TestSetup { + sendWebSocketRequest("Sec-WebSocket-Extensions: permessage-deflate\r\n") + + val request = expectRequest() + val upgrade = request.attribute(webSocketUpgrade) + val response = upgrade.get.handleMessagesWith( + Flow.fromSinkAndSource( + Sink.foreach[JavaMessage](_ => ()), + Source(List[JavaMessage]( + JavaTextMessage.create("plain"), + JavaTextMessage.create("compressed Java server message")))), + new Predicate[JavaMessage] { + override def test(message: JavaMessage): Boolean = + message.asTextMessage.getStrictText.startsWith("compressed") + }) + responses.sendNext(response) + + expectResponseWithWipedDate( + """HTTP/1.1 101 Switching Protocols + |Upgrade: websocket + |Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo= + |Sec-WebSocket-Extensions: permessage-deflate + |Server: pekko-http/test + |Date: XXXX + |Connection: upgrade + | + |""") + + expectWSFrame(Protocol.Opcode.Text, ByteString("plain"), fin = true) + expectCompressedTextFrame("compressed Java server message") + expectWSCloseFrame(Protocol.CloseCodes.Regular) + + sendWSCloseFrame(Protocol.CloseCodes.Regular, mask = true) + closeNetworkInput() + expectNetworkClose() + } + } + + "not evaluate the outbound compression filter when permessage-deflate is not negotiated" in + Utils.assertAllStagesStopped { + new TestSetup { + sendWebSocketRequest("") + + val request = expectRequest() + val upgrade = request.attribute(webSocketUpgrade) + val evaluated = new AtomicInteger() + val response = upgrade.get.handleMessages( + Flow.fromSinkAndSource(Sink.ignore, Source.single(TextMessage.Strict("plain server message"))), + None, + _ => { + evaluated.incrementAndGet() + true + }) + responses.sendNext(response) + + expectResponseWithWipedDate( + """HTTP/1.1 101 Switching Protocols + |Upgrade: websocket + |Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo= + |Server: pekko-http/test + |Date: XXXX + |Connection: upgrade + | + |""") + + expectWSFrame(Protocol.Opcode.Text, ByteString("plain server message"), fin = true) + evaluated.get() shouldEqual 0 + expectWSCloseFrame(Protocol.CloseCodes.Regular) + + sendWSCloseFrame(Protocol.CloseCodes.Regular, mask = true) + closeNetworkInput() + expectNetworkClose() + } + } + "deflate empty outbound messages sent by the application" in Utils.assertAllStagesStopped { new TestSetup { sendWebSocketRequest("Sec-WebSocket-Extensions: permessage-deflate\r\n") @@ -718,6 +845,48 @@ class WebSocketServerSpec extends PekkoSpecWithMaterializer("pekko.http.server.w tracking.awaitAllEnded() } + "keep one uncompressed selection across a streamed outbound message" in Utils.assertAllStagesStopped { + new TestSetup { + sendWebSocketRequest("Sec-WebSocket-Extensions: permessage-deflate\r\n") + + val request = expectRequest() + val upgrade = request.attribute(webSocketUpgrade) + val evaluated = new AtomicInteger() + val response = upgrade.get.handleMessages( + Flow.fromSinkAndSource( + Sink.ignore, + Source.single(TextMessage(Source(List("streamed ", "plain ", "message"))))), + None, + _ => { + evaluated.incrementAndGet() + false + }) + responses.sendNext(response) + + expectResponseWithWipedDate( + """HTTP/1.1 101 Switching Protocols + |Upgrade: websocket + |Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo= + |Sec-WebSocket-Extensions: permessage-deflate + |Server: pekko-http/test + |Date: XXXX + |Connection: upgrade + | + |""") + + expectWSFrame(Protocol.Opcode.Text, ByteString("streamed "), fin = false) + expectWSFrame(Protocol.Opcode.Continuation, ByteString("plain "), fin = false) + expectWSFrame(Protocol.Opcode.Continuation, ByteString("message"), fin = false) + expectWSFrame(Protocol.Opcode.Continuation, ByteString.empty, fin = true) + evaluated.get() shouldEqual 1 + expectWSCloseFrame(Protocol.CloseCodes.Regular) + + sendWSCloseFrame(Protocol.CloseCodes.Regular, mask = true) + closeNetworkInput() + expectNetworkClose() + } + } + "fail invalid compressed messages with a protocol error" in Utils.assertAllStagesStopped { new TestSetup { sendWebSocketRequest("Sec-WebSocket-Extensions: permessage-deflate\r\n") @@ -1290,6 +1459,143 @@ class WebSocketServerSpec extends PekkoSpecWithMaterializer("pekko.http.server.w } } + "select compression separately for low-level outbound messages" in Utils.assertAllStagesStopped { + new TestSetup { + sendWebSocketRequest("Sec-WebSocket-Extensions: permessage-deflate\r\n") + + val request = expectRequest() + val upgrade = request.attribute(webSocketUpgrade) + val evaluated = new AtomicInteger() + val firstMessage = ByteString("compressed first low-level message") + val secondMessage = ByteString("plain") + val thirdMessage = ByteString("compressed third low-level message") + val response = upgrade.get.asInstanceOf[UpgradeToWebSocketLowLevel].handleFrames( + Flow.fromSinkAndSource( + Sink.ignore, + Source(List( + FrameEvent.fullFrame(Protocol.Opcode.Text, None, firstMessage, fin = true), + FrameEvent.fullFrame(Protocol.Opcode.Text, None, secondMessage, fin = true), + FrameEvent.fullFrame(Protocol.Opcode.Text, None, thirdMessage, fin = true)))), + None, + compressionEnabled = true, + shouldCompress = start => { + evaluated.incrementAndGet() + start.data.utf8String.startsWith("compressed") + }) + responses.sendNext(response) + + expectResponseWithWipedDate( + """HTTP/1.1 101 Switching Protocols + |Upgrade: websocket + |Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo= + |Sec-WebSocket-Extensions: permessage-deflate + |Server: pekko-http/test + |Date: XXXX + |Connection: upgrade + | + |""") + + val firstPayload = expectCompressedFrame(Protocol.Opcode.Text, fin = true, rsv1 = true) + expectWSFrame(Protocol.Opcode.Text, secondMessage, fin = true) + val thirdPayload = expectCompressedFrame(Protocol.Opcode.Text, fin = true, rsv1 = true) + inflatePerMessagesWithContext(firstPayload, thirdPayload) shouldEqual Seq(firstMessage, thirdMessage) + evaluated.get() shouldEqual 3 + netOut.expectComplete() + + sendWSCloseFrame(Protocol.CloseCodes.Regular, mask = true) + closeNetworkInput() + } + } + + "retain a low-level compression decision across continuation and control frames" in + Utils.assertAllStagesStopped { + new TestSetup { + sendWebSocketRequest("Sec-WebSocket-Extensions: permessage-deflate\r\n") + + val request = expectRequest() + val upgrade = request.attribute(webSocketUpgrade) + val evaluated = new AtomicInteger() + val response = upgrade.get.asInstanceOf[UpgradeToWebSocketLowLevel].handleFrames( + Flow.fromSinkAndSource( + Sink.ignore, + Source(List( + FrameEvent.fullFrame(Protocol.Opcode.Text, None, ByteString("plain "), fin = false), + FrameEvent.fullFrame(Protocol.Opcode.Ping, None, ByteString("ping"), fin = true), + FrameEvent.fullFrame(Protocol.Opcode.Continuation, None, ByteString("message"), fin = true)))), + None, + compressionEnabled = true, + shouldCompress = _ => { + evaluated.incrementAndGet() + false + }) + responses.sendNext(response) + + expectResponseWithWipedDate( + """HTTP/1.1 101 Switching Protocols + |Upgrade: websocket + |Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo= + |Sec-WebSocket-Extensions: permessage-deflate + |Server: pekko-http/test + |Date: XXXX + |Connection: upgrade + | + |""") + + expectWSFrame(Protocol.Opcode.Text, ByteString("plain "), fin = false) + expectWSFrame(Protocol.Opcode.Ping, ByteString("ping"), fin = true) + expectWSFrame(Protocol.Opcode.Continuation, ByteString("message"), fin = true) + evaluated.get() shouldEqual 1 + netOut.expectComplete() + + sendWSCloseFrame(Protocol.CloseCodes.Regular, mask = true) + closeNetworkInput() + } + } + + "pass every frame-data event through for an uncompressed low-level message" in + Utils.assertAllStagesStopped { + new TestSetup { + sendWebSocketRequest("Sec-WebSocket-Extensions: permessage-deflate\r\n") + + val request = expectRequest() + val upgrade = request.attribute(webSocketUpgrade) + val evaluated = new AtomicInteger() + val message = ByteString("plain split frame data") + val (firstPart, secondPart) = message.splitAt(8) + val response = upgrade.get.asInstanceOf[UpgradeToWebSocketLowLevel].handleFrames( + Flow.fromSinkAndSource( + Sink.ignore, + Source(List( + FrameStart(FrameHeader(Protocol.Opcode.Text, None, message.length, fin = true), firstPart), + FrameData(secondPart, lastPart = true)))), + None, + compressionEnabled = true, + shouldCompress = _ => { + evaluated.incrementAndGet() + false + }) + responses.sendNext(response) + + expectResponseWithWipedDate( + """HTTP/1.1 101 Switching Protocols + |Upgrade: websocket + |Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo= + |Sec-WebSocket-Extensions: permessage-deflate + |Server: pekko-http/test + |Date: XXXX + |Connection: upgrade + | + |""") + + expectWSFrame(Protocol.Opcode.Text, message, fin = true) + evaluated.get() shouldEqual 1 + netOut.expectComplete() + + sendWSCloseFrame(Protocol.CloseCodes.Regular, mask = true) + closeNetworkInput() + } + } + "apply compression to low-level frame handlers" in Utils.assertAllStagesStopped { new TestSetup { sendWebSocketRequest("Sec-WebSocket-Extensions: permessage-deflate\r\n") @@ -1493,6 +1799,14 @@ class WebSocketServerSpec extends PekkoSpecWithMaterializer("pekko.http.server.w } } + private def inflatePerMessagesWithContext(messages: ByteString*): Seq[ByteString] = { + val inflater = new Inflater(true) + try messages.map(message => inflateFrame(inflater, message ++ ByteString(0x00, 0x00, 0xFF.toByte, 0xFF.toByte))) + finally { + inflater.end() + } + } + private def inflateFrame(inflater: Inflater, data: ByteString): ByteString = { inflater.setInput(data.toArray) val output = new ByteArrayOutputStream() diff --git a/http-tests/src/test/scala/org/apache/pekko/http/scaladsl/server/directives/WebSocketDirectivesSpec.scala b/http-tests/src/test/scala/org/apache/pekko/http/scaladsl/server/directives/WebSocketDirectivesSpec.scala index ad993b9ba3..58f6c4da9d 100644 --- a/http-tests/src/test/scala/org/apache/pekko/http/scaladsl/server/directives/WebSocketDirectivesSpec.scala +++ b/http-tests/src/test/scala/org/apache/pekko/http/scaladsl/server/directives/WebSocketDirectivesSpec.scala @@ -13,9 +13,15 @@ package org.apache.pekko.http.scaladsl.server.directives +import java.util.function.Predicate + +import scala.collection.immutable + import org.apache.pekko -import pekko.http.scaladsl.model.StatusCodes -import pekko.http.scaladsl.model.headers.`Sec-WebSocket-Protocol` +import pekko.http.impl.engine.server.InternalCustomHeader +import pekko.http.scaladsl.model.AttributeKeys.webSocketUpgrade +import pekko.http.scaladsl.model.{ HttpRequest, HttpResponse, StatusCodes } +import pekko.http.scaladsl.model.headers.{ `Sec-WebSocket-Protocol`, Upgrade, UpgradeProtocol } import pekko.http.scaladsl.model.ws._ import pekko.http.scaladsl.server.{ ExpectedWebSocketRequestRejection, @@ -24,7 +30,7 @@ import pekko.http.scaladsl.server.{ UnsupportedWebSocketSubprotocolRejection } import pekko.http.scaladsl.testkit.WSProbe -import pekko.stream.OverflowStrategy +import pekko.stream.{ FlowShape, Graph, OverflowStrategy } import pekko.stream.scaladsl.{ Flow, Sink, Source } import pekko.util.ByteString @@ -49,6 +55,48 @@ class WebSocketDirectivesSpec extends RoutingSpec { wsClient.expectCompletion() } } + "pass the outbound compression selector to the WebSocket upgrade" in { + var selectedSubprotocol = Option.empty[Option[String]] + var compressionDecision = Option.empty[Boolean] + val request = recordingWebSocketRequest() { (subprotocol, shouldCompress) => + selectedSubprotocol = Some(subprotocol) + compressionDecision = Some(shouldCompress(TextMessage("compress"))) + } + + request ~> handleWebSocketMessages(Flow[Message], + { + case TextMessage.Strict("compress") => true + case _ => false + }) ~> check { + isWebSocketUpgrade shouldEqual true + } + + selectedSubprotocol shouldEqual Some(None) + compressionDecision shouldEqual Some(true) + } + "pass the Java outbound compression predicate and selected subprotocol to the WebSocket upgrade" in { + var selectedSubprotocol = Option.empty[Option[String]] + var compressionDecision = Option.empty[Boolean] + val request = recordingWebSocketRequest("echo") { (subprotocol, shouldCompress) => + selectedSubprotocol = Some(subprotocol) + compressionDecision = Some(shouldCompress(TextMessage("compress"))) + } + val shouldCompress = new Predicate[pekko.http.javadsl.model.ws.Message] { + override def test(message: pekko.http.javadsl.model.ws.Message): Boolean = message.isText + } + val javaRoute = pekko.http.javadsl.server.Directives.handleWebSocketMessagesForProtocol( + Flow[pekko.http.javadsl.model.ws.Message].asJava, + "echo", + shouldCompress) + + request ~> javaRoute.asScala ~> check { + isWebSocketUpgrade shouldEqual true + header[`Sec-WebSocket-Protocol`].get.protocols shouldEqual immutable.Seq("echo") + } + + selectedSubprotocol shouldEqual Some(Some("echo")) + compressionDecision shouldEqual Some(true) + } "choose subprotocol from offered ones" in { val wsClient = WSProbe() @@ -113,4 +161,33 @@ class WebSocketDirectivesSpec extends RoutingSpec { def echo: Flow[Message, Message, Any] = Flow[Message] .buffer(1, OverflowStrategy.backpressure) // needed because a noop flow hasn't any buffer that would start processing + + private def recordingWebSocketRequest(offeredProtocols: String*)( + onHandle: (Option[String], Message => Boolean) => Unit): HttpRequest = { + val upgrade = + new InternalCustomHeader("UpgradeToWebSocketTestHeader") with WebSocketUpgrade { + override def requestedProtocols: immutable.Seq[String] = offeredProtocols.toList + + override def handleMessages( + handlerFlow: Graph[FlowShape[Message, Message], Any], + subprotocol: Option[String]): HttpResponse = + throw new AssertionError("The selective-compression overload was not called") + + override def handleMessages( + handlerFlow: Graph[FlowShape[Message, Message], Any], + subprotocol: Option[String], + shouldCompress: Message => Boolean): HttpResponse = { + onHandle(subprotocol, shouldCompress) + HttpResponse( + StatusCodes.SwitchingProtocols, + headers = + Upgrade(UpgradeProtocol("websocket") :: Nil) :: + subprotocol.map(protocol => `Sec-WebSocket-Protocol`(protocol :: Nil)).toList) + } + } + + HttpRequest() + .addAttribute(webSocketUpgrade, upgrade) + .addHeader(upgrade) + } } diff --git a/http/src/main/scala/org/apache/pekko/http/javadsl/server/directives/WebSocketDirectives.scala b/http/src/main/scala/org/apache/pekko/http/javadsl/server/directives/WebSocketDirectives.scala index 71a0f8ecd2..d4a1c9097e 100644 --- a/http/src/main/scala/org/apache/pekko/http/javadsl/server/directives/WebSocketDirectives.scala +++ b/http/src/main/scala/org/apache/pekko/http/javadsl/server/directives/WebSocketDirectives.scala @@ -16,7 +16,7 @@ package directives import java.util.{ List => JList } import java.util.Optional -import java.util.function.{ Function => JFunction } +import java.util.function.{ Function => JFunction, Predicate } import org.apache.pekko import pekko.NotUsed @@ -59,6 +59,17 @@ abstract class WebSocketDirectives extends SecurityDirectives { D.handleWebSocketMessages(adapt(handler)) } + /** + * Handles WebSocket requests with the given handler and selectively compresses outbound messages for which + * {@code shouldCompress} returns {@code true} when {@code permessage-deflate} was negotiated. + * + * @since 2.0.0 + */ + def handleWebSocketMessages[T](handler: Flow[Message, Message, T], shouldCompress: Predicate[Message]): Route = + RouteAdapter { + D.handleWebSocketMessages(adapt(handler), message => shouldCompress.test(message)) + } + /** * Handles WebSocket requests with the given handler if the given subprotocol is offered in the request and * rejects other requests with an [[ExpectedWebSocketRequestRejection]] or an [[UnsupportedWebSocketSubprotocolRejection]]. @@ -68,6 +79,21 @@ abstract class WebSocketDirectives extends SecurityDirectives { D.handleWebSocketMessagesForProtocol(adapt(handler), subprotocol) } + /** + * Handles WebSocket requests with the given handler if the given subprotocol is offered and selectively compresses + * outbound messages for which {@code shouldCompress} returns {@code true} when {@code permessage-deflate} was + * negotiated. + * + * @since 2.0.0 + */ + def handleWebSocketMessagesForProtocol[T]( + handler: Flow[Message, Message, T], + subprotocol: String, + shouldCompress: Predicate[Message]): Route = + RouteAdapter { + D.handleWebSocketMessagesForProtocol(adapt(handler), subprotocol, message => shouldCompress.test(message)) + } + /** * Handles WebSocket requests with the given handler and rejects other requests with an * [[ExpectedWebSocketRequestRejection]]. @@ -84,6 +110,22 @@ abstract class WebSocketDirectives extends SecurityDirectives { D.handleWebSocketMessagesForOptionalProtocol(adapt(handler), subprotocol.asScala) } + /** + * Handles WebSocket requests with the given handler and selectively compresses outbound messages for which + * {@code shouldCompress} returns {@code true} when {@code permessage-deflate} was negotiated. + * + * @since 2.0.0 + */ + def handleWebSocketMessagesForOptionalProtocol[T]( + handler: Flow[Message, Message, T], + subprotocol: Optional[String], + shouldCompress: Predicate[Message]): Route = RouteAdapter { + D.handleWebSocketMessagesForOptionalProtocol( + adapt(handler), + subprotocol.asScala, + message => shouldCompress.test(message)) + } + private def adapt[T](handler: Flow[Message, Message, T]): scaladsl.Flow[s.Message, s.Message, NotUsed] = { scaladsl.Flow[s.Message].map(_.asJava).via(handler).map(_.asScala) } diff --git a/http/src/main/scala/org/apache/pekko/http/scaladsl/server/directives/WebSocketDirectives.scala b/http/src/main/scala/org/apache/pekko/http/scaladsl/server/directives/WebSocketDirectives.scala index faff9f10ea..5730ab5687 100644 --- a/http/src/main/scala/org/apache/pekko/http/scaladsl/server/directives/WebSocketDirectives.scala +++ b/http/src/main/scala/org/apache/pekko/http/scaladsl/server/directives/WebSocketDirectives.scala @@ -60,6 +60,18 @@ trait WebSocketDirectives { def handleWebSocketMessages(handler: Flow[Message, Message, Any]): Route = handleWebSocketMessagesForOptionalProtocol(handler, None) + /** + * Handles WebSocket requests with the given handler and selectively compresses outbound messages for which + * `shouldCompress` returns `true` when `permessage-deflate` was negotiated. + * + * @group websocket + * @since 2.0.0 + */ + def handleWebSocketMessages( + handler: Flow[Message, Message, Any], + shouldCompress: Message => Boolean): Route = + handleWebSocketMessagesForOptionalProtocol(handler, None, shouldCompress) + /** * Handles WebSocket requests with the given handler if the given subprotocol is offered in the request and * rejects other requests with an [[ExpectedWebSocketRequestRejection]] or an [[UnsupportedWebSocketSubprotocolRejection]]. @@ -69,6 +81,19 @@ trait WebSocketDirectives { def handleWebSocketMessagesForProtocol(handler: Flow[Message, Message, Any], subprotocol: String): Route = handleWebSocketMessagesForOptionalProtocol(handler, Some(subprotocol)) + /** + * Handles WebSocket requests with the given handler if the given subprotocol is offered and selectively compresses + * outbound messages for which `shouldCompress` returns `true` when `permessage-deflate` was negotiated. + * + * @group websocket + * @since 2.0.0 + */ + def handleWebSocketMessagesForProtocol( + handler: Flow[Message, Message, Any], + subprotocol: String, + shouldCompress: Message => Boolean): Route = + handleWebSocketMessagesForOptionalProtocol(handler, Some(subprotocol), shouldCompress) + /** * Handles WebSocket requests with the given handler and rejects other requests with an * [[ExpectedWebSocketRequestRejection]]. @@ -90,4 +115,22 @@ trait WebSocketDirectives { else reject(UnsupportedWebSocketSubprotocolRejection(subprotocol.get)) // None.forall == true } + + /** + * Handles WebSocket requests with the given handler and selectively compresses outbound messages for which + * `shouldCompress` returns `true` when `permessage-deflate` was negotiated. + * + * @group websocket + * @since 2.0.0 + */ + def handleWebSocketMessagesForOptionalProtocol( + handler: Flow[Message, Message, Any], + subprotocol: Option[String], + shouldCompress: Message => Boolean): Route = + extractWebSocketUpgrade { upgrade => + if (subprotocol.forall(sub => upgrade.requestedProtocols.exists(_.equalsIgnoreCase(sub)))) + complete(upgrade.handleMessages(handler, subprotocol, shouldCompress)) + else + reject(UnsupportedWebSocketSubprotocolRejection(subprotocol.get)) // None.forall == true + } }