check body length

Author: artemredkin

Sources/AsyncHTTPClient/HTTPClient.swift

@@ -925,6 +925,7 @@ public struct HTTPClientError: Error, Equatable, CustomStringConvertible {
         case uncleanShutdown
         case traceRequestWithBody
         case invalidHeaderFieldNames([String])
+        case bodyLengthMismatch
     }
 
     private var code: Code
@@ -969,10 +970,12 @@ public struct HTTPClientError: Error, Equatable, CustomStringConvertible {
     public static let redirectLimitReached = HTTPClientError(code: .redirectLimitReached)
     /// Redirect Cycle detected.
     public static let redirectCycleDetected = HTTPClientError(code: .redirectCycleDetected)
-    /// Unclean shutdown
+    /// Unclean shutdown.
     public static let uncleanShutdown = HTTPClientError(code: .uncleanShutdown)
-    /// A body was sent in a request with method TRACE
+    /// A body was sent in a request with method TRACE.
     public static let traceRequestWithBody = HTTPClientError(code: .traceRequestWithBody)
-    /// Header field names contain invalid characters
+    /// Header field names contain invalid characters.
     public static func invalidHeaderFieldNames(_ names: [String]) -> HTTPClientError { return HTTPClientError(code: .invalidHeaderFieldNames(names)) }
+    /// Body length is not equal to `Content-Length`.
+    public static let bodyLengthMismatch = HTTPClientError(code: .bodyLengthMismatch)
 }

Sources/AsyncHTTPClient/HTTPHandler.swift

@@ -651,6 +651,8 @@ internal class TaskHandler<Delegate: HTTPClientResponseDelegate>: RemovableChann
     let logger: Logger // We are okay to store the logger here because a TaskHandler is just for one request.
 
     var state: State = .idle
+    var expectedBodyLength: Int? = nil
+    var actualBodyLength: Int = 0
     var pendingRead = false
     var mayRead = true
     var closing = false {
@@ -794,11 +796,19 @@ extension TaskHandler: ChannelDuplexHandler {
         assert(head.version == HTTPVersion(major: 1, minor: 1),
                "Sending a request in HTTP version \(head.version) which is unsupported by the above `if`")
 
+        self.expectedBodyLength = head.headers[canonicalForm: "content-length"].first.flatMap { Int($0) }
+
         context.write(wrapOutboundOut(.head(head))).map {
             self.callOutToDelegateFireAndForget(value: head, self.delegate.didSendRequestHead)
         }.flatMap {
             self.writeBody(request: request, context: context)
         }.flatMap {
+            if let expectedBodyLength = self.expectedBodyLength, expectedBodyLength != self.actualBodyLength {
+                self.state = .end
+                let error = HTTPClientError.bodyLengthMismatch
+                self.failTaskAndNotifyDelegate(error: error, self.delegate.didReceiveError)
+                return context.eventLoop.makeFailedFuture(error)
+            }
             context.eventLoop.assertInEventLoop()
             return context.writeAndFlush(self.wrapOutboundOut(.end(nil)))
         }.map {
@@ -836,6 +846,7 @@ extension TaskHandler: ChannelDuplexHandler {
                 }
 
                 return promise.futureResult.map {
+                    self.actualBodyLength += part.readableBytes
                     self.callOutToDelegateFireAndForget(value: part, self.delegate.didSendRequestPart)
                 }
             })

Tests/AsyncHTTPClientTests/HTTPClientTests.swift

@@ -2051,4 +2051,35 @@ class HTTPClientTests: XCTestCase {
 
         XCTAssertNoThrow(try future.wait())
     }
+
+    func testContentLengthTooLongFails() {
+        let url = self.defaultHTTPBinURLPrefix + "/post"
+        XCTAssertThrowsError(
+            try self.defaultClient.execute(request:
+                                            Request(url: url,
+                                                    body: .stream(length: 10) { streamWriter in
+                                                        streamWriter.write(.byteBuffer(ByteBuffer(string: "1")))
+                                                    })).wait()) { error in
+            XCTAssertEqual(error as! HTTPClientError, HTTPClientError.bodyLengthMismatch)
+        }
+        // Quickly try another request and check that it works.
+        XCTAssertNoThrow(try self.defaultClient.get(url: self.defaultHTTPBinURLPrefix + "/get").wait())
+    }
+
+    // currently gets stuck because of #250 the server just never replies
+    func testContentLengthTooShortFails() {
+        let url = self.defaultHTTPBinURLPrefix + "/post"
+        let tooLong = "XBAD BAD BAD NOT HTTP/1.1\r\n\r\n"
+        XCTAssertThrowsError(
+            try self.defaultClient.execute(request:
+                                            Request(url: url,
+                                                    body: .stream(length: 1) { streamWriter in
+                                                        streamWriter.write(.byteBuffer(ByteBuffer(string: tooLong)))
+                                                    })).wait()) { error in
+            XCTAssertEqual(error as! HTTPClientError, HTTPClientError.bodyLengthMismatch)
+        }
+        // Quickly try another request and check that it works. If we by accident wrote some extra bytes into the
+        // stream (and reuse the connection) that could cause problems.
+        XCTAssertNoThrow(try self.defaultClient.get(url: self.defaultHTTPBinURLPrefix + "/get").wait())
+    }
 }