From eaca6dde6aa72154143f20d2b612994b7f6d62f5 Mon Sep 17 00:00:00 2001
From: Matthias Neeracher <microtherion@gmail.com>
Date: Sun, 16 Feb 2020 22:20:02 +0100
Subject: [PATCH] Add shouldGzipResponse

---
 Sources/GzipMiddleware/GzipMiddleware.swift | 27 ++++++++++++++++-----
 1 file changed, 21 insertions(+), 6 deletions(-)

diff --git a/Sources/GzipMiddleware/GzipMiddleware.swift b/Sources/GzipMiddleware/GzipMiddleware.swift
index e365edf..ce51682 100644
--- a/Sources/GzipMiddleware/GzipMiddleware.swift
+++ b/Sources/GzipMiddleware/GzipMiddleware.swift
@@ -3,7 +3,7 @@ import Vapor
 import Foundation
 import ZIPFoundation
 
-/// Server Gzip middlere:
+/// Server Gzip middleware:
 /// 1. checks if the "Accept-Encoding" header contains "gzip"
 /// 2. if so, compresses the body and sets the response header "Content-Encoding" to "gzip",
 public struct GzipServerMiddleware: Middleware, ServiceType {
@@ -11,22 +11,37 @@ public struct GzipServerMiddleware: Middleware, ServiceType {
         return .init()
     }
 
-    private let shouldGzip: (_ request: Request) -> Bool
+    private let shouldGzipRequest: (_ request: Request) -> Bool
+    private let shouldGzipResponse: (_ response: Response) -> Bool
 
     /// The `shouldGzip` closure is asked for every request whether that request
     /// should allow response gzipping. Returns `true` always by default.
     public init(shouldGzip: @escaping (_ request: Request) -> Bool = { _ in true }) {
-        self.shouldGzip = shouldGzip
+        self.shouldGzipRequest = shouldGzip
+        self.shouldGzipResponse = { _ in true }
+    }
+
+    /// The `shouldGzipRequest` closure is asked for every request whether that request
+    /// should allow response gzipping.  `shouldGzipResponse` asks the same for the response.
+    /// Both return`true` always by default.
+    public init(shouldGzipRequest: @escaping (_ request: Request) -> Bool = { _ in true },
+                shouldGzipResponse: @escaping (_ response: Response) -> Bool = { _ in true }
+    ) {
+        self.shouldGzipRequest = shouldGzipRequest
+        self.shouldGzipResponse = shouldGzipResponse
     }
 
     public func respond(to request: Request, chainingTo next: Responder) throws -> EventLoopFuture<Response> {
         let acceptsGzip = request.http.headers[.acceptEncoding].first?.contains("gzip") ?? false
-        guard acceptsGzip && shouldGzip(request) else {
-            return try next.respond(to: request)
+        let response = try next.respond(to: request)
+        guard acceptsGzip && shouldGzipRequest(request) else {
+            return response
         }
 
-        let response = try next.respond(to: request)
         return response.flatMap { response in
+            guard self.shouldGzipResponse(response) else {
+                return request.future(response)
+            }
             var headers = response.http.headers
             headers.replaceOrAdd(name: .contentEncoding, value: "gzip")
             return response.http.body.consumeData(on: request).map { data in