├── .codecov.yml ├── .gitignore ├── Sources └── Gatekeeper │ ├── Request+Hostname.swift │ ├── KeyMaker │ ├── GatekeeperKeyMaker.swift │ └── GatekeeperHostnameKeyMaker.swift │ ├── Gatekeeper+Vapor │ ├── Request+Gatekeeper.swift │ └── Application+Gatekeeper.swift │ ├── GatekeeperEntry.swift │ ├── GatekeeperConfig.swift │ ├── GatekeeperMiddleware.swift │ └── Gatekeeper.swift ├── .swiftlint.yml ├── .github └── workflows │ ├── documentation.yml │ └── test.yml ├── Package.swift ├── LICENSE ├── README.md └── Tests └── GatekeeperTests └── GatekeeperTests.swift /.codecov.yml: -------------------------------------------------------------------------------- 1 | coverage: 2 | range: "0...100" 3 | ignore: 4 | - "Tests" 5 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | Packages 2 | .build 3 | .idea 4 | xcuserdata 5 | *.xcodeproj 6 | Config/secrets/ 7 | .DS_Store 8 | .swift-version 9 | CMakeLists.txt 10 | Package.resolved 11 | .swiftpm -------------------------------------------------------------------------------- /Sources/Gatekeeper/Request+Hostname.swift: -------------------------------------------------------------------------------- 1 | import Vapor 2 | 3 | extension Request { 4 | var hostname: String? { 5 | headers.forwarded.first?.for ?? headers.first(name: .xForwardedFor) ?? remoteAddress?.hostname 6 | } 7 | } 8 | -------------------------------------------------------------------------------- /Sources/Gatekeeper/KeyMaker/GatekeeperKeyMaker.swift: -------------------------------------------------------------------------------- 1 | import Vapor 2 | 3 | /// Reponsible for generating a cache key for a specific `Request` 4 | public protocol GatekeeperKeyMaker { 5 | func make(for req: Request) -> EventLoopFuture 6 | } 7 | -------------------------------------------------------------------------------- /.swiftlint.yml: -------------------------------------------------------------------------------- 1 | included: 2 | - Sources 3 | function_body_length: 4 | warning: 60 5 | identifier_name: 6 | min_length: 7 | warning: 2 8 | line_length: 100 9 | disabled_rules: 10 | - opening_brace 11 | - nesting 12 | colon: 13 | flexible_right_spacing: true -------------------------------------------------------------------------------- /Sources/Gatekeeper/KeyMaker/GatekeeperHostnameKeyMaker.swift: -------------------------------------------------------------------------------- 1 | import Vapor 2 | 3 | /// Uses the hostname of the client to create a cache key. 4 | public struct GatekeeperHostnameKeyMaker: GatekeeperKeyMaker { 5 | public func make(for req: Request) -> EventLoopFuture { 6 | guard let hostname = req.hostname else { 7 | return req.eventLoop.future(error: Abort(.forbidden, reason: "Unable to verify peer")) 8 | } 9 | 10 | return req.eventLoop.future("gatekeeper_" + hostname) 11 | } 12 | } 13 | -------------------------------------------------------------------------------- /Sources/Gatekeeper/Gatekeeper+Vapor/Request+Gatekeeper.swift: -------------------------------------------------------------------------------- 1 | import Vapor 2 | 3 | public extension Request { 4 | func gatekeeper( 5 | config: GatekeeperConfig? = nil, 6 | cache: Cache? = nil, 7 | keyMaker: GatekeeperKeyMaker? = nil 8 | ) -> Gatekeeper { 9 | .init( 10 | cache: cache ?? application.gatekeeper.caches.cache.for(self), 11 | config: config ?? application.gatekeeper.config, 12 | identifier: keyMaker ?? application.gatekeeper.keyMakers.keyMaker 13 | ) 14 | } 15 | } 16 | -------------------------------------------------------------------------------- /.github/workflows/documentation.yml: -------------------------------------------------------------------------------- 1 | name: Documentation 2 | 3 | on: 4 | release: 5 | types: [published] 6 | 7 | jobs: 8 | build: 9 | runs-on: ubuntu-latest 10 | 11 | steps: 12 | - uses: actions/checkout@v1 13 | - name: Generate Documentation 14 | uses: SwiftDocOrg/swift-doc@master 15 | with: 16 | inputs: "Sources" 17 | module-name: Gatekeeper 18 | output: "Documentation" 19 | - name: Upload Documentation to Wiki 20 | uses: SwiftDocOrg/github-wiki-publish-action@v1 21 | with: 22 | path: "Documentation" 23 | env: 24 | GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.WIKI_ACCESS_TOKEN }} -------------------------------------------------------------------------------- /Sources/Gatekeeper/GatekeeperEntry.swift: -------------------------------------------------------------------------------- 1 | import Vapor 2 | 3 | extension Gatekeeper { 4 | /// A model representing a entry in the cache for a specific client 5 | public struct Entry: Codable { 6 | let hostname: String 7 | var createdAt: Date 8 | var requestsLeft: Int 9 | } 10 | } 11 | 12 | extension Gatekeeper.Entry { 13 | func hasExpired(within interval: Double) -> Bool { 14 | Date().timeIntervalSince1970 - createdAt.timeIntervalSince1970 >= interval 15 | } 16 | 17 | mutating func touch() { 18 | if requestsLeft > 0 { 19 | requestsLeft -= 1 20 | } else { 21 | requestsLeft = 0 22 | } 23 | } 24 | } 25 | -------------------------------------------------------------------------------- /Sources/Gatekeeper/GatekeeperConfig.swift: -------------------------------------------------------------------------------- 1 | import Vapor 2 | 3 | public struct GatekeeperConfig { 4 | public enum Interval { 5 | case second 6 | case minute 7 | case hour 8 | case day 9 | } 10 | 11 | public let limit: Int 12 | public let interval: Interval 13 | 14 | public init(maxRequests limit: Int, per interval: Interval) { 15 | self.limit = limit 16 | self.interval = interval 17 | } 18 | 19 | var refreshInterval: Double { 20 | switch interval { 21 | case .second: 22 | return 1 23 | case .minute: 24 | return 60 25 | case .hour: 26 | return 3_600 27 | case .day: 28 | return 86_400 29 | } 30 | } 31 | } 32 | -------------------------------------------------------------------------------- /.github/workflows/test.yml: -------------------------------------------------------------------------------- 1 | name: test 2 | on: 3 | pull_request: 4 | push: 5 | branches: 6 | - master 7 | jobs: 8 | linux: 9 | runs-on: ubuntu-latest 10 | container: swift:5.3-focal 11 | steps: 12 | - name: Check out code 13 | uses: actions/checkout@v2 14 | - name: Run tests with Thread Sanitizer 15 | run: swift test --enable-test-discovery --sanitize=thread 16 | macOS: 17 | runs-on: macos-latest 18 | steps: 19 | - name: Select latest available Xcode 20 | uses: maxim-lobanov/setup-xcode@v1 21 | with: 22 | xcode-version: latest 23 | - name: Check out code 24 | uses: actions/checkout@v2 25 | - name: Run tests with Thread Sanitizer 26 | run: swift test --enable-test-discovery --sanitize=thread -------------------------------------------------------------------------------- /Package.swift: -------------------------------------------------------------------------------- 1 | // swift-tools-version:5.3 2 | import PackageDescription 3 | 4 | let package = Package( 5 | name: "gatekeeper", 6 | platforms: [ 7 | .macOS(.v10_15), 8 | ], 9 | products: [ 10 | .library( 11 | name: "Gatekeeper", 12 | targets: ["Gatekeeper"]), 13 | ], 14 | dependencies: [ 15 | .package(url: "https://github.com/vapor/vapor.git", from: "4.44.0"), 16 | ], 17 | targets: [ 18 | .target( 19 | name: "Gatekeeper", 20 | dependencies: [ 21 | .product(name: "Vapor", package: "vapor") 22 | ]), 23 | .testTarget( 24 | name: "GatekeeperTests", 25 | dependencies: [ 26 | "Gatekeeper", 27 | .product(name: "XCTVapor", package: "vapor") 28 | ]), 29 | ] 30 | ) 31 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2017-2021 Nodes 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /Sources/Gatekeeper/GatekeeperMiddleware.swift: -------------------------------------------------------------------------------- 1 | import Vapor 2 | 3 | /// Middleware used to rate-limit a single route or a group of routes. 4 | public struct GatekeeperMiddleware: Middleware { 5 | private let config: GatekeeperConfig? 6 | private let keyMaker: GatekeeperKeyMaker? 7 | private let error: Error? 8 | 9 | /// Initialize a new middleware for rate-limiting routes, by optionally overriding default configurations. 10 | /// 11 | /// - Parameters: 12 | /// - config: Override `GatekeeperConfig` instead of using the default `app.gatekeeper.config` 13 | /// - keyMaker: Override `GatekeeperKeyMaker` instead of using the default `app.gatekeeper.keyMaker` 14 | /// - config: Override the `Error` thrown when the user is rate-limited instead of using the default error. 15 | public init(config: GatekeeperConfig? = nil, keyMaker: GatekeeperKeyMaker? = nil, error: Error? = nil) { 16 | self.config = config 17 | self.keyMaker = keyMaker 18 | self.error = error 19 | } 20 | 21 | public func respond(to request: Request, chainingTo next: Responder) -> EventLoopFuture { 22 | let gatekeeper = request.gatekeeper(config: config, keyMaker: keyMaker) 23 | 24 | let gatekeep: EventLoopFuture 25 | if let error = error { 26 | gatekeep = gatekeeper.gatekeep(on: request, throwing: error) 27 | } else { 28 | gatekeep = gatekeeper.gatekeep(on: request) 29 | } 30 | 31 | return gatekeep.flatMap { next.respond(to: request) } 32 | } 33 | } 34 | -------------------------------------------------------------------------------- /Sources/Gatekeeper/Gatekeeper.swift: -------------------------------------------------------------------------------- 1 | import Vapor 2 | 3 | public struct Gatekeeper { 4 | private let cache: Cache 5 | private let config: GatekeeperConfig 6 | private let keyMaker: GatekeeperKeyMaker 7 | 8 | public init(cache: Cache, config: GatekeeperConfig, identifier: GatekeeperKeyMaker) { 9 | self.cache = cache 10 | self.config = config 11 | self.keyMaker = identifier 12 | } 13 | 14 | public func gatekeep( 15 | on req: Request, 16 | throwing error: Error = Abort(.tooManyRequests, reason: "Slow down. You sent too many requests.") 17 | ) -> EventLoopFuture { 18 | keyMaker 19 | .make(for: req) 20 | .flatMap { cacheKey in 21 | fetchOrCreateEntry(for: cacheKey, on: req) 22 | .guard( 23 | { $0.requestsLeft > 0 }, 24 | else: error 25 | ) 26 | .map(updateEntry) 27 | .flatMap { entry in 28 | // The amount of time the entry has existed. 29 | let entryLifetime = Int(Date().timeIntervalSince1970 - entry.createdAt.timeIntervalSince1970) 30 | // Remaining time until the entry expires. The entry would be expired by cache if it was negative. 31 | let timeRemaining = Int(config.refreshInterval) - entryLifetime 32 | return cache.set(cacheKey, to: entry, expiresIn: .seconds(timeRemaining)) 33 | } 34 | } 35 | } 36 | 37 | private func updateEntry(_ entry: Entry) -> Entry { 38 | var newEntry = entry 39 | newEntry.touch() 40 | return newEntry 41 | } 42 | 43 | private func fetchOrCreateEntry(for key: String, on req: Request) -> EventLoopFuture { 44 | guard let hostname = req.hostname else { 45 | return req.eventLoop.future(error: Abort(.forbidden, reason: "Unable to verify peer")) 46 | } 47 | 48 | return cache 49 | .get(key, as: Entry.self) 50 | .unwrap(orReplace: Entry(hostname: hostname, createdAt: Date(), requestsLeft: config.limit)) 51 | } 52 | } 53 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Gatekeeper 👮 2 | [![Swift Version](https://img.shields.io/badge/Swift-5.3-brightgreen.svg)](http://swift.org) 3 | [![Vapor Version](https://img.shields.io/badge/Vapor-4-30B6FC.svg)](http://vapor.codes) 4 | [![GitHub license](https://img.shields.io/badge/license-MIT-blue.svg)](https://raw.githubusercontent.com/nodes-vapor/gatekeeper/master/LICENSE) 5 | 6 | Gatekeeper is a middleware that restricts the number of requests from clients, based on their IP address **(can be customized)**. 7 | It works by adding the clients identifier to the cache and count how many requests the clients can make during the Gatekeeper's defined lifespan and give back an HTTP 429(Too Many Requests) if the limit has been reached. The number of requests left will be reset when the defined timespan has been reached. 8 | 9 | **Please take into consideration that multiple clients can be using the same IP address. eg. public wifi** 10 | 11 | 12 | ## 📦 Installation 13 | 14 | Update your `Package.swift` dependencies: 15 | 16 | ```swift 17 | .package(url: "https://github.com/nodes-vapor/gatekeeper.git", from: "4.0.0"), 18 | ``` 19 | 20 | as well as to your target (e.g. "App"): 21 | 22 | ```swift 23 | targets: [ 24 | .target(name: "App", dependencies: [..., "Gatekeeper", ...]), 25 | // ... 26 | ] 27 | ``` 28 | 29 | ## Getting started 🚀 30 | 31 | ### Configuration 32 | 33 | in configure.swift: 34 | ```swift 35 | import Gatekeeper 36 | 37 | // [...] 38 | 39 | app.caches.use(.memory) 40 | app.gatekeeper.config = .init(maxRequests: 10, per: .second) 41 | ``` 42 | 43 | ### Add to routes 44 | 45 | You can add the `GatekeeperMiddleware` to specific routes or to all. 46 | 47 | **Specific routes** 48 | in routes.swift: 49 | ```swift 50 | let protectedRoutes = router.grouped(GatekeeperMiddleware()) 51 | protectedRoutes.get("protected/hello") { req in 52 | return "Protected Hello, World!" 53 | } 54 | ``` 55 | 56 | **For all requests** 57 | in configure.swift: 58 | ```swift 59 | // Register middleware 60 | app.middlewares.use(GatekeeperMiddleware()) 61 | ``` 62 | 63 | #### Customizing config 64 | By default `GatekeeperMiddleware` uses `app.gatekeeper.config` as its configuration. 65 | However, you can pass a custom configuration to each `GatekeeperMiddleware` type via the initializer 66 | `GatekeeperMiddleware(config:)`. This allows you to set configuration on a per-route basis. 67 | 68 | ## Key Makers 🔑 69 | By default Gatekeeper uses the client's hostname (IP address) to identify them. This can cause issues where multiple clients are connected from the same network. Therefore, you can customize how Gatekeeper should identify the client by using the `GatekeeperKeyMaker` protocol. 70 | 71 | `GatekeeperHostnameKeyMaker` is used by default. 72 | 73 | You can configure which key maker Gatekeeper should use in `configure.swift`: 74 | ```swift 75 | app.gatekeeper.keyMakers.use(.hostname) // default 76 | ``` 77 | 78 | ### Custom key maker 79 | This is an example of a key maker that uses the user's ID to identify them. 80 | ```swift 81 | struct UserIDKeyMaker: GatekeeperKeyMaker { 82 | public func make(for req: Request) -> EventLoopFuture { 83 | let userID = try req.auth.require(User.self).requireID() 84 | return req.eventLoop.future("gatekeeper_" + userID.uuidString) 85 | } 86 | } 87 | ``` 88 | 89 | ```swift 90 | extension Application.Gatekeeper.KeyMakers.Provider { 91 | public static var userID: Self { 92 | .init { app in 93 | app.gatekeeper.keyMakers.use { _ in UserIDKeyMaker() } 94 | } 95 | } 96 | } 97 | ``` 98 | **configure.swift:** 99 | ```swift 100 | app.gatekeeper.keyMakers.use(.userID) 101 | ``` 102 | 103 | ## Cache 🗄 104 | Gatekeeper uses the same cache as configured by `app.caches.use()` from Vapor, by default. 105 | Therefore it is **important** to set up Vapor's cache if you're using this default behaviour. You can use an in-memory cache for Vapor like so: 106 | 107 | **configure.swift**: 108 | ```swift 109 | app.cache.use(.memory) 110 | ``` 111 | 112 | ### Custom cache 113 | You can override which cache to use by creating your own type that conforms to the `Cache` protocol from Vapor. Use `app.gatekeeper.caches.use()` to configure which cache to use. 114 | 115 | 116 | ## Credits 🏆 117 | 118 | This package is developed and maintained by the Vapor team at [Nodes](https://www.nodesagency.com). 119 | The package owner for this project is [Christian](https://github.com/cweinberger). 120 | Special thanks goes to [madsodgaard](https://github.com/madsodgaard) for his work on the Vapor 4 version! 121 | 122 | ## License 📄 123 | 124 | This package is open-sourced software licensed under the [MIT license](http://opensource.org/licenses/MIT) 125 | -------------------------------------------------------------------------------- /Sources/Gatekeeper/Gatekeeper+Vapor/Application+Gatekeeper.swift: -------------------------------------------------------------------------------- 1 | import Vapor 2 | 3 | // MARK: - Application+Gatekeeper 4 | extension Application { 5 | public struct Gatekeeper { 6 | private let app: Application 7 | 8 | init(app: Application) { 9 | self.app = app 10 | } 11 | 12 | private final class Storage { 13 | var config: GatekeeperConfig? = nil 14 | var makeCache: ((Application) -> Cache)? = nil 15 | var makeKeyMaker: ((Application) -> GatekeeperKeyMaker)? = nil 16 | } 17 | 18 | private struct Key: StorageKey { 19 | typealias Value = Storage 20 | } 21 | 22 | private var storage: Storage { 23 | if app.storage[Key.self] == nil { 24 | initialize() 25 | } 26 | 27 | return app.storage[Key.self]! 28 | } 29 | 30 | private func initialize() { 31 | app.storage[Key.self] = Storage() 32 | app.gatekeeper.caches.use(.default) 33 | app.gatekeeper.keyMakers.use(.hostname) 34 | } 35 | 36 | /// The default config used for middlewares. 37 | public var config: GatekeeperConfig { 38 | get { 39 | guard let config = storage.config else { 40 | fatalError("Gatekeeper not configured, use: app.gatekeeper.config = ...") 41 | } 42 | 43 | return config 44 | } 45 | nonmutating set { storage.config = newValue } 46 | } 47 | } 48 | 49 | public var gatekeeper: Gatekeeper { 50 | .init(app: self) 51 | } 52 | } 53 | 54 | // MARK: - Gatekeeper+Caches 55 | extension Application.Gatekeeper { 56 | public struct Caches { 57 | private let gatekeeper: Application.Gatekeeper 58 | 59 | public init(_ gatekeeper: Application.Gatekeeper) { 60 | self.gatekeeper = gatekeeper 61 | } 62 | 63 | public struct Provider { 64 | public let run: (Application) -> Void 65 | 66 | public init(_ run: @escaping (Application) -> Void) { 67 | self.run = run 68 | } 69 | 70 | /// A provider that uses the default Vapor cache. 71 | public static var `default`: Self { 72 | .init { app in 73 | app.gatekeeper.caches.use { $0.cache } 74 | } 75 | } 76 | } 77 | 78 | public func use(_ makeCache: @escaping (Application) -> Cache) { 79 | gatekeeper.storage.makeCache = makeCache 80 | } 81 | 82 | public func use(_ provider: Provider) { 83 | provider.run(gatekeeper.app) 84 | } 85 | 86 | public var cache: Cache { 87 | guard let factory = gatekeeper.storage.makeCache else { 88 | fatalError("Gatekeeper not configured, use: app.gatekeeper.caches.use(...)") 89 | } 90 | 91 | return factory(gatekeeper.app) 92 | } 93 | } 94 | 95 | public var caches: Caches { 96 | .init(self) 97 | } 98 | } 99 | 100 | // MARK: - Gatekeeper+Keymakers 101 | extension Application.Gatekeeper { 102 | public struct KeyMakers { 103 | private let gatekeeper: Application.Gatekeeper 104 | 105 | public init(_ gatekeeper: Application.Gatekeeper) { 106 | self.gatekeeper = gatekeeper 107 | } 108 | 109 | public struct Provider { 110 | public let run: (Application) -> Void 111 | 112 | public init(_ run: @escaping (Application) -> Void) { 113 | self.run = run 114 | } 115 | 116 | /// A provider that the request hostname to generate a cache key. 117 | public static var hostname: Self { 118 | .init { app in 119 | app.gatekeeper.keyMakers.use { _ in GatekeeperHostnameKeyMaker() } 120 | } 121 | } 122 | } 123 | 124 | public func use(_ makeKeyMaker: @escaping (Application) -> GatekeeperKeyMaker) { 125 | gatekeeper.storage.makeKeyMaker = makeKeyMaker 126 | } 127 | 128 | public func use(_ provider: Provider) { 129 | provider.run(gatekeeper.app) 130 | } 131 | 132 | public var keyMaker: GatekeeperKeyMaker { 133 | guard let factory = gatekeeper.storage.makeKeyMaker else { 134 | fatalError("Gatekeeper not configured, use: app.gatekeeper.keyMakers.use(...)") 135 | } 136 | 137 | return factory(gatekeeper.app) 138 | } 139 | } 140 | 141 | public var keyMakers: KeyMakers { 142 | .init(self) 143 | } 144 | } 145 | -------------------------------------------------------------------------------- /Tests/GatekeeperTests/GatekeeperTests.swift: -------------------------------------------------------------------------------- 1 | import XCTest 2 | import XCTVapor 3 | @testable import Gatekeeper 4 | 5 | class GatekeeperTests: XCTestCase { 6 | func testGateKeeper() throws { 7 | let app = Application(.testing) 8 | defer { app.shutdown() } 9 | app.gatekeeper.config = .init(maxRequests: 10, per: .second) 10 | 11 | app.grouped(GatekeeperMiddleware()).get("test") { req -> HTTPStatus in 12 | return .ok 13 | } 14 | 15 | for i in 1...11 { 16 | try app.test(.GET, "test", headers: ["X-Forwarded-For": "::1"], afterResponse: { res in 17 | if i == 11 { 18 | XCTAssertEqual(res.status, .tooManyRequests) 19 | } else { 20 | XCTAssertEqual(res.status, .ok, "failed for request \(i) with status: \(res.status)") 21 | } 22 | }) 23 | } 24 | } 25 | 26 | func testGateKeeperNoPeerReturnsForbidden() throws { 27 | let app = Application(.testing) 28 | defer { app.shutdown() } 29 | app.gatekeeper.config = .init(maxRequests: 10, per: .second) 30 | 31 | app.grouped(GatekeeperMiddleware()).get("test") { req -> HTTPStatus in 32 | return .ok 33 | } 34 | 35 | try app.test(.GET, "test", afterResponse: { res in 36 | XCTAssertEqual(res.status, .forbidden) 37 | }) 38 | } 39 | 40 | func testGateKeeperForwardedSupported() throws { 41 | let app = Application(.testing) 42 | defer { app.shutdown() } 43 | app.gatekeeper.config = .init(maxRequests: 10, per: .second) 44 | 45 | app.grouped(GatekeeperMiddleware()).get("test") { req -> HTTPStatus in 46 | return .ok 47 | } 48 | 49 | try app.test( 50 | .GET, 51 | "test", 52 | beforeRequest: { req in 53 | req.headers.forwarded = [HTTPHeaders.Forwarded(for: "\"[::1]\"")] 54 | }, 55 | afterResponse: { res in 56 | XCTAssertEqual(res.status, .ok) 57 | } 58 | ) 59 | } 60 | 61 | func testGateKeeperCountRefresh() throws { 62 | let app = Application(.testing) 63 | defer { app.shutdown() } 64 | app.gatekeeper.config = .init(maxRequests: 100, per: .second) 65 | app.grouped(GatekeeperMiddleware()).get("test") { req -> HTTPStatus in 66 | return .ok 67 | } 68 | 69 | for _ in 0..<50 { 70 | try app.test(.GET, "test", headers: ["X-Forwarded-For": "::1"], afterResponse: { res in 71 | XCTAssertEqual(res.status, .ok) 72 | }) 73 | } 74 | 75 | let entryBefore = try app.gatekeeper.caches.cache.get("gatekeeper_::1", as: Gatekeeper.Entry.self).wait() 76 | XCTAssertEqual(entryBefore!.requestsLeft, 50) 77 | 78 | Thread.sleep(forTimeInterval: 1) 79 | 80 | try app.test(.GET, "test", headers: ["X-Forwarded-For": "::1"], afterResponse: { res in 81 | XCTAssertEqual(res.status, .ok) 82 | }) 83 | 84 | let entryAfter = try app.gatekeeper.caches.cache .get("gatekeeper_::1", as: Gatekeeper.Entry.self).wait() 85 | XCTAssertEqual(entryAfter!.requestsLeft, 99, "Requests left should've reset") 86 | } 87 | 88 | func testGatekeeperCacheExpiry() throws { 89 | let app = Application(.testing) 90 | defer { app.shutdown() } 91 | app.gatekeeper.config = .init(maxRequests: 5, per: .second) 92 | app.grouped(GatekeeperMiddleware()).get("test") { req -> HTTPStatus in 93 | return .ok 94 | } 95 | 96 | for _ in 1...5 { 97 | try app.test(.GET, "test", headers: ["X-Forwarded-For": "::1"], afterResponse: { res in 98 | XCTAssertEqual(res.status, .ok) 99 | }) 100 | } 101 | 102 | let entryBefore = try app.gatekeeper.caches.cache.get("gatekeeper_::1", as: Gatekeeper.Entry.self).wait() 103 | XCTAssertEqual(entryBefore!.requestsLeft, 0) 104 | 105 | Thread.sleep(forTimeInterval: 1) 106 | 107 | try XCTAssertNil(app.gatekeeper.caches.cache.get("gatekeeper_::1", as: Gatekeeper.Entry.self).wait()) 108 | } 109 | 110 | func testRefreshIntervalValues() { 111 | let expected: [(GatekeeperConfig.Interval, Double)] = [ 112 | (.second, 1), 113 | (.minute, 60), 114 | (.hour, 3_600), 115 | (.day, 86_400) 116 | ] 117 | 118 | expected.forEach { interval, expected in 119 | let rate = GatekeeperConfig(maxRequests: 1, per: interval) 120 | XCTAssertEqual(rate.refreshInterval, expected) 121 | } 122 | } 123 | 124 | func testGatekeeperUsesKeyMaker() throws { 125 | struct DummyKeyMaker: GatekeeperKeyMaker { 126 | func make(for req: Request) -> EventLoopFuture { 127 | req.eventLoop.future("dummy") 128 | } 129 | } 130 | 131 | let app = Application(.testing) 132 | defer { app.shutdown() } 133 | app.gatekeeper.config = .init(maxRequests: 10, per: .second) 134 | app.gatekeeper.keyMakers.use { _ in 135 | DummyKeyMaker() 136 | } 137 | 138 | app.grouped(GatekeeperMiddleware()).get("test") { req -> HTTPStatus in 139 | return .ok 140 | } 141 | 142 | try app.test(.GET, "test", headers: ["X-Forwarded-For": "::1"], afterResponse: { _ in }) 143 | 144 | let entry = try app.gatekeeper.caches.cache.get("dummy", as: Gatekeeper.Entry.self).wait() 145 | XCTAssertNotNil(entry) 146 | } 147 | 148 | func testGatekeeperDefaultProviders() throws { 149 | let app = Application(.testing) 150 | defer { app.shutdown() } 151 | 152 | XCTAssertTrue(app.gatekeeper.keyMakers.keyMaker is GatekeeperHostnameKeyMaker) 153 | } 154 | } 155 | --------------------------------------------------------------------------------