Default max body size for streaming request body collection

Sources/Vapor/Request/Request+Body.swift

@@ -42,6 +42,10 @@ extension Request {
                     return buffer
                 }
             case .collected(let buffer):
+                if let max = max, buffer.readableBytes > max {
+                    return self.request.eventLoop.future(error: Abort(.payloadTooLarge))
+                }
+
                 return self.request.eventLoop.makeSucceededFuture(buffer)
             case .none:
                 return self.request.eventLoop.makeSucceededFuture(nil)

Sources/Vapor/Routing/Routes.swift

@@ -1,11 +1,13 @@
 public final class Routes: RoutesBuilder, CustomStringConvertible {
+    public var defaultMaxBodySize: Int?
     public var all: [Route]
 
     public var description: String {
         return self.all.description
     }
-    
+
     public init() {
+        self.defaultMaxBodySize = 1 << 20 //1MB
         self.all = []
     }
 

Sources/Vapor/Routing/RoutesBuilder+Group.swift

@@ -13,7 +13,7 @@ extension RoutesBuilder {
     ///     - path: Group path components separated by commas.
     /// - returns: Newly created `Router` wrapped in the path.
     public func grouped(_ path: PathComponent...) -> RoutesBuilder {
-        return HTTPRoutesGroup(root: self, path: path)
+        return HTTPRoutesGroup(root: self, path: path, defaultMaxBodySize: self.defaultMaxBodySize)
     }
 
     /// Creates a new `Router` that will automatically prepend the supplied path components.
@@ -29,7 +29,36 @@ extension RoutesBuilder {
     ///     - path: Group path components separated by commas.
     ///     - configure: Closure to configure the newly created `Router`.
     public func group(_ path: PathComponent..., configure: (RoutesBuilder) throws -> ()) rethrows {
-        try configure(HTTPRoutesGroup(root: self, path: path))
+        try configure(HTTPRoutesGroup(root: self, path: path, defaultMaxBodySize: self.defaultMaxBodySize))
+    }
+
+    /// Creates a new `Router` with a different default max body size for its routes than the rest of the application.
+    ///
+    ///     let large = router.group(maxSize: 1_000_000)
+    ///     large.post("image", use: self.uploadImage)
+    ///
+    /// - parameters:
+    ///     - defaultMaxBodySize: The maximum number of bytes that a request body can contain in the new group
+    ///     `nil` means there is no limit.
+    /// - returns: A newly created `Router` with a new max body size.
+    public func grouped(defaultMaxBodySize: Int?) -> RoutesBuilder {
+        return HTTPRoutesGroup(root: self, defaultMaxBodySize: defaultMaxBodySize)
+    }
+
+    /// Creates a new `Router` with a different default max body size for its routes than the rest of the application.
+    ///
+    ///     router.grouped(maxSize: 1_000_000) { large
+    ///         large.post("image", use: self.uploadImage)
+    ///     }
+    ///
+    /// - parameters:
+    ///     - defaultMaxBodySize: The maximum number of bytes that a request body can contain in the new group
+    ///     `nil` means there is no limit.
+    ///     - configure: Closure to configure the newly created `Router`.
+    ///     - builder: The new builder with the new max body size.
+    /// - returns: A newly created `Router` with a new max body size.
+    public func group(defaultMaxBodySize: Int?, configure: (_ builder: RoutesBuilder) throws -> ()) rethrows {
+        try configure(HTTPRoutesGroup(root: self, defaultMaxBodySize: defaultMaxBodySize))
     }
 }
 
@@ -40,11 +69,15 @@ private final class HTTPRoutesGroup: RoutesBuilder {
 
     /// Additional components.
     let path: [PathComponent]
-
+
+    /// The default max body size for requests in the group.
+    let defaultMaxBodySize: Int?
+
     /// Creates a new `PathGroup`.
-    init(root: RoutesBuilder, path: [PathComponent]) {
+    init(root: RoutesBuilder, path: [PathComponent] = [], defaultMaxBodySize: Int? = nil) {
         self.root = root
         self.path = path
+        self.defaultMaxBodySize = defaultMaxBodySize
     }
 
     /// See `HTTPRoutesBuilder`.

Sources/Vapor/Routing/RoutesBuilder+Method.swift

@@ -2,11 +2,11 @@
 public enum HTTPBodyStreamStrategy {
     /// The HTTP request's body will be collected into memory before the route handler is
     /// called. The max size will determine how much data can be collected. The default is
-    /// `1 << 14`.
+    /// 1,000,000 bytes, or 1 megabyte.
     ///
     /// See `collect(maxSize:)` to set a lower max body size.
     public static var collect: HTTPBodyStreamStrategy {
-        return .collect(maxSize: 1 << 14)
+        return .collect(maxSize: nil)
     }
 
     /// The HTTP request's body will not be collected first before the route handler is called
@@ -18,7 +18,7 @@ public enum HTTPBodyStreamStrategy {
     ///
     /// If a `maxSize` is supplied, the request body size in bytes will be limited. Requests
     /// exceeding that size will result in an error.
-    case collect(maxSize: Int?)
+    case collect(maxSize: ByteCount?)
 }
 
 extension RoutesBuilder {
@@ -145,16 +145,28 @@ extension RoutesBuilder {
     ) -> Route
         where Response: ResponseEncodable
     {
+        let streamStrategy = body
         let responder = BasicResponder { request in
-            if case .collect(let max) = body, request.body.data == nil {
-                return request.body.collect(max: max).flatMapThrowing { _ in
+            switch streamStrategy {
+            case let .collect(maxSize):
+                let collected: EventLoopFuture<Void>
+
+                if let data = request.body.data {
+                    collected = request.eventLoop.tryFuture {
+                        if let max = maxSize, data.readableBytes > max.value { throw Abort(.payloadTooLarge) }
+                    }
+                } else {
+                    collected = request.body.collect(max: maxSize?.value).transform(to: ())
+                }
+
+                return collected.flatMapThrowing {
                     return try closure(request)
                 }.encodeResponse(for: request)
-            } else {
-                return try closure(request)
-                    .encodeResponse(for: request)
+            case .stream:
+                return try closure(request).encodeResponse(for: request)
             }
         }
+
         let route = Route(
             method: method,
             path: path,

Sources/Vapor/Routing/RoutesBuilder.swift

@@ -1,7 +1,13 @@
 public protocol RoutesBuilder {
+    var defaultMaxBodySize: Int? { get }
+
     func add(_ route: Route)
 }
 
+extension RoutesBuilder {
+    public var defaultMaxBodySize: Int? { 1_000_000 }
+}
+
 extension UUID: LosslessStringConvertible {
     public init?(_ description: String) {
         self.init(uuidString: description)

Sources/Vapor/Utilities/ByteCount.swift

@@ -0,0 +1,71 @@
+//
+//  ByteCount.swift
+//  
+//
+//  Created by Jimmy McDermott on 4/24/20.
+//
+
+import Foundation
+
+/// Represents a number of bytes:
+///
+/// let bytes: ByteCount = "1mb"
+/// print(bytes.value) // 1048576
+///
+/// let bytes: ByteCount = 1_000_000
+/// print(bytes.value) // 1000000
+
+/// let bytes: ByteCount = "2kib"
+/// print(bytes.value) // 2048
+public struct ByteCount {
+
+    /// The value in Bytes
+    public let value: Int
+}
+
+extension ByteCount: ExpressibleByIntegerLiteral {
+
+    /// Initializes the `ByteCount` with the raw byte count
+    /// - Parameter value: The number of bytes
+    public init(integerLiteral value: Int) {
+         self.value = value
+    }
+}
+
+extension ByteCount: ExpressibleByStringLiteral {
+
+    /// Initializes the `ByteCount` via a descriptive string. Available suffixes are:
+    /// `kib`, `mb`, `gb`, `tb`
+    /// - Parameter value: The string value (`1mb`)
+    public init(stringLiteral value: String) {
+        // Short path if it's an int wrapped in a string
+        if let intValue = Int(value) {
+            self.value = intValue
+            return
+        }
+
+        let validSuffixes = [
+            "kib": 10,
+            "mb": 20,
+            "gb": 30,
+            "tb": 40
+        ]
+
+        for suffix in validSuffixes {
+            guard value.hasSuffix(suffix.key) else { continue }
+            guard let stringIntValue = value.components(separatedBy: suffix.key).first else {
+                fatalError("Invalid string format")
+            }
+
+            guard let intValue = Int(stringIntValue) else {
+                fatalError("Invalid int value: \(stringIntValue)")
+            }
+
+            self.value = intValue << suffix.value
+            return
+        }
+
+        // Assert failure here because all cases are handled in the above loop
+        fatalError("Could not parse byte count string: \(value)")
+    }
+}

Tests/VaporTests/RouteTests.swift

@@ -268,4 +268,59 @@ final class RouteTests: XCTestCase {
             XCTAssertEqual(res.body.string, "bar")
         }
     }
+
+    func testConfigurableMaxBodySize() throws {
+        let app = Application(.testing)
+        defer { app.shutdown() }
+
+        XCTAssertEqual(app.routes.defaultMaxBodySize, 1048576)
+        app.routes.defaultMaxBodySize = 50
+        XCTAssertEqual(app.routes.defaultMaxBodySize, 50)
+
+        let group = app.routes.grouped(defaultMaxBodySize: 1_000_000_000).grouped("uploads")
+        XCTAssertEqual(group.defaultMaxBodySize, 1_000_000_000)
+
+        group.on(.POST, "small", body: .collect(maxSize: 1_000) , use: { request in return "small" })
+        group.on(.POST, "medium", body: .collect(maxSize: 1_000_000) , use: { request in return "medium" })
+        group.post("large", use: { request in return "large" })
+
+
+        var smallBuffer = ByteBufferAllocator().buffer(capacity: 1_001)
+        smallBuffer.writeBytes(Array(repeating: 48, count: 1_000))
+        try app.test(.POST, "/uploads/small", body: smallBuffer, afterResponse: { response in
+            XCTAssertEqual(response.status, .ok)
+        })
+
+        smallBuffer.clear()
+        smallBuffer.writeBytes(Array(repeating: 48, count: 1_001))
+        try app.test(.POST, "/uploads/small", body: smallBuffer, afterResponse: { response in
+            XCTAssertEqual(response.status, .payloadTooLarge)
+        })
+
+
+        var mediumBuffer = ByteBufferAllocator().buffer(capacity: 1_000_001)
+        mediumBuffer.writeBytes(Array(repeating: 48, count: 1_000_000))
+        try app.test(.POST, "/uploads/medium", body: mediumBuffer, afterResponse: { response in
+            XCTAssertEqual(response.status, .ok)
+        })
+
+        mediumBuffer.clear()
+        mediumBuffer.writeBytes(Array(repeating: 48, count: 1_000_001))
+        try app.test(.POST, "/uploads/medium", body: mediumBuffer, afterResponse: { response in
+            XCTAssertEqual(response.status, .payloadTooLarge)
+        })
+
+
+        var largeBuffer = ByteBufferAllocator().buffer(capacity: 1_000_000_001)
+        largeBuffer.writeBytes(Array(repeating: 48, count: 1_000_000_000))
+        try app.test(.POST, "/uploads/large", body: largeBuffer, afterResponse: { response in
+            XCTAssertEqual(response.status, .ok)
+        })
+
+        largeBuffer.clear()
+        largeBuffer.writeBytes(Array(repeating: 48, count: 1_000_000_001))
+        try app.test(.POST, "/uploads/large", body: largeBuffer, afterResponse: { response in
+            XCTAssertEqual(response.status, .payloadTooLarge)
+        })
+    }
 }

Tests/VaporTests/UtilityTests.swift

@@ -21,4 +21,21 @@ final class UtilityTests: XCTestCase {
         XCTAssertEqual(data.base32EncodedString(), "AEBAGBA")
         XCTAssertEqual(Data(base32Encoded: "AEBAGBA"), data)
     }
+
+    func testByteCount() throws {
+        let twoKib: ByteCount = "2kib"
+        XCTAssertEqual(twoKib.value, 2_048)
+
+        let oneMb: ByteCount = "1mb"
+        XCTAssertEqual(oneMb.value, 1_048_576)
+
+        let oneGb: ByteCount = "1gb"
+        XCTAssertEqual(oneGb.value, 1_073_741_824)
+
+        let oneTb: ByteCount = "1tb"
+        XCTAssertEqual(oneTb.value, 1_099_511_627_776)
+
+        let intBytes: ByteCount = 1_000_000
+        XCTAssertEqual(intBytes.value, 1_000_000)
+    }
 }