From 4f2f77d63dfca2b204a5cef85fab4b0b1b90bc92 Mon Sep 17 00:00:00 2001 From: Mads Odgaard Date: Mon, 12 Apr 2021 17:38:58 +0200 Subject: [PATCH 1/3] Add expiry to cache entry --- Package.swift | 2 +- .../Gatekeeper+Vapor/Request+Gatekeeper.swift | 2 +- Sources/Gatekeeper/Gatekeeper.swift | 26 ++++++++++------- Sources/Gatekeeper/GatekeeperEntry.swift | 11 ++++---- Tests/GatekeeperTests/GatekeeperTests.swift | 28 +++++++++++++++++-- 5 files changed, 48 insertions(+), 21 deletions(-) diff --git a/Package.swift b/Package.swift index cab3258..8d5f1ec 100644 --- a/Package.swift +++ b/Package.swift @@ -12,7 +12,7 @@ let package = Package( targets: ["Gatekeeper"]), ], dependencies: [ - .package(url: "https://github.com/vapor/vapor.git", from: "4.38.0"), + .package(url: "https://github.com/vapor/vapor.git", from: "4.44.0"), ], targets: [ .target( diff --git a/Sources/Gatekeeper/Gatekeeper+Vapor/Request+Gatekeeper.swift b/Sources/Gatekeeper/Gatekeeper+Vapor/Request+Gatekeeper.swift index b432bf0..9afe91d 100644 --- a/Sources/Gatekeeper/Gatekeeper+Vapor/Request+Gatekeeper.swift +++ b/Sources/Gatekeeper/Gatekeeper+Vapor/Request+Gatekeeper.swift @@ -1,6 +1,6 @@ import Vapor -extension Request { +public extension Request { func gatekeeper( config: GatekeeperConfig? = nil, cache: Cache? = nil, diff --git a/Sources/Gatekeeper/Gatekeeper.swift b/Sources/Gatekeeper/Gatekeeper.swift index 5e67b17..2ad9ab4 100644 --- a/Sources/Gatekeeper/Gatekeeper.swift +++ b/Sources/Gatekeeper/Gatekeeper.swift @@ -11,29 +11,35 @@ public struct Gatekeeper { self.keyMaker = identifier } - public func gatekeep(on req: Request) -> EventLoopFuture { + public func gatekeep( + on req: Request, + throwing error: Error = Abort(.tooManyRequests, reason: "Slow down. You sent too many requests.") + ) -> EventLoopFuture { keyMaker .make(for: req) .flatMap { cacheKey in fetchOrCreateEntry(for: cacheKey, on: req) + .guard( + { $0.requestsLeft > 0 }, + else: error + ) .map(updateEntry) - .flatMap { entry in - cache - .set(cacheKey, to: entry) + .flatMap { entry -> EventLoopFuture in + // The amount of time the entry has existed. + let entryLifetime = Int(Date().timeIntervalSince1970 - entry.createdAt.timeIntervalSince1970) + // Remaining time until the entry expires. The entry would be expired by cache if it was negative. + let timeRemaining = Int(config.refreshInterval) - entryLifetime + + return cache + .set(cacheKey, to: entry, expiresIn: .seconds(timeRemaining)) .transform(to: entry) } } - .guard( - { $0.requestsLeft > 0 }, - else: Abort(.tooManyRequests, reason: "Slow down. You sent too many requests.")) .transform(to: ()) } private func updateEntry(_ entry: Entry) -> Entry { var newEntry = entry - if newEntry.hasExpired(within: config.refreshInterval) { - newEntry.reset(remainingRequests: config.limit) - } newEntry.touch() return newEntry } diff --git a/Sources/Gatekeeper/GatekeeperEntry.swift b/Sources/Gatekeeper/GatekeeperEntry.swift index a1bf908..3302a09 100644 --- a/Sources/Gatekeeper/GatekeeperEntry.swift +++ b/Sources/Gatekeeper/GatekeeperEntry.swift @@ -14,12 +14,11 @@ extension Gatekeeper.Entry { Date().timeIntervalSince1970 - createdAt.timeIntervalSince1970 >= interval } - mutating func reset(remainingRequests: Int) { - createdAt = Date() - requestsLeft = remainingRequests - } - mutating func touch() { - requestsLeft -= 1 + if requestsLeft > 0 { + requestsLeft -= 1 + } else { + requestsLeft = 0 + } } } diff --git a/Tests/GatekeeperTests/GatekeeperTests.swift b/Tests/GatekeeperTests/GatekeeperTests.swift index bca7120..300c898 100644 --- a/Tests/GatekeeperTests/GatekeeperTests.swift +++ b/Tests/GatekeeperTests/GatekeeperTests.swift @@ -12,9 +12,9 @@ class GatekeeperTests: XCTestCase { return .ok } - for i in 1...10 { + for i in 1...11 { try app.test(.GET, "test", headers: ["X-Forwarded-For": "::1"], afterResponse: { res in - if i == 10 { + if i == 11 { XCTAssertEqual(res.status, .tooManyRequests) } else { XCTAssertEqual(res.status, .ok, "failed for request \(i) with status: \(res.status)") @@ -51,7 +51,7 @@ class GatekeeperTests: XCTestCase { }) } - let entryBefore = try app.gatekeeper.caches.cache .get("gatekeeper_::1", as: Gatekeeper.Entry.self).wait() + let entryBefore = try app.gatekeeper.caches.cache.get("gatekeeper_::1", as: Gatekeeper.Entry.self).wait() XCTAssertEqual(entryBefore!.requestsLeft, 50) Thread.sleep(forTimeInterval: 1) @@ -63,6 +63,28 @@ class GatekeeperTests: XCTestCase { let entryAfter = try app.gatekeeper.caches.cache .get("gatekeeper_::1", as: Gatekeeper.Entry.self).wait() XCTAssertEqual(entryAfter!.requestsLeft, 99, "Requests left should've reset") } + + func testGatekeeperCacheExpiry() throws { + let app = Application(.testing) + defer { app.shutdown() } + app.gatekeeper.config = .init(maxRequests: 5, per: .second) + app.grouped(GatekeeperMiddleware()).get("test") { req -> HTTPStatus in + return .ok + } + + for _ in 1...5 { + try app.test(.GET, "test", headers: ["X-Forwarded-For": "::1"], afterResponse: { res in + XCTAssertEqual(res.status, .ok) + }) + } + + let entryBefore = try app.gatekeeper.caches.cache.get("gatekeeper_::1", as: Gatekeeper.Entry.self).wait() + XCTAssertEqual(entryBefore!.requestsLeft, 0) + + Thread.sleep(forTimeInterval: 1) + + try XCTAssertNil(app.gatekeeper.caches.cache.get("gatekeeper_::1", as: Gatekeeper.Entry.self).wait()) + } func testRefreshIntervalValues() { let expected: [(GatekeeperConfig.Interval, Double)] = [ From 3e78119b858aa29152df081b079959ae22736677 Mon Sep 17 00:00:00 2001 From: Mads Odgaard Date: Mon, 12 Apr 2021 17:48:34 +0200 Subject: [PATCH 2/3] Add customization to middleware --- Sources/Gatekeeper/GatekeeperMiddleware.swift | 27 ++++++++++++++----- 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/Sources/Gatekeeper/GatekeeperMiddleware.swift b/Sources/Gatekeeper/GatekeeperMiddleware.swift index e6f6f01..23026da 100644 --- a/Sources/Gatekeeper/GatekeeperMiddleware.swift +++ b/Sources/Gatekeeper/GatekeeperMiddleware.swift @@ -3,16 +3,31 @@ import Vapor /// Middleware used to rate-limit a single route or a group of routes. public struct GatekeeperMiddleware: Middleware { private let config: GatekeeperConfig? + private let keyMaker: GatekeeperKeyMaker? + private let error: Error? - /// Initialize with a custom `GatekeeperConfig` instead of using the default `app.gatekeeper.config` - public init(config: GatekeeperConfig? = nil) { + /// Initialize a new middleware for rate-limiting routes, by optionally overriding default configurations. + /// + /// - Parameters: + /// - config: Override `GatekeeperConfig` instead of using the default `app.gatekeeper.config` + /// - keyMaker: Override `GatekeeperKeyMaker` instead of using the default `app.gatekeeper.keyMaker` + /// - config: Override the `Error` thrown when the user is rate-limited instead of using the default error. + public init(config: GatekeeperConfig? = nil, keyMaker: GatekeeperKeyMaker? = nil, error: Error? = nil) { self.config = config + self.keyMaker = keyMaker + self.error = error } public func respond(to request: Request, chainingTo next: Responder) -> EventLoopFuture { - request - .gatekeeper(config: config) - .gatekeep(on: request) - .flatMap { next.respond(to: request) } + let gatekeeper = request.gatekeeper(config: config, keyMaker: keyMaker) + + let gatekeep: EventLoopFuture + if let error = error { + gatekeep = gatekeeper.gatekeep(on: request, throwing: error) + } else { + gatekeep = gatekeeper.gatekeep(on: request) + } + + return gatekeep.flatMap { next.respond(to: request) } } } From d0a827323360784d56b8710ab7e562aac286f0ff Mon Sep 17 00:00:00 2001 From: Mads Odgaard Date: Tue, 13 Apr 2021 08:51:58 +0200 Subject: [PATCH 3/3] Simplify return --- Sources/Gatekeeper/Gatekeeper.swift | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/Sources/Gatekeeper/Gatekeeper.swift b/Sources/Gatekeeper/Gatekeeper.swift index 2ad9ab4..893f273 100644 --- a/Sources/Gatekeeper/Gatekeeper.swift +++ b/Sources/Gatekeeper/Gatekeeper.swift @@ -24,18 +24,14 @@ public struct Gatekeeper { else: error ) .map(updateEntry) - .flatMap { entry -> EventLoopFuture in + .flatMap { entry in // The amount of time the entry has existed. let entryLifetime = Int(Date().timeIntervalSince1970 - entry.createdAt.timeIntervalSince1970) // Remaining time until the entry expires. The entry would be expired by cache if it was negative. let timeRemaining = Int(config.refreshInterval) - entryLifetime - - return cache - .set(cacheKey, to: entry, expiresIn: .seconds(timeRemaining)) - .transform(to: entry) + return cache.set(cacheKey, to: entry, expiresIn: .seconds(timeRemaining)) } } - .transform(to: ()) } private func updateEntry(_ entry: Entry) -> Entry {