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

Sources/KituraNet/HTTP/HTTPServer.swift

@@ -106,6 +106,13 @@ public class HTTPServer: Server {
      */
     public var keepAliveState: KeepAliveState = .unlimited
 
+    /// Controls policies relating to incoming connections and requests.
+    public var options: ServerOptions = ServerOptions() {
+        didSet {
+            self.socketManager?.serverOptions = options
+        }
+    }
+
     /// Incoming socket handler
     private var socketManager: IncomingSocketManager?
 
@@ -220,7 +227,7 @@ public class HTTPServer: Server {
                 listenerDescription = "path \(path)"
             }
 
-            let socketManager = IncomingSocketManager()
+            let socketManager = IncomingSocketManager(options: self.options)
             self.socketManager = socketManager
 
             if let delegate = socket.delegate {
@@ -349,7 +356,24 @@ public class HTTPServer: Server {
             do {
                 let clientSocket = try listenSocket.acceptClientConnection(invokeDelegate: false)
                 let clientSource = "\(clientSocket.remoteHostname):\(clientSocket.remotePort)"
-                Log.debug("Accepted HTTP connection from: \(clientSource)")
+                if let connectionLimit = self.options.connectionLimit, socketManager.socketHandlerCount >= connectionLimit {
+                    // See if any idle sockets can be removed before rejecting this connection
+                    socketManager.removeIdleSockets(runNow: true)
+                }
+                if let connectionLimit = self.options.connectionLimit, socketManager.socketHandlerCount >= connectionLimit {
+                    // Connections still at limit, this connection must be rejected
+                    if let (httpStatus, response) = self.options.connectionResponseGenerator(connectionLimit, clientSource) {
+                        let statusCode = httpStatus.rawValue
+                        let statusDescription = HTTP.statusCodes[statusCode] ?? ""
+                        let contentLength = response.utf8.count
+                        let httpResponse = "HTTP/1.1 \(statusCode) \(statusDescription)\r\nConnection: Close\r\nContent-Length: \(contentLength)\r\n\r\n".appending(response)
+                        _ = try? clientSocket.write(from: httpResponse)
+                    }
+                    clientSocket.close()
+                    continue
+                } else {
+                    Log.debug("Accepted HTTP connection from: \(clientSource)")
+                }
 
                 if listenSocket.delegate != nil {
                     DispatchQueue.global().async { [weak self] in

Sources/KituraNet/HTTP/IncomingHTTPSocketProcessor.swift

@@ -111,6 +111,12 @@ public class IncomingHTTPSocketProcessor: IncomingSocketProcessor {
 
     /// Location in the buffer to start parsing from
     private var parseStartingFrom = 0
+
+    /// Running total of body data bytes that we have parsed. Used to determine
+    /// whether this request is within the defined request size limit.
+    /// This number is only meaningful once the HTTP parser has completed
+    /// parsing of headers.
+    private var bodyBytesParsed = 0
 
     init(socket: Socket, using: ServerDelegate, keepalive: KeepAliveState) {
         delegate = using
@@ -236,7 +242,29 @@ public class IncomingHTTPSocketProcessor: IncomingSocketProcessor {
         }
 
         let bytes = buffer.bytes.assumingMemoryBound(to: Int8.self) + from
+        let headersAlreadyComplete = httpParser.headersComplete
         let (numberParsed, upgrade) = httpParser.execute(bytes, length: length)
+
+        // Count the number of body bytes parsed from the current buffer
+        if headersAlreadyComplete {
+            // Headers were completed while parsing a previous buffer. This entire buffer
+            // represents body data.
+            bodyBytesParsed += numberParsed
+        } else if httpParser.headersComplete {
+            // Headers complete while parsing current buffer. Subtract header length so that
+            // bodyBytesParsed is an accurate measure the number of bytes of body data.
+            bodyBytesParsed += (numberParsed - httpParser.headersLength)
+        } else {
+            // Entire buffer represents headers. We will subtract this later once headers
+            // are complete.
+            bodyBytesParsed += numberParsed
+        }
+
+        if httpParser.headersComplete, let handler = self.handler, let limit = handler.options.requestSizeLimit, bodyBytesParsed > limit {
+            handler.handleOversizedRead(limit)
+            status.error = .parsedLessThanRead
+            return status
+        }
 
         if completeBuffer && numberParsed == length {
             // Tell parser we reached the end
@@ -293,6 +321,7 @@ public class IncomingHTTPSocketProcessor: IncomingSocketProcessor {
         case .initial:
             break
         case .messageComplete:
+            bodyBytesParsed = 0
             isUpgrade = parsingStatus.upgrade
             clientRequestedKeepAlive = parsingStatus.keepAlive && !isUpgrade
             parsingComplete()

Sources/KituraNet/HTTPParser/HTTPParser.swift

@@ -41,9 +41,15 @@ class HTTPParser {
 
     /// Chunk of body read in by the http_parser
     var bodyChunk: BufferList { return parseResults.bodyChunk }
-    
+
     /// Parsing of message completed
     var completed: Bool { return parseResults.completed }
+
+    /// Parsing of headers completed
+    var headersComplete: Bool { return parseResults.headersComplete }
+
+    /// Length of headers section (in bytes)
+    var headersLength: Int { return parseResults.headersLength }
 
     /// A Handle to the HTTPParser C-library
     var parser: http_parser
@@ -169,7 +175,7 @@ fileprivate func onHeadersComplete(_ parser: UnsafeMutableRawPointer?) {
     let results = getResults(httpParser)
 
     results?.onHeadersComplete(method: method, versionMajor: (httpParser?.pointee.http_major)!,
-        versionMinor: (httpParser?.pointee.http_minor)!)
+                               versionMinor: (httpParser?.pointee.http_minor)!)
 }
 
 fileprivate func getResults(_ parser: UnsafeMutableRawPointer?) -> ParseResults? {

Sources/KituraNet/HTTPParser/ParseResults.swift

@@ -20,6 +20,12 @@ import Foundation
 class ParseResults {
     /// Have the parsing completed? (MessageComplete)
     var completed = false
+
+    /// Has the header parsing completed?
+    var headersComplete = false
+
+    /// The approximate size in bytes of the message headers.
+    var headersLength: Int = 0
 
     /// HTTP Method of the incoming message.
     private(set) var method = ""
@@ -69,6 +75,7 @@ class ParseResults {
     func onHeadersComplete(method: String, versionMajor: UInt16, versionMinor: UInt16) {
         httpVersionMajor = versionMajor
         httpVersionMinor = versionMinor
+        headersLength += method.count + 14  // length of method, HTTP version, spaces, newlines
         self.method = method
         if  lastHeaderWasAValue  {
             addHeader()
@@ -78,6 +85,7 @@ class ParseResults {
         url.append(&zero, length: 1)
         urlString = String(cString: url.bytes.assumingMemoryBound(to: CChar.self))
         url.length -= 1
+        headersComplete = true
     }
 
     /// Callback for when a piece of a header key was parsed
@@ -90,7 +98,7 @@ class ParseResults {
             addHeader()
         }
         lastHeaderField.append(bytes, length: count)
-
+        headersLength += count + 2  // Extra bytes for colon and space
         lastHeaderWasAValue = false
 
     }
@@ -101,7 +109,7 @@ class ParseResults {
     /// - Parameter count: The number of bytes parsed
     func onHeaderValue (_ bytes: UnsafePointer<UInt8>, count: Int) {
         lastHeaderValue.append(bytes, length: count)
-
+        headersLength += count + 2  // Extra bytes for newline
         lastHeaderWasAValue = true
     }
 
@@ -116,6 +124,7 @@ class ParseResults {
     /// - Parameter count: The number of bytes parsed
     func onURL(_ bytes: UnsafePointer<UInt8>, count: Int) {
         url.append(bytes, length: count)
+        headersLength += count
     }
 
     func reset() {

Sources/KituraNet/IncomingSocketHandler.swift

@@ -83,6 +83,11 @@ public class IncomingSocketHandler {
     private let writeBuffer = NSMutableData()
     private var writeBufferPosition = 0
 
+    /// Provides an ability to limit the maximum amount of data that can be read from a socket before rejecting a request and closing
+    /// the connection.
+    /// This is to protect against accidental or malicious requests from exhausting available memory.
+    let options: ServerOptions
+
     /// preparingToClose is set when prepareToClose() gets called or anytime we detect the socket has errored or was closed,
     /// so we try to close and cleanup as long as there is no data waiting to be written and a socket read/write is not in progress.
     private var preparingToClose = false
@@ -108,7 +113,7 @@ public class IncomingSocketHandler {
     /// The file descriptor of the incoming socket
     var fileDescriptor: Int32 { return socket.socketfd }
 
-    init(socket: Socket, using: IncomingSocketProcessor) {
+    init(socket: Socket, using: IncomingSocketProcessor, options: ServerOptions) {
         self.socket = socket
         processor = using
 
@@ -118,7 +123,8 @@ public class IncomingSocketHandler {
             readerSource = DispatchSource.makeReadSource(fileDescriptor: socket.socketfd,
                                                          queue: socketReaderQueue)
         #endif
-
+
+        self.options = options
         processor?.handler = self
 
         #if os(OSX) || os(iOS) || os(tvOS) || os(watchOS) || GCD_ASYNCH
@@ -228,6 +234,19 @@ public class IncomingSocketHandler {
             }
         #endif
     }
+
+    /// Handle the situation where we have received data that exceeds the configured requestSizeLimit.
+    func handleOversizedRead(_ limit: Int) {
+        let clientSource = "\(socket.remoteHostname):\(socket.remotePort)"
+        if let (httpStatus, response) = self.options.requestSizeResponseGenerator(limit, clientSource) {
+            let statusCode = httpStatus.rawValue
+            let statusDescription = HTTP.statusCodes[statusCode] ?? ""
+            let contentLength = response.utf8.count
+            let httpResponse = "HTTP/1.1 \(statusCode) \(statusDescription)\r\nConnection: Close\r\nContent-Length: \(contentLength)\r\n\r\n".appending(response)
+            _ = try? socket.write(from: httpResponse)
+        }
+        preparingToClose = true
+    }
 
     /// Write out any buffered data now that the socket can accept more data
     func handleWrite() {

Sources/KituraNet/IncomingSocketManager.swift

@@ -60,13 +60,16 @@ public class IncomingSocketManager  {
         }
         return result
     }
-    
+
     /// Interval at which to check for idle sockets to close
     let keepAliveIdleCheckingInterval: TimeInterval = 5.0
 
     /// The last time we checked for an idle socket
     var keepAliveIdleLastTimeChecked = Date()
 
+    /// Defines the limit on the size of requests. This may be modified by the user at runtime.
+    var serverOptions: ServerOptions = ServerOptions()
+
     /// Flag indicating when we are done using this socket manager, so we can clean up
     private var stopped = false
 
@@ -88,7 +91,8 @@ public class IncomingSocketManager  {
     /**
      IncomingSocketManager initializer
      */
-        public init() {
+        public init(options: ServerOptions = ServerOptions()) {
+            self.serverOptions = options
             var t1 = [Int32]()
             var t2 = [DispatchQueue]()
             for i in 0 ..< numberOfEpollTasks {
@@ -114,8 +118,8 @@ public class IncomingSocketManager  {
     /**
      IncomingSocketManager initializer
      */
-        public init() {
-
+        public init(options: ServerOptions = ServerOptions()) {
+            self.serverOptions = options
         }
     #endif
 
@@ -159,7 +163,7 @@ public class IncomingSocketManager  {
         do {
             try socket.setBlocking(mode: false)
 
-            let handler = IncomingSocketHandler(socket: socket, using: processor)
+            let handler = IncomingSocketHandler(socket: socket, using: processor, options: self.serverOptions)
             shQueue.sync(flags: .barrier) {
                 socketHandlers[socket.socketfd] = handler
             }
@@ -270,10 +274,11 @@ public class IncomingSocketManager  {
     ///   2. Removing the reference to the IncomingHTTPSocketHandler
     ///   3. Have the IncomingHTTPSocketHandler close the socket
     ///
-    /// - Parameter allSockets: flag indicating if the manager is shutting down, and we should cleanup all sockets, not just idle ones
-    private func removeIdleSockets(removeAll: Bool = false) {
+    /// - Parameter removeAll: flag indicating if the manager is shutting down, and we should cleanup all sockets, not just idle ones
+    /// - Parameter runNow: indicates that the removal should be performed immediately, rather than at the next scheduled interval
+    internal func removeIdleSockets(removeAll: Bool = false, runNow: Bool = false) {
         let now = Date()
-        guard removeAll || now.timeIntervalSince(keepAliveIdleLastTimeChecked) > keepAliveIdleCheckingInterval  else { return }
+        guard removeAll || runNow || now.timeIntervalSince(keepAliveIdleLastTimeChecked) > keepAliveIdleCheckingInterval  else { return }
         shQueue.sync(flags: .barrier) {
             let maxInterval = now.timeIntervalSinceReferenceDate
             for (fileDescriptor, handler) in socketHandlers {