Akka & Pekko backends

softwaremill · Jan 21, 2025 · 78fb6c3 · 78fb6c3
1 parent 2dda137
commit 78fb6c3
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 3 deletions.
diff --git a/akka-http-backend/src/main/scala/sttp/client4/akkahttp/AkkaHttpBackend.scala b/akka-http-backend/src/main/scala/sttp/client4/akkahttp/AkkaHttpBackend.scala
@@ -131,12 +131,19 @@ class AkkaHttpBackend private (
     val body = bodyFromAkka(
       r.response,
       responseMetadata,
-      wsFlow.map(Right(_)).getOrElse(Left(decodeAkkaResponse(hr, r.autoDecompressionEnabled)))
+      wsFlow
+        .map(Right(_))
+        .getOrElse(
+          Left(decodeAkkaResponse(limitPekkoResponseIfNeeded(hr, r.maxResponseBodyLength), r.autoDecompressionEnabled))
+        )
     )
 
     body.map(sttp.client4.Response(_, code, statusText, headers, Nil, r.onlyMetadata))
   }
 
+  private def limitPekkoResponseIfNeeded(response: HttpResponse, limit: Option[Long]): HttpResponse =
+    limit.fold(response)(l => response.withEntity(response.entity.withSizeLimit(l)))
+
   // http://doc.akka.io/docs/akka-http/10.0.7/scala/http/common/de-coding.html
   private def decodeAkkaResponse(response: HttpResponse, autoDecompressionEnabled: Boolean): HttpResponse =
     if (!response.status.allowsEntity() || !autoDecompressionEnabled) response

diff --git a/akka-http-backend/src/main/scala/sttp/client4/akkahttp/FromAkka.scala b/akka-http-backend/src/main/scala/sttp/client4/akkahttp/FromAkka.scala
@@ -1,6 +1,7 @@
 package sttp.client4.akkahttp
 
 import akka.http.scaladsl.model.HttpResponse
+import akka.http.scaladsl.model.EntityStreamSizeException
 import sttp.client4.{GenericRequest, SttpClientException}
 import sttp.model.{Header, HeaderNames}
 
@@ -25,6 +26,7 @@ private[akkahttp] object FromAkka {
           case _ => Some(new SttpClientException.ReadException(request, e))
         }
       case e: akka.stream.scaladsl.TcpIdleTimeoutException => Some(new SttpClientException.TimeoutException(request, e))
+      case e: EntityStreamSizeException                    => Some(new SttpClientException.ReadException(request, e))
       case e: Exception => SttpClientException.defaultExceptionToSttpClientException(request, e)
     }
 }
diff --git a/pekko-http-backend/src/main/scala/sttp/client4/pekkohttp/FromPekko.scala b/pekko-http-backend/src/main/scala/sttp/client4/pekkohttp/FromPekko.scala
@@ -2,6 +2,7 @@ package sttp.client4.pekkohttp
 
 import org.apache.pekko
 import pekko.http.scaladsl.model.HttpResponse
+import pekko.http.scaladsl.model.EntityStreamSizeException
 import sttp.client4.{GenericRequest, SttpClientException}
 import sttp.model.{Header, HeaderNames}
 
@@ -27,6 +28,7 @@ private[pekkohttp] object FromPekko {
         }
       case e: pekko.stream.scaladsl.TcpIdleTimeoutException =>
         Some(new SttpClientException.TimeoutException(request, e))
-      case e: Exception => SttpClientException.defaultExceptionToSttpClientException(request, e)
+      case e: EntityStreamSizeException => Some(new SttpClientException.ReadException(request, e))
+      case e: Exception                 => SttpClientException.defaultExceptionToSttpClientException(request, e)
     }
 }
diff --git a/pekko-http-backend/src/main/scala/sttp/client4/pekkohttp/PekkoHttpBackend.scala b/pekko-http-backend/src/main/scala/sttp/client4/pekkohttp/PekkoHttpBackend.scala
@@ -133,12 +133,19 @@ class PekkoHttpBackend private (
     val body = bodyFromPekko(
       r.response,
       responseMetadata,
-      wsFlow.map(Right(_)).getOrElse(Left(decodePekkoResponse(hr, r.autoDecompressionEnabled)))
+      wsFlow
+        .map(Right(_))
+        .getOrElse(
+          Left(decodePekkoResponse(limitPekkoResponseIfNeeded(hr, r.maxResponseBodyLength), r.autoDecompressionEnabled))
+        )
     )
 
     body.map(sttp.client4.Response(_, code, statusText, headers, Nil, r.onlyMetadata))
   }
 
+  private def limitPekkoResponseIfNeeded(response: HttpResponse, limit: Option[Long]): HttpResponse =
+    limit.fold(response)(l => response.withEntity(response.entity.withSizeLimit(l)))
+
   // http://doc.akka.io/docs/akka-http/10.0.7/scala/http/common/de-coding.html
   private def decodePekkoResponse(response: HttpResponse, enableAutoDecompression: Boolean): HttpResponse =
     if (!response.status.allowsEntity() || !enableAutoDecompression) response