Support for limiting streams (Pekko, Akka, fs2, zio-streams) (#350)

Author: kciesielski

akka/src/main/scala/sttp/capabilities/akka/AkkaStreams.scala

@@ -1,11 +1,23 @@
 package sttp.capabilities.akka
 
-import akka.stream.scaladsl.{Flow, Source}
+import akka.stream.scaladsl.Flow
+import akka.stream.scaladsl.Source
 import akka.util.ByteString
+import sttp.capabilities.StreamMaxLengthExceededException
 import sttp.capabilities.Streams
 
 trait AkkaStreams extends Streams[AkkaStreams] {
   override type BinaryStream = Source[ByteString, Any]
   override type Pipe[A, B] = Flow[A, B, Any]
 }
-object AkkaStreams extends AkkaStreams
+
+object AkkaStreams extends AkkaStreams {
+
+  def limitBytes(stream: Source[ByteString, Any], maxBytes: Long): Source[ByteString, Any] = {
+    stream
+      .limitWeighted(maxBytes)(_.length.toLong)
+      .mapError { case _: akka.stream.StreamLimitReachedException =>
+        StreamMaxLengthExceededException(maxBytes)
+      }
+  }
+}

akka/src/test/scala/sttp/capabilities/akka/AkkaStreamsTest.scala

@@ -0,0 +1,66 @@
+package sttp.capabilities.akka
+
+import akka.actor.ActorSystem
+import akka.stream.scaladsl._
+import akka.stream.testkit.scaladsl.TestSink
+import akka.util.ByteString
+import org.scalatest.BeforeAndAfterAll
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.should.Matchers
+import sttp.capabilities.StreamMaxLengthExceededException
+
+import scala.concurrent.Await
+import scala.concurrent.duration._
+
+class AkkaStreamsTest extends AnyFlatSpec with Matchers with BeforeAndAfterAll {
+
+  behavior of "AkkaStreams"
+  implicit lazy val system: ActorSystem = ActorSystem()
+
+  override def afterAll(): Unit = {
+    val _ = Await.result(system.terminate(), 10.seconds)
+  }
+
+  it should "Pass all bytes if limit is not exceeded" in {
+    // given
+    val inputByteCount = 8192
+    val maxBytes = 8192L
+
+    val iterator = Iterator.fill[Byte](inputByteCount)('5'.toByte)
+    val chunkSize = 256
+
+    val inputStream: Source[ByteString, Any] =
+      Source.fromIterator(() => iterator.grouped(chunkSize).map(group => ByteString(group.toArray)))
+
+    // when
+    val stream = AkkaStreams.limitBytes(inputStream, maxBytes)
+
+    // then
+    stream
+      .fold(0L)((acc, bs) => acc + bs.length)
+      .runWith(TestSink[Long]())
+      .request(1)
+      .expectNext(inputByteCount.toLong)
+      .expectComplete()
+  }
+
+  it should "Fail stream if limit is exceeded" in {
+    // given
+    val inputByteCount = 8192
+    val maxBytes = 8191L
+
+    val iterator = Iterator.fill[Byte](inputByteCount)('5'.toByte)
+    val chunkSize = 256
+
+    val inputStream: Source[ByteString, Any] =
+      Source.fromIterator(() => iterator.grouped(chunkSize).map(group => ByteString(group.toArray)))
+
+    // when
+    val stream = AkkaStreams.limitBytes(inputStream, maxBytes)
+    val probe = stream.runWith(TestSink[ByteString]())
+    val _ = for (_ <- 1 to 31) yield probe.requestNext()
+
+    // then
+    probe.request(1).expectError(StreamMaxLengthExceededException(maxBytes))
+  }
+}

build.sbt

@@ -9,7 +9,8 @@ val scala2_13 = "2.13.12"
 val scala2 = List(scala2_11, scala2_12, scala2_13)
 val scala2alive = List(scala2_12, scala2_13)
 val scala3 = List("3.3.1")
-
+val akkaVersion = "2.6.20"
+val pekkoVersion = "1.0.1"
 val sttpModelVersion = "1.6.0"
 
 val scalaTestVersion = "3.2.17"
@@ -39,6 +40,7 @@ val commonSettings = commonSmlBuildSettings ++ ossPublishSettings ++ Seq(
 val commonJvmSettings = commonSettings ++ Seq(
   scalacOptions ++= Seq("-target:jvm-1.8"),
   ideSkipProject := (scalaVersion.value != scala2_13),
+  bspEnabled := !ideSkipProject.value,
   mimaPreviousArtifacts := previousStableVersion.value.map(organization.value %% moduleName.value % _).toSet,
   mimaReportBinaryIssues := { if ((publish / skip).value) {} else mimaReportBinaryIssues.value }
 )
@@ -141,7 +143,10 @@ lazy val akka = (projectMatrix in file("akka"))
   .jvmPlatform(
     scalaVersions = scala2alive ++ scala3,
     settings = commonJvmSettings ++ Seq(
-      libraryDependencies += "com.typesafe.akka" %% "akka-stream" % "2.6.20" % "provided"
+      libraryDependencies ++= Seq(
+        "com.typesafe.akka" %% "akka-stream" % akkaVersion % "provided",
+        "com.typesafe.akka" %% "akka-stream-testkit" % akkaVersion % Test
+      )
     )
   )
   .dependsOn(core)
@@ -153,7 +158,10 @@ lazy val pekko = (projectMatrix in file("pekko"))
   .jvmPlatform(
     scalaVersions = scala2alive ++ scala3,
     settings = commonJvmSettings ++ Seq(
-      libraryDependencies += "org.apache.pekko" %% "pekko-stream" % "1.0.1" % "provided"
+      libraryDependencies ++= Seq(
+      "org.apache.pekko" %% "pekko-stream" % pekkoVersion % "provided",
+      "org.apache.pekko" %% "pekko-stream-testkit" % pekkoVersion % Test
+      )
     )
   )
   .dependsOn(core)
@@ -196,7 +204,10 @@ lazy val fs2 = (projectMatrix in file("fs2"))
   .jvmPlatform(
     scalaVersions = scala2alive ++ scala3,
     settings = commonJvmSettings ++ Seq(
-      libraryDependencies += "co.fs2" %% "fs2-io" % fs2_3_version
+      libraryDependencies ++= Seq(
+        "co.fs2" %% "fs2-io" % fs2_3_version,
+        "org.scalatest" %% "scalatest" % scalaTestVersion % Test
+      )
     )
   )
   .jsPlatform(

core/src/main/scala/sttp/capabilities/StreamMaxLimitExceededException.scala

@@ -0,0 +1,5 @@
+package sttp.capabilities
+
+case class StreamMaxLengthExceededException(maxBytes: Long) extends Exception {
+  override def getMessage(): String = s"Stream length limit of $maxBytes bytes exceeded"
+}

fs2/src/main/scala/sttp/capabilities/fs2/Fs2Streams.scala

@@ -1,12 +1,33 @@
 package sttp.capabilities.fs2
 
+import cats.MonadThrow
+import fs2.Pull
 import fs2.Stream
+import sttp.capabilities.StreamMaxLengthExceededException
 import sttp.capabilities.Streams
 
 trait Fs2Streams[F[_]] extends Streams[Fs2Streams[F]] {
   override type BinaryStream = Stream[F, Byte]
   override type Pipe[A, B] = fs2.Pipe[F, A, B]
 }
+
 object Fs2Streams {
   def apply[F[_]]: Fs2Streams[F] = new Fs2Streams[F] {}
+
+  def limitBytes[F[_]](stream: Stream[F, Byte], maxBytes: Long)(implicit mErr: MonadThrow[F]): Stream[F, Byte] = {
+    def go(s: Stream[F, Byte], remaining: Long): Pull[F, Byte, Unit] = {
+      if (remaining < 0) Pull.raiseError(new StreamMaxLengthExceededException(maxBytes))
+      else
+        s.pull.uncons.flatMap {
+          case Some((chunk, tail)) =>
+            val chunkSize = chunk.size.toLong
+            if (chunkSize <= remaining)
+              Pull.output(chunk) >> go(tail, remaining - chunkSize)
+            else
+              Pull.raiseError(new StreamMaxLengthExceededException(maxBytes))
+          case None => Pull.done
+        }
+    }
+    go(stream, maxBytes).stream
+  }
 }

fs2/src/test/scala/sttp/capabilities/fs2/Fs2StreamsTest.scala

@@ -0,0 +1,48 @@
+package sttp.capabilities.fs2
+
+import cats.effect.IO
+import cats.effect.unsafe.implicits.global
+import fs2._
+import org.scalatest.flatspec.AsyncFlatSpec
+import org.scalatest.matchers.should.Matchers
+import sttp.capabilities.StreamMaxLengthExceededException
+
+class Fs2StreamsTest extends AsyncFlatSpec with Matchers {
+  behavior of "Fs2Streams"
+
+  it should "Pass all bytes if limit is not exceeded" in {
+    // given
+    val inputByteCount = 8192
+    val maxBytes = 8192L
+    val inputStream = Stream.fromIterator[IO](Iterator.fill[Byte](inputByteCount)('5'.toByte), chunkSize = 1024)
+
+    // when
+    val stream = Fs2Streams.limitBytes(inputStream, maxBytes)
+
+    // then
+    stream.fold(0L)((acc, _) => acc + 1).compile.lastOrError.unsafeToFuture().map { count =>
+      count shouldBe inputByteCount
+    }
+  }
+
+  it should "Fail stream if limit is exceeded" in {
+    // given
+    val inputByteCount = 8192
+    val maxBytes = 8191L
+    val inputStream = Stream.fromIterator[IO](Iterator.fill[Byte](inputByteCount)('5'.toByte), chunkSize = 1024)
+
+    // when
+    val stream = Fs2Streams.limitBytes(inputStream, maxBytes)
+
+    // then
+    stream.compile.drain
+      .map(_ => fail("Unexpected end of stream."))
+      .handleErrorWith {
+        case StreamMaxLengthExceededException(limit) =>
+          IO(limit shouldBe maxBytes)
+        case other =>
+          IO(fail(s"Unexpected failure cause: $other"))
+      }
+      .unsafeToFuture()
+  }
+}

pekko/src/main/scala/sttp/capabilities/pekko/PekkoStreams.scala

@@ -1,12 +1,22 @@
 package sttp.capabilities.pekko
 
 import org.apache.pekko
+import sttp.capabilities.StreamMaxLengthExceededException
+import sttp.capabilities.Streams
+
 import pekko.stream.scaladsl.{Flow, Source}
 import pekko.util.ByteString
-import sttp.capabilities.Streams
 
 trait PekkoStreams extends Streams[PekkoStreams] {
   override type BinaryStream = Source[ByteString, Any]
   override type Pipe[A, B] = Flow[A, B, Any]
 }
-object PekkoStreams extends PekkoStreams
+object PekkoStreams extends PekkoStreams {
+  def limitBytes(stream: Source[ByteString, Any], maxBytes: Long): Source[ByteString, Any] = {
+    stream
+      .limitWeighted(maxBytes)(_.length.toLong)
+      .mapError { 
+        case _: pekko.stream.StreamLimitReachedException => StreamMaxLengthExceededException(maxBytes) 
+      }
+  }
+}

pekko/src/test/scala/sttp/capabilities/pekko/PekkoStreamsTest.scala

@@ -0,0 +1,61 @@
+package sttp.capabilities.pekko
+
+import org.apache.pekko.actor.ActorSystem
+import org.apache.pekko.stream.scaladsl._
+import org.apache.pekko.stream.testkit.scaladsl.TestSink
+import org.apache.pekko.util.ByteString
+import org.scalatest.BeforeAndAfterAll
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.should.Matchers
+import sttp.capabilities.StreamMaxLengthExceededException
+
+import scala.concurrent.Await
+import scala.concurrent.duration._
+
+class PekkoStreamsTest extends AnyFlatSpec with Matchers with BeforeAndAfterAll {
+
+  behavior of "PekkoStreams"
+  implicit lazy val system: ActorSystem = ActorSystem()
+
+  override def afterAll(): Unit = {
+    val _ = Await.result(system.terminate(), 10.seconds)
+  }
+
+  it should "Pass all bytes if limit is not exceeded" in {
+    // given
+    val inputByteCount = 8192
+    val maxBytes = 8192L
+
+    val iterator = Iterator.fill[Byte](inputByteCount)('5'.toByte)
+    val chunkSize = 256
+
+    val inputStream: Source[ByteString, Any] =
+      Source.fromIterator(() => iterator.grouped(chunkSize).map(group => ByteString(group.toArray)))
+
+    // when
+    val stream = PekkoStreams.limitBytes(inputStream, maxBytes)
+
+    // then
+    stream.fold(0L)((acc, bs) => acc + bs.length).runWith(TestSink[Long]()).request(1).expectNext(inputByteCount.toLong).expectComplete()
+  }
+
+  it should "Fail stream if limit is exceeded" in {
+    // given
+    val inputByteCount = 8192
+    val maxBytes = 8191L
+
+    val iterator = Iterator.fill[Byte](inputByteCount)('5'.toByte)
+    val chunkSize = 256
+
+    val inputStream: Source[ByteString, Any] =
+      Source.fromIterator(() => iterator.grouped(chunkSize).map(group => ByteString(group.toArray)))
+
+    // when
+    val stream = PekkoStreams.limitBytes(inputStream, maxBytes)    
+    val probe = stream.runWith(TestSink[ByteString]())
+    val _ = for (_ <- 1 to 31) yield probe.requestNext()
+
+    // then
+    probe.request(1).expectError(StreamMaxLengthExceededException(maxBytes))
+  }
+}

zio/src/main/scala/sttp/capabilities/zio/ZioStreams.scala

@@ -1,10 +1,49 @@
 package sttp.capabilities.zio
 
+import sttp.capabilities.StreamMaxLengthExceededException
 import sttp.capabilities.Streams
+import zio.Chunk
+import zio.Trace
 import zio.stream.Stream
+import zio.stream.ZChannel
+import zio.stream.ZStream
+
+import scala.util.control.NonFatal
 
 trait ZioStreams extends Streams[ZioStreams] {
   override type BinaryStream = Stream[Throwable, Byte]
   override type Pipe[A, B] = Stream[Throwable, A] => Stream[Throwable, B]
 }
-object ZioStreams extends ZioStreams
+
+object ZioStreams extends ZioStreams {
+
+  def limitBytes(stream: Stream[Throwable, Byte], maxBytes: Long): Stream[Throwable, Byte] =
+    scanChunksAccum(stream, initState = 0L) { case (accumulatedBytes, chunk) =>
+      val byteCount = accumulatedBytes + chunk.size
+      if (byteCount > maxBytes)
+        throw new StreamMaxLengthExceededException(maxBytes)
+      else
+        byteCount
+    }
+
+  private def scanChunksAccum[S, R, A](inputStream: ZStream[R, Throwable, A], initState: => S)(
+      f: (S, Chunk[A]) => S
+  )(implicit trace: Trace): ZStream[R, Throwable, A] =
+    ZStream.succeed(initState).flatMap { state =>
+      def accumulator(currS: S): ZChannel[Any, Throwable, Chunk[A], Any, Throwable, Chunk[A], Unit] =
+        ZChannel.readWith(
+          (in: Chunk[A]) => {
+            try {
+              val nextS = f(currS, in)
+              ZChannel.write(in) *> accumulator(nextS)
+            } catch {
+              case NonFatal(err) => ZChannel.fail(err)
+            }
+          },
+          (err: Throwable) => ZChannel.fail(err),
+          (_: Any) => ZChannel.unit
+        )
+
+      ZStream.fromChannel(inputStream.channel >>> accumulator(state))
+    }
+}