Skip to content

Commit

Permalink
feat: Ability to limit request size and connection count (#221)
Browse files Browse the repository at this point in the history
  • Loading branch information
RudraniW authored and djones6 committed Oct 8, 2019
1 parent 52b4235 commit 29b0dde
Show file tree
Hide file tree
Showing 8 changed files with 328 additions and 23 deletions.
2 changes: 1 addition & 1 deletion Package.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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"])
Expand Down
55 changes: 53 additions & 2 deletions Sources/KituraNet/HTTP/HTTPRequestHandler.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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?

Expand Down Expand Up @@ -66,7 +68,6 @@ internal class HTTPRequestHandler: ChannelInboundHandler, RemovableChannelHandle
self.enableSSLVerification = true
}
}

public typealias InboundIn = HTTPServerRequestPart
public typealias OutboundOut = HTTPServerResponsePart

Expand All @@ -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
Expand All @@ -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 {
Expand Down Expand Up @@ -152,4 +199,8 @@ internal class HTTPRequestHandler: ChannelInboundHandler, RemovableChannelHandle
func updateKeepAliveState() {
keepAliveState.decrement()
}

func channelInactive(context: ChannelHandlerContext, httpServer: HTTPServer) {
httpServer.connectionCount.sub(1)
}
}
16 changes: 13 additions & 3 deletions Sources/KituraNet/HTTP/HTTPServer.swift
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,9 @@ import SSLService
import LoggerAPI
import NIOWebSocket
import CLinuxHelpers
import Foundation
import NIOExtras
import NIOConcurrencyHelpers

#if os(Linux)
import Glibc
Expand Down Expand Up @@ -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
}

/**
Expand Down Expand Up @@ -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 {
Expand Down
120 changes: 120 additions & 0 deletions Sources/KituraNet/HTTP/ServerOptions.swift
Original file line number Diff line number Diff line change
@@ -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
}

}
27 changes: 27 additions & 0 deletions Tests/KituraNetTests/ClientE2ETests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ class ClientE2ETests: KituraNetTest {
("testQueryParameters", testQueryParameters),
("testRedirect", testRedirect),
("testPercentEncodedQuery", testPercentEncodedQuery),
("testRequestSize",testRequestSize),
]
}

Expand All @@ -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
Expand Down
Loading

0 comments on commit 29b0dde

Please sign in to comment.