Implementation of HTTPServerConfigurationHandler and RequestSizeTests

Kitura · Aug 27, 2019 · 204d34c · 204d34c
1 parent 9710e0a
commit 204d34c
Show file tree

Hide file tree

Showing 8 changed files with 52 additions and 23 deletions.
diff --git a/Package.swift b/Package.swift
@@ -32,7 +32,7 @@ let package = Package(
         .package(url: "https://github.com/apple/swift-nio-ssl.git", from: "2.0.0"),
         .package(url: "https://github.com/apple/swift-nio-extras.git", from: "1.0.0"),
         .package(url: "https://github.com/IBM-Swift/BlueSSLService.git", from: "1.0.0"),
-        .package(url: "https://github.com/IBM-Swift/LoggerAPI.git", from: "1.7.3")
+        .package(url: "https://github.com/IBM-Swift/LoggerAPI.git", from: "1.7.3"),
     ],
     targets: [
         // Targets are the basic building blocks of a package. A target can define a module or a test suite.
@@ -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", "NIOExtras", "NIOConcurrencyHelpers"]),
         .testTarget(
             name: "KituraNetTests",
             dependencies: ["KituraNet"])

diff --git a/Sources/KituraNet/HTTP/HTTP.swift b/Sources/KituraNet/HTTP/HTTP.swift
@@ -87,8 +87,9 @@ public class HTTP {
     let server = HTTP.createServer()
     ````
     */
-    public static func createServer() -> HTTPServer {
-        return HTTPServer()
+    public static func createServer(serverConfig: HTTPServerConfiguration = .default) -> HTTPServer {
+        let serverConfig = serverConfig
+        return HTTPServer(serverConfig: serverConfig)
     }
 
     /**

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(channel: context.channel, 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/Sources/KituraNet/HTTP/HTTPServer.swift b/Sources/KituraNet/HTTP/HTTPServer.swift
@@ -23,6 +23,8 @@ import LoggerAPI
 import NIOWebSocket
 import CLinuxHelpers
 import NIOExtras
+import Foundation
+import NIOConcurrencyHelpers
 
 #if os(Linux)
 import Glibc
@@ -127,6 +129,12 @@ public class HTTPServer: Server {
 
     var quiescingHelper: ServerQuiescingHelper?
 
+    /// server configuration
+    public var serverConfig: HTTPServerConfiguration
+
+    //counter for no of connections
+    var connectionCount = Atomic(value: 0)
+
     /**
      Creates an HTTP server object.
 
@@ -136,13 +144,14 @@ public class HTTPServer: Server {
      server.listen(on: 8080)
      ````
     */
-    public init() {
+    public init(serverConfig: HTTPServerConfiguration = .default) {
 #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.serverConfig = serverConfig
     }
 
     /**
@@ -309,13 +318,15 @@ public class HTTPServer: Server {
             }
             .childChannelInitializer { channel in
                 let httpHandler = HTTPRequestHandler(for: self)
+                let serverConfigurationHandler = HTTPServerConfigurationHandler(for: self)
                 let config: NIOHTTPServerUpgradeConfiguration = (upgraders: upgraders, completionHandler: { _ in
                     _ = channel.pipeline.removeHandler(httpHandler)
                 })
                 return channel.pipeline.configureHTTPServerPipeline(withServerUpgrade: config, withErrorHandling: true).flatMap {
                     if let nioSSLServerHandler = self.createNIOSSLServerHandler() {
                         _ = channel.pipeline.addHandler(nioSSLServerHandler, position: .first)
                     }
+                    _ = channel.pipeline.addHandler(serverConfigurationHandler, position: .first)
                     return channel.pipeline.addHandler(httpHandler)
                 }
             }

diff --git a/Tests/KituraNetTests/HTTPResponseTests.swift b/Tests/KituraNetTests/HTTPResponseTests.swift
@@ -40,7 +40,7 @@ class HTTPResponseTests: KituraNetTest {
         let headers = HeadersContainer()
 
         headers.append("Content-Type", value: "text/html")
-        var values = headers["Content-Type"]
+        let values = headers["Content-Type"]
         XCTAssertNotNil(values, "Couldn't retrieve just set Content-Type header")
         XCTAssertEqual(values?.count, 1, "Content-Type header should only have one value")
         XCTAssertEqual(values?[0], "text/html")

diff --git a/Tests/KituraNetTests/KituraNIOTest.swift b/Tests/KituraNetTests/KituraNIOTest.swift
@@ -85,9 +85,9 @@ class KituraNetTest: XCTestCase {
         }
     }
 
-    func startServer(_ delegate: ServerDelegate?, unixDomainSocketPath: String? = nil, port: Int = portDefault, useSSL: Bool = useSSLDefault, allowPortReuse: Bool = portReuseDefault) throws -> HTTPServer {
-
-        let server = HTTP.createServer()
+    func startServer(_ delegate: ServerDelegate?, unixDomainSocketPath: String? = nil, port: Int = portDefault, useSSL: Bool = useSSLDefault, allowPortReuse: Bool = portReuseDefault, serverConfig: HTTPServerConfiguration = .default) throws -> HTTPServer {
+        let serverConfig = serverConfig
+        let server = HTTP.createServer(serverConfig: serverConfig)
         server.delegate = delegate
         if useSSL {
             server.sslConfig = KituraNetTest.sslConfig
@@ -103,8 +103,9 @@ class KituraNetTest: XCTestCase {
 
     /// Convenience function for starting an HTTPServer on an ephemeral port,
     /// returning the a tuple containing the server and the port it is listening on.
-    func startEphemeralServer(_ delegate: ServerDelegate?, useSSL: Bool = useSSLDefault, allowPortReuse: Bool = portReuseDefault) throws -> (server: HTTPServer, port: Int) {
-        let server = try startServer(delegate, port: 0, useSSL: useSSL, allowPortReuse: allowPortReuse)
+    func startEphemeralServer(_ delegate: ServerDelegate?, useSSL: Bool = useSSLDefault, allowPortReuse: Bool = portReuseDefault, serverConfig: HTTPServerConfiguration = .default) throws -> (server: HTTPServer, port: Int) {
+        let serverConfig = serverConfig
+        let server = try startServer(delegate, port: 0, useSSL: useSSL, allowPortReuse: allowPortReuse, serverConfig: serverConfig)
         guard let serverPort = server.port else {
             throw KituraNetTestError(message: "Server port was not initialized")
         }
@@ -121,26 +122,26 @@ class KituraNetTest: XCTestCase {
         case both
     }
 
-    func performServerTest(_ delegate: ServerDelegate?, socketType: SocketType = .both, useSSL: Bool = useSSLDefault, allowPortReuse: Bool = portReuseDefault, line: Int = #line, asyncTasks: (XCTestExpectation) -> Void...) {
+    func performServerTest(serverConfig: HTTPServerConfiguration = .default, _ delegate: ServerDelegate?, socketType: SocketType = .both, useSSL: Bool = useSSLDefault, allowPortReuse: Bool = portReuseDefault, line: Int = #line, asyncTasks: (XCTestExpectation) -> Void...) {
+        let serverConfig = serverConfig
         self.socketType = socketType
         if socketType != .tcp {
-            performServerTestWithUnixSocket(delegate: delegate, useSSL: useSSL, allowPortReuse: allowPortReuse, line: line, asyncTasks: asyncTasks)
+            performServerTestWithUnixSocket(serverConfig: serverConfig, delegate: delegate, useSSL: useSSL, allowPortReuse: allowPortReuse, line: line, asyncTasks: asyncTasks)
         }
         if socketType != .unixDomainSocket {
-            performServerTestWithTCPPort(delegate: delegate, useSSL: useSSL, allowPortReuse:  allowPortReuse, line: line, asyncTasks: asyncTasks)
+            performServerTestWithTCPPort(serverConfig: serverConfig, delegate: delegate, useSSL: useSSL, allowPortReuse:  allowPortReuse, line: line, asyncTasks: asyncTasks)
         }
     }
 
-    func performServerTestWithUnixSocket(delegate: ServerDelegate?, useSSL: Bool = useSSLDefault, allowPortReuse: Bool = portReuseDefault, line: Int = #line, asyncTasks: [(XCTestExpectation) -> Void]) { 
+    func performServerTestWithUnixSocket(serverConfig: HTTPServerConfiguration = .default, delegate: ServerDelegate?, useSSL: Bool = useSSLDefault, allowPortReuse: Bool = portReuseDefault, line: Int = #line, asyncTasks: [(XCTestExpectation) -> Void]) {
         do {
             var server: HTTPServer
             self.useSSL = useSSL
             self.unixDomainSocketPath = self.socketFilePath
-            server = try startServer(delegate, unixDomainSocketPath: self.unixDomainSocketPath, useSSL: useSSL, allowPortReuse: allowPortReuse)
+            server = try startServer(delegate, unixDomainSocketPath: self.unixDomainSocketPath, useSSL: useSSL, allowPortReuse: allowPortReuse, serverConfig: serverConfig)
             defer {
                 server.stop()
             }
-
             let requestQueue = DispatchQueue(label: "Request queue")
             for (index, asyncTask) in asyncTasks.enumerated() {
                 let expectation = self.expectation(line: line, index: index)
@@ -158,18 +159,17 @@ class KituraNetTest: XCTestCase {
         }
     }
 
-    func performServerTestWithTCPPort(delegate: ServerDelegate?, useSSL: Bool = useSSLDefault, allowPortReuse: Bool = portReuseDefault, line: Int = #line, asyncTasks: [(XCTestExpectation) -> Void]) { 
+    func performServerTestWithTCPPort(serverConfig: HTTPServerConfiguration = .default, delegate: ServerDelegate?, useSSL: Bool = useSSLDefault, allowPortReuse: Bool = portReuseDefault, line: Int = #line, asyncTasks: [(XCTestExpectation) -> Void]) {
         do {
             var server: HTTPServer
             var ephemeralPort: Int = 0
             self.useSSL = useSSL
-            (server, ephemeralPort) = try startEphemeralServer(delegate, useSSL: useSSL, allowPortReuse: allowPortReuse)
+            (server, ephemeralPort) = try startEphemeralServer(delegate, useSSL: useSSL, allowPortReuse: allowPortReuse, serverConfig: serverConfig)
             self.port = ephemeralPort
             self.unixDomainSocketPath = nil
             defer {
                 server.stop()
             }
-
             let requestQueue = DispatchQueue(label: "Request queue")
             for (index, asyncTask) in asyncTasks.enumerated() {
                 let expectation = self.expectation(line: line, index: index)

diff --git a/Tests/KituraNetTests/LargePayloadTests.swift b/Tests/KituraNetTests/LargePayloadTests.swift
@@ -40,7 +40,7 @@ class LargePayloadTests: KituraNetTest {
     private let delegate = TestServerDelegate()
 
     func testLargePosts() {
-        performServerTest(delegate, useSSL: false, asyncTasks: { expectation in
+        performServerTest(serverConfig: HTTPServerConfiguration(requestSizeLimit: 100*10000, connectionLimit: 1),delegate,socketType: .tcp, useSSL: false, asyncTasks: { expectation in
             let payload = "[" + contentTypesString + "," + contentTypesString + contentTypesString + "," + contentTypesString + "]"
             self.performRequest("post", path: "/largepost", callback: {response in
                 XCTAssertEqual(response?.statusCode, HTTPStatusCode.OK, "Status code wasn't .Ok was \(String(describing: response?.statusCode))")

diff --git a/Tests/KituraNetTests/PipeliningTests.swift b/Tests/KituraNetTests/PipeliningTests.swift
@@ -80,7 +80,7 @@ class PipeliningTests: KituraNetTest {
     /// spanning several packets. It is necessary to sleep between writes to allow
     /// the server time to receive and process the data.
     func testPipeliningSpanningPackets() {
-        let server = HTTPServer()
+        let server = HTTPServer(serverConfig: HTTPServerConfiguration(requestSizeLimit: 100*1024, connectionLimit: 1))
         server.delegate = Delegate()
         server.delegate = Delegate()
         do {