feat: Ability to limit request size and connection count (#221)

Author: RudraniW

Package.swift

@@ -42,7 +42,7 @@ let package = Package(
             dependencies: []),
         .target(
             name: "KituraNet",
-            dependencies: ["NIO", "NIOFoundationCompat", "NIOHTTP1", "NIOSSL", "SSLService", "LoggerAPI", "NIOWebSocket", "CLinuxHelpers", "NIOExtras"]),
+            dependencies: ["NIO", "NIOFoundationCompat", "NIOHTTP1", "NIOSSL", "SSLService", "LoggerAPI", "NIOWebSocket", "CLinuxHelpers", "NIOConcurrencyHelpers", "NIOExtras"]),
         .testTarget(
             name: "KituraNetTests",
             dependencies: ["KituraNet"])

Sources/KituraNet/HTTP/HTTPRequestHandler.swift

@@ -26,6 +26,8 @@ internal class HTTPRequestHandler: ChannelInboundHandler, RemovableChannelHandle
     /// The HTTPServer instance on which this handler is installed
     var server: HTTPServer
 
+    var requestSize: Int = 0
+
     /// The serverRequest related to this handler instance
     var serverRequest: HTTPServerRequest?
 
@@ -66,7 +68,6 @@ internal class HTTPRequestHandler: ChannelInboundHandler, RemovableChannelHandle
             self.enableSSLVerification = true
         }
     }
-
     public typealias InboundIn = HTTPServerRequestPart
     public typealias OutboundOut = HTTPServerResponsePart
 
@@ -76,12 +77,42 @@ internal class HTTPRequestHandler: ChannelInboundHandler, RemovableChannelHandle
         // If an upgrade to WebSocket fails, both `errorCaught` and `channelRead` are triggered.
         // We'd want to return the error via `errorCaught`.
         if errorResponseSent { return }
-
         switch request {
         case .head(let header):
+            serverRequest = HTTPServerRequest(channel: context.channel, requestHead: header, enableSSL: enableSSLVerification)
+            if let requestSizeLimit = server.options.requestSizeLimit,
+                let contentLength = header.headers["Content-Length"].first,
+                let contentLengthValue = Int(contentLength) {
+                if contentLengthValue > requestSizeLimit {
+                    do {
+                        if let (httpStatus, response) = server.options.requestSizeResponseGenerator(requestSizeLimit, serverRequest?.remoteAddress ?? "") {
+                            serverResponse = HTTPServerResponse(channel: context.channel, handler: self)
+                            errorResponseSent = true
+                            try serverResponse?.end(with: httpStatus, message: response)
+                        }
+                    } catch {
+                        Log.error("Failed to send error response")
+                    }
+                    context.close()
+                }
+            }
             serverRequest = HTTPServerRequest(channel: context.channel, requestHead: header, enableSSL: enableSSLVerification)
             self.clientRequestedKeepAlive = header.isKeepAlive
         case .body(var buffer):
+            requestSize += buffer.readableBytes
+            if let requestSizeLimit = server.options.requestSizeLimit {
+                if requestSize > requestSizeLimit {
+                    do {
+                        if let (httpStatus, response) = server.options.requestSizeResponseGenerator(requestSizeLimit,serverRequest?.remoteAddress ?? "") {
+                            serverResponse = HTTPServerResponse(channel: context.channel, handler: self)
+                            errorResponseSent = true
+                            try serverResponse?.end(with: httpStatus, message: response)
+                        }
+                    } catch {
+                        Log.error("Failed to send error response")
+                    }
+                }
+            }
             guard let serverRequest = serverRequest else {
                 Log.error("No ServerRequest available")
                 return
@@ -91,7 +122,23 @@ internal class HTTPRequestHandler: ChannelInboundHandler, RemovableChannelHandle
             } else {
                 serverRequest.buffer!.byteBuffer.writeBuffer(&buffer)
             }
+
         case .end:
+            requestSize = 0
+            server.connectionCount.add(1)
+            if let connectionLimit = server.options.connectionLimit {
+                if server.connectionCount.load() > connectionLimit {
+                    do {
+                        if let (httpStatus, response) = server.options.connectionResponseGenerator(connectionLimit,serverRequest?.remoteAddress ?? "") {
+                            serverResponse = HTTPServerResponse(channel: context.channel, handler: self)
+                            errorResponseSent = true
+                            try serverResponse?.end(with: httpStatus, message: response)
+                        }
+                    } catch {
+                        Log.error("Failed to send error response")
+                    }
+                }
+            }
             serverResponse = HTTPServerResponse(channel: context.channel, handler: self)
             //Make sure we use the latest delegate registered with the server
             DispatchQueue.global().async {
@@ -152,4 +199,8 @@ internal class HTTPRequestHandler: ChannelInboundHandler, RemovableChannelHandle
     func updateKeepAliveState() {
         keepAliveState.decrement()
     }
+
+    func channelInactive(context: ChannelHandlerContext, httpServer: HTTPServer) {
+        httpServer.connectionCount.sub(1)
+    }
 }

Sources/KituraNet/HTTP/HTTPServer.swift

@@ -22,7 +22,9 @@ import SSLService
 import LoggerAPI
 import NIOWebSocket
 import CLinuxHelpers
+import Foundation
 import NIOExtras
+import NIOConcurrencyHelpers
 
 #if os(Linux)
 import Glibc
@@ -127,22 +129,30 @@ public class HTTPServer: Server {
 
     var quiescingHelper: ServerQuiescingHelper?
 
+    /// server configuration
+    public var options: ServerOptions = ServerOptions()
+
+    //counter for no of connections
+    var connectionCount = Atomic(value: 0)
+
     /**
      Creates an HTTP server object.
 
      ### Usage Example: ###
      ````swift
-     let server = HTTPServer()
+     let config =HTTPServerConfiguration(requestSize: 1000, coonectionLimit: 100)
+     let server = HTTPServer(serverconfig: config)
      server.listen(on: 8080)
      ````
     */
-    public init() {
+    public init(options: ServerOptions = ServerOptions()) {
 #if os(Linux)
         let numberOfCores = Int(linux_sched_getaffinity())
         self.eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: numberOfCores > 0 ? numberOfCores : System.coreCount)
 #else
         self.eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: System.coreCount)
 #endif
+        self.options = options
     }
 
     /**
@@ -309,7 +319,7 @@ public class HTTPServer: Server {
             }
             .childChannelInitializer { channel in
                 let httpHandler = HTTPRequestHandler(for: self)
-                let config: NIOHTTPServerUpgradeConfiguration = (upgraders: upgraders, completionHandler: { _ in
+                let config: HTTPUpgradeConfiguration = (upgraders: upgraders, completionHandler: {_ in 
                     _ = channel.pipeline.removeHandler(httpHandler)
                 })
                 return channel.pipeline.configureHTTPServerPipeline(withServerUpgrade: config, withErrorHandling: true).flatMap {

Sources/KituraNet/HTTP/ServerOptions.swift

@@ -0,0 +1,120 @@
+/*
+ * Copyright IBM Corporation 2019
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import Foundation
+import LoggerAPI
+
+/**
+ ServerOptions allows customization of default server policies, including:
+
+ - `requestSizeLimit`: Defines the maximum size for the body of an incoming request, in bytes. If a request body is larger than this limit, it will be rejected and the connection will be closed. A value of `nil` means no limit.
+ - `connectionLimit`: Defines the maximum number of concurrent connections that a server should accept. Clients attempting to connect when this limit has been reached will be rejected. A value of `nil` means no limit.
+
+ The server can optionally respond to the client with a message in either of these cases. This message can be customized by defining `requestSizeResponseGenerator` and `connectionResponseGenerator`.
+
+ Example usage:
+ ```
+ let server = HTTP.createServer()
+ server.options = ServerOptions(requestSizeLimit: 1000, connectionLimit: 10)
+ ```
+ */
+public struct ServerOptions {
+
+    /// A default limit of 100mb on the size of the request body that a server should accept.
+    public static let defaultRequestSizeLimit = 104857600
+
+    /// A default limit of 10,000 on the number of concurrent connections that a server should accept.
+    public static let defaultConnectionLimit = 10000
+
+    /// Defines a default response to an over-sized request of HTTP 413: Request Too Long. A message is also
+    /// logged at debug level.
+    public static let defaultRequestSizeResponseGenerator: (Int, String) -> (HTTPStatusCode, String)? = { (limit, clientSource) in
+        Log.debug("Request from \(clientSource) exceeds size limit of \(limit) bytes. Connection will be closed.")
+        return (.requestTooLong, "")
+    }
+
+    /// Defines a default response when refusing a new connection of HTTP 503: Service Unavailable. A message is
+    /// also logged at debug level.
+    public static let defaultConnectionResponseGenerator: (Int, String) -> (HTTPStatusCode, String)? = { (limit, clientSource) in
+        Log.debug("Rejected connection from \(clientSource): Maximum connection limit of \(limit) reached.")
+        return (.serviceUnavailable, "")
+    }
+
+    /// Defines the maximum size for the body of an incoming request, in bytes. If a request body is larger
+    /// than this limit, it will be rejected and the connection will be closed.
+    ///
+    /// A value of `nil` means no limit.
+    public let requestSizeLimit: Int?
+
+    /// Defines the maximum number of concurrent connections that a server should accept. Clients attempting
+    /// to connect when this limit has been reached will be rejected.
+    public let connectionLimit: Int?
+
+    /**
+     Determines the response message and HTTP status code used to respond to clients whose request exceeds
+     the `requestSizeLimit`. The current limit and client's address are provided as parameters to enable a
+     message to be logged, and/or a response to be provided back to the client.
+
+     The returned tuple indicates the HTTP status code and response body to send to the client. If `nil` is
+     returned, then no response will be sent.
+
+     Example usage:
+     ```
+     let oversizeResponse: (Int, String) -> (HTTPStatusCode, String)? = { (limit, client) in
+         Log.debug("Rejecting request from \(client): Exceeds limit of \(limit) bytes")
+         return (.requestTooLong, "Your request exceeds the limit of \(limit) bytes.\r\n")
+     }
+     ```
+     */
+    public let requestSizeResponseGenerator: (Int, String) -> (HTTPStatusCode, String)?
+
+    /**
+     Determines the response message and HTTP status code used to respond to clients that attempt to connect
+     while the server is already servicing the maximum number of connections, as defined by `connectionLimit`.
+     The current limit and client's address are provided as parameters to enable a  message to be logged,
+     and/or a response to be provided back to the client.
+
+     The returned tuple indicates the HTTP status code and response body to send to the client. If `nil` is
+     returned, then no response will be sent.
+
+     Example usage:
+     ```
+     let connectionResponse: (Int, String) -> (HTTPStatusCode, String)? = { (limit, client) in
+         Log.debug("Rejecting request from \(client): Connection limit \(limit) reached")
+         return (.serviceUnavailable, "Service busy - please try again later.\r\n")
+     }
+     ```
+     */
+    public let connectionResponseGenerator: (Int, String) -> (HTTPStatusCode, String)?
+
+    /// Create a `ServerOptions` to determine the behaviour of a `Server`.
+    ///
+    /// - parameter requestSizeLimit: The maximum size of an incoming request body. Defaults to `ServerOptions.defaultRequestSizeLimit`.
+    /// - parameter connectionLimit: The maximum number of concurrent connections. Defaults to `ServerOptions.defaultConnectionLimit`.
+    /// - parameter requestSizeResponseGenerator: A closure producing a response to send to a client when an over-sized request is rejected. Defaults to `ServerOptions.defaultRequestSizeResponseGenerator`.
+    /// - parameter defaultConnectionResponseGenerator: A closure producing a response to send to a client when a the server is busy and new connections are not being accepted. Defaults to `ServerOptions.defaultConnectionResponseGenerator`.
+    public init(requestSizeLimit: Int? = ServerOptions.defaultRequestSizeLimit,
+                connectionLimit: Int? = ServerOptions.defaultConnectionLimit,
+                requestSizeResponseGenerator: @escaping (Int, String) -> (HTTPStatusCode, String)? = ServerOptions.defaultRequestSizeResponseGenerator,
+                connectionResponseGenerator: @escaping (Int, String) -> (HTTPStatusCode, String)? = ServerOptions.defaultConnectionResponseGenerator)
+    {
+        self.requestSizeLimit = requestSizeLimit
+        self.connectionLimit = connectionLimit
+        self.requestSizeResponseGenerator = requestSizeResponseGenerator
+        self.connectionResponseGenerator = connectionResponseGenerator
+    }
+
+}

Tests/KituraNetTests/ClientE2ETests.swift

@@ -37,6 +37,7 @@ class ClientE2ETests: KituraNetTest {
             ("testQueryParameters", testQueryParameters),
             ("testRedirect", testRedirect),
             ("testPercentEncodedQuery", testPercentEncodedQuery),
+            ("testRequestSize",testRequestSize),
         ]
     }
 
@@ -52,6 +53,32 @@ class ClientE2ETests: KituraNetTest {
 
     let delegate = TestServerDelegate()
 
+    func testRequestSize() {
+        performServerTest(serverConfig: ServerOptions(requestSizeLimit: 10000, connectionLimit: 100),delegate, useSSL: false, asyncTasks: { expectation in
+            let payload = "[" + contentTypesString + "," + contentTypesString + contentTypesString + "," + contentTypesString + "]"
+            self.performRequest("post", path: "/largepost", callback: {response in
+                XCTAssertEqual(response?.statusCode, HTTPStatusCode.requestTooLong)
+                do {
+                    let expectedResult = ""
+                    var data = Data()
+                    let count = try response?.readAllData(into: &data)
+                    XCTAssertEqual(count, expectedResult.count, "Result should have been \(expectedResult.count) bytes, was \(String(describing: count)) bytes")
+                    let postValue = String(data: data, encoding: .utf8)
+                    if  let postValue = postValue {
+                        XCTAssertEqual(postValue, expectedResult)
+                    } else {
+                        XCTFail("postValue's value wasn't an UTF8 string")
+                    }
+                } catch {
+                    XCTFail("Failed reading the body of the response")
+                }
+                expectation.fulfill()
+            }) {request in
+                request.write(from: payload)
+            }
+        })
+    }
+
     func testHeadRequests() {
         performServerTest(delegate) { expectation in
             self.performRequest("head", path: "/headtest", callback: {response in