Merge pull request #4 from MahdiBM/mmbm-guard-against-multiple-iterat…

…ions Guard against multiple iterations
swift-server · Oct 6, 2023 · 7100ea1 · 7100ea1
2 parents d5ac556 + 8c619d9
commit 7100ea1
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 11 deletions.
diff --git a/Package.swift b/Package.swift
@@ -15,14 +15,16 @@ let package = Package(
     ],
     dependencies: [
         .package(url: "https://github.com/apple/swift-openapi-runtime.git", .upToNextMinor(from: "0.3.0")),
-        .package(url: "https://github.com/vapor/vapor.git", from: "4.0.0")
+        .package(url: "https://github.com/vapor/vapor.git", from: "4.84.4"),
+        .package(url: "https://github.com/apple/swift-atomics", from: "1.2.0")
     ],
     targets: [
         .target(
             name: "OpenAPIVapor",
             dependencies: [
                 .product(name: "Vapor", package: "vapor"),
-                .product(name: "OpenAPIRuntime", package: "swift-openapi-runtime")
+                .product(name: "OpenAPIRuntime", package: "swift-openapi-runtime"),
+                .product(name: "Atomics", package: "swift-atomics")
             ]
         ),
         .testTarget(

diff --git a/Sources/OpenAPIVapor/VaporTransport.swift b/Sources/OpenAPIVapor/VaporTransport.swift
@@ -17,6 +17,7 @@ import OpenAPIRuntime
 import HTTPTypes
 import Vapor
 import NIOFoundationCompat
+import Atomics
 
 public final class VaporTransport {
 
@@ -58,6 +59,7 @@ enum VaporTransportError: Error {
     case unsupportedHTTPMethod(String)
     case duplicatePathParameter([String])
     case missingRequiredPathParameter(String)
+    case multipleBodyIteration
 }
 
 extension [Vapor.PathComponent] {
@@ -143,21 +145,26 @@ extension Vapor.Response.Body {
             self = .empty
             return
         }
+        /// Used to guard the body from being iterated multiple times.
+        /// https://github.com/vapor/vapor/issues/3002
+        let iterated = ManagedAtomic(false)
         let stream: @Sendable (any Vapor.BodyStreamWriter) -> () = { writer in
+            guard iterated.compareExchange(
+                expected: false,
+                desired: true,
+                ordering: .relaxed
+            ).exchanged else {
+                _ = writer.write(.error(VaporTransportError.multipleBodyIteration))
+                return
+            }
             _ = writer.eventLoop.makeFutureWithTask {
                 do {
                     for try await chunk in body {
-                        try await writer.eventLoop.flatSubmit {
-                            writer.write(.buffer(ByteBuffer(bytes: chunk)))
-                        }.get()
+                        try await writer.write(.buffer(ByteBuffer(bytes: chunk))).get()
                     }
-                    try await writer.eventLoop.flatSubmit {
-                        writer.write(.end)
-                    }.get()
+                    try await writer.write(.end).get()
                 } catch {
-                    try await writer.eventLoop.flatSubmit {
-                        writer.write(.error(error))
-                    }.get()
+                    try await writer.write(.error(error)).get()
                 }
             }
         }