Merge pull request #21 from madsodgaard/master

Add expiration to cache entry

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(

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,

Sources/Gatekeeper/Gatekeeper.swift

@@ -11,29 +11,31 @@ public struct Gatekeeper {
         self.keyMaker = identifier
     }
 
-    public func gatekeep(on req: Request) -> EventLoopFuture<Void> {
+    public func gatekeep(
+        on req: Request,
+        throwing error: Error = Abort(.tooManyRequests, reason: "Slow down. You sent too many requests.")
+    ) -> EventLoopFuture<Void> {
         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)
-                            .transform(to: entry)
+                        // 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))
                     }
             }
-            .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
     }

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
+        }
     }
 }

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<Response> {
-        request
-            .gatekeeper(config: config)
-            .gatekeep(on: request)
-            .flatMap { next.respond(to: request) }
+        let gatekeeper = request.gatekeeper(config: config, keyMaker: keyMaker)
+
+        let gatekeep: EventLoopFuture<Void>
+        if let error = error {
+            gatekeep = gatekeeper.gatekeep(on: request, throwing: error)
+        } else {
+            gatekeep = gatekeeper.gatekeep(on: request)
+        }
+
+        return gatekeep.flatMap { next.respond(to: request) }
     }
 }

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)] = [