check content length against requestSizeLimit

Kitura · Aug 27, 2019 · 765ef49 · 765ef49
1 parent 7be264b
commit 765ef49
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 2 deletions.
diff --git a/Sources/KituraNet/HTTP/HTTPRequestHandler.swift b/Sources/KituraNet/HTTP/HTTPRequestHandler.swift
@@ -22,7 +22,6 @@ import Foundation
 import Dispatch
 
 internal class HTTPRequestHandler: ChannelInboundHandler, RemovableChannelHandler {
-
     /// The HTTPServer instance on which this handler is installed
     var server: HTTPServer
 
@@ -79,6 +78,13 @@ internal class HTTPRequestHandler: ChannelInboundHandler, RemovableChannelHandle
 
         switch request {
         case .head(let header):
+            if let requestSizeLimit = server.serverConfig.requestSizeLimit,
+                let contentLength = header.headers["Content-Length"].first,
+                    let contentLengthValue = Int(contentLength) {
+                    if contentLengthValue > requestSizeLimit {
+                        sendStatus(context: context)
+                    }
+            }
             serverRequest = HTTPServerRequest(ctx: context, requestHead: header, enableSSL: enableSSLVerification)
             self.clientRequestedKeepAlive = header.isKeepAlive
         case .body(var buffer):
@@ -152,4 +158,15 @@ internal class HTTPRequestHandler: ChannelInboundHandler, RemovableChannelHandle
     func updateKeepAliveState() {
         keepAliveState.decrement()
     }
+
+    func sendStatus(context: ChannelHandlerContext) {
+        let statusDescription = HTTP.statusCodes[HTTPStatusCode.requestTooLong.rawValue] ?? ""
+        do {
+            serverResponse = HTTPServerResponse(channel: context.channel, handler: self)
+            errorResponseSent = true
+            try serverResponse?.end(with: .requestTooLong, message: statusDescription)
+        } catch {
+            Log.error("Failed to send error response")
+        }
+    }
 }
diff --git a/Tests/KituraNetTests/ClientE2ETests.swift b/Tests/KituraNetTests/ClientE2ETests.swift
@@ -132,7 +132,6 @@ class ClientE2ETests: KituraNetTest {
     func testPostRequests() {
         performServerTest(delegate, asyncTasks: { expectation in
             self.performRequest("post", path: "/posttest", callback: {response in
-                print("HI")
                 XCTAssertEqual(response?.statusCode, HTTPStatusCode.OK, "Status code wasn't .Ok was \(String(describing: response?.statusCode))")
                 do {
                     let postValue = try response?.readString()