Vapor 4 Gzip middleware

GzipMiddleware.swift

//
//  GzipMiddleware.swift
//  GzipMiddleware
//
//  Created by 周楷雯 on 2021/9/8.
//

// MARK: use https://docs.vapor.codes/4.0/server/#response-compression instead of this script

// MARK: Use this at end of app.middleware
// app.middleware.use(gzipMiddleware, at: .end)

import Vapor
import Gzip  // .package(url: "https://github.com/1024jp/GzipSwift.git", from: "5.1.1"),

struct GzipMiddleware: Middleware {
    public func respond(to req: Request, chainingTo next: Responder) -> EventLoopFuture<Response> {
        /// Parse Gzip request
        if req.headers["Content-Encoding"].joined(separator: ",").contains("gzip") {
            if let dataBytes = req.body.data {
                let data = Data(dataBytes.readableBytesView)
                guard data.isGzipped else {
                    return req.eventLoop.makeFailedFuture(Abort(.notAcceptable, reason: "gzip data is not zipped"))
                }
                do {
                    let unzipped = try data.gunzipped()
                    let buffer = ByteBuffer.init(data: unzipped)
                    let newRequest = Request(application: req.application, method: req.method, url: req.url, version: req.version, headers: req.headers, collectedBody: buffer, remoteAddress: req.remoteAddress, logger: req.logger, on: req.eventLoop)
                    return next.respond(to: newRequest)
                } catch {
                    return req.eventLoop.makeFailedFuture(Abort(.notAcceptable, reason: "gzip parse error"))
                }
            } else {
                return req.eventLoop.makeFailedFuture(Abort(.notAcceptable, reason: "gzip data is empty"))
            }
        }

        // Gzip Response

        let acceptsGzip = req.headers["Accept-Encoding"].joined(separator: ",").contains("gzip")

        guard acceptsGzip && shouldGzip(req) else {
            return next.respond(to: req)
        }

        return next.respond(to: req).flatMapThrowing { response throws in
            if let unzipped = response.body.data {
                response.headers.replaceOrAdd(name: "Content-Encoding", value: "gzip")
                response.body = .init(data: try unzipped.gzipped())
            }
            return response
        }
    }
    
    private let shouldGzip: (_ request: Request) -> 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
    }

}