Update Gatekeeper for Vapor-3

Author: cweinberger

Sources/Gatekeeper/Gatekeeper.swift

@@ -1,90 +1,77 @@
-import HTTP
-import Cache
 import Vapor
-import Foundation
 
-public struct Rate {
-    public enum Interval {
-        case second
-        case minute
-        case hour
-        case day
-    }
-    
-    public let limit: Int
-    public let interval: Interval
-    
-    public init(_ limit: Int, per interval: Interval) {
-        self.limit = limit
-        self.interval = interval
-    }
-    
-    internal var refreshInterval: Double {
-        switch interval {
-        case .second:
-            return 1
-        case .minute:
-            return 60
-        case .hour:
-            return 3_600
-        case .day:
-            return 86_400
-        }
-    }
-}
+public struct Gatekeeper: Service {
+
+    internal let config: GatekeeperConfig
+    internal let cacheFactory: ((Container) throws -> KeyedCache)
 
-public struct Gatekeeper: Middleware {
-    internal var cache: CacheProtocol
-    
-    internal let limit: Int
-    internal let refreshInterval: Double
-    
-    public init(rate: Rate, cache: CacheProtocol = MemoryCache()) {
-        self.cache = cache
-        self.limit = rate.limit
-        self.refreshInterval = rate.refreshInterval
+    public init(
+        config: GatekeeperConfig,
+        cacheFactory: @escaping ((Container) throws -> KeyedCache) = { container in try container.make() }
+    ) {
+        self.config = config
+        self.cacheFactory = cacheFactory
     }
-    
-    public func respond(to request: Request, chainingTo next: Responder) throws -> Response {
-        guard let peer = request.peerHostname else {
+
+    public func accessEndpoint(
+        on request: Request
+    ) throws -> Future<Gatekeeper.Entry> {
+
+        guard let peerHostName = request.http.remotePeer.hostname else {
             throw Abort(
                 .forbidden,
-                metadata: nil,
-                reason: "Unable to verify peer."
+                reason: "Unable to verify peer"
             )
         }
-        
-        var entry = try cache.get(peer)
-        var createdAt = entry?["createdAt"]?.double ?? Date().timeIntervalSince1970
-        var requestsLeft = entry?["requestsLeft"]?.int ?? limit
-        
-        let now = Date().timeIntervalSince1970
-        if now - createdAt >= refreshInterval {
-            createdAt = now
-            requestsLeft = limit
-        }
-        
-        defer {
-            do {
-                try cache.set(peer, Node(node: [
-                    "createdAt": createdAt,
-                    "requestsLeft": requestsLeft
-                ]))
-            } catch {
-                print("WARNING: cache failed: \(error)")
+
+        let peerCacheKey = cacheKey(for: peerHostName)
+        let cache = try cacheFactory(request)
+
+        return cache.get(peerCacheKey, as: Entry.self)
+            .map(to: Entry.self) { entry in
+                if let entry = entry {
+                    return entry
+                } else {
+                    return Entry(
+                        peerHostname: peerHostName,
+                        createdAt: Date(),
+                        requestsLeft: self.config.limit
+                    )
+                }
             }
-        }
-        
-        requestsLeft -= 1
-        guard requestsLeft >= 0 else {
-            throw Abort(
-                .tooManyRequests,
-                metadata: nil,
-                reason: "Slow down."
-            )
-        }
-        
-        let response = try next.respond(to: request)
-        return response
+            .map(to: Entry.self) { entry in
+
+                let now = Date()
+                var mutableEntry = entry
+                if now.timeIntervalSince1970 - entry.createdAt.timeIntervalSince1970 >= self.config.refreshInterval {
+                    mutableEntry.createdAt = now
+                    mutableEntry.requestsLeft = self.config.limit
+                }
+                mutableEntry.requestsLeft -= 1
+                return mutableEntry
+            }.then { entry in
+                return cache.set(peerCacheKey, to: entry).transform(to: entry)
+            }.map(to: Entry.self) { entry in
+
+                if entry.requestsLeft <= 0 {
+                    throw Abort(
+                        .tooManyRequests,
+                        reason: "Patience you must have, my young Padawan."
+                    )
+                }
+                return entry
+            }
+    }
+
+    private func cacheKey(for hostname: String) -> String {
+        return "gatekeeper_\(hostname)"
+    }
+}
+
+extension Gatekeeper {
+    public struct Entry: Codable {
+        let peerHostname: String
+        var createdAt: Date
+        var requestsLeft: Int
     }
 }

Sources/Gatekeeper/GatekeeperConfig.swift

@@ -0,0 +1,32 @@
+import Vapor
+
+public struct GatekeeperConfig: Service {
+
+    public enum Interval {
+        case second
+        case minute
+        case hour
+        case day
+    }
+
+    public let limit: Int
+    public let interval: Interval
+
+    public init(maxRequests limit: Int, per interval: Interval) {
+        self.limit = limit
+        self.interval = interval
+    }
+
+    internal var refreshInterval: Double {
+        switch interval {
+        case .second:
+            return 1
+        case .minute:
+            return 60
+        case .hour:
+            return 3_600
+        case .day:
+            return 86_400
+        }
+    }
+}

Sources/Gatekeeper/GatekeeperMiddleware.swift

@@ -0,0 +1,25 @@
+import Vapor
+
+public struct GatekeeperMiddleware {
+    let gatekeeper: Gatekeeper
+}
+
+extension GatekeeperMiddleware: Middleware {
+    public func respond(
+        to request: Request,
+        chainingTo next: Responder
+    ) throws -> EventLoopFuture<Response> {
+
+        return try gatekeeper.accessEndpoint(on: request).do { entry in
+            print("Gatekeeper Entry: \(entry)")
+        }.flatMap { _ in
+            return try next.respond(to: request)
+        }
+    }
+}
+
+extension GatekeeperMiddleware: ServiceType {
+    public static func makeService(for container: Container) throws -> GatekeeperMiddleware {
+        return try .init(gatekeeper: container.make())
+    }
+}

Sources/Gatekeeper/GatekeeperProvider.swift

@@ -0,0 +1,33 @@
+import Vapor
+
+public final class GatekeeperProvider {
+
+    internal let config: GatekeeperConfig
+    internal let cacheFactory: ((Container) throws -> KeyedCache)
+
+    public init(
+        config: GatekeeperConfig,
+        cacheFactory: @escaping ((Container) throws -> KeyedCache) = { container in try container.make() }
+    ) {
+        self.config = config
+        self.cacheFactory = cacheFactory
+    }
+}
+
+extension GatekeeperProvider: Provider {
+    public func register(_ services: inout Services) throws {
+        services.register(config)
+        services.register(
+            Gatekeeper(
+                config: config,
+                cacheFactory: cacheFactory
+            ),
+            as: Gatekeeper.self
+        )
+        services.register(GatekeeperMiddleware.self)
+    }
+
+    public func didBoot(_ container: Container) throws -> EventLoopFuture<Void> {
+        return .done(on: container)
+    }
+}