JSONFilter.swift

// https://github.com/vapor/vapor
import Vapor

// https://github.com/Ikiga/IkigaJSON
import IkigaJSON

public struct JSONSafetyMiddleware: Middleware {
    public init() {}
    
    public func respond(to request: Request, chainingTo next: Responder) throws -> EventLoopFuture<Response> {
        return try next.respond(to: request).flatMap { response in
            return response.modifyingJSONObject { object in
                try request.filter().apply(to: &object)
            }
        }
    }
}


public final class JSONFilterProvider: Provider {
    public func register(_ services: inout Services) throws {
        services.register(factory: JSONFilter.init)
    }
    
    public func didBoot(_ container: Container) throws -> EventLoopFuture<Void> {
        return container.future()
    }
    
    public init() {}
}

extension Vapor.Request {
    public func filter() throws -> JSONFilter {
        return try privateContainer.make(JSONFilter.self)
    }
}

public final class JSONFilter: Service {
    public enum FilterMode {
        case whitelistKeys(Set<String>)
        case blacklistKeys(Set<String>)
    }
    
    internal init(container: Container) {}
    internal init() {}
    
    public var mode = FilterMode.whitelistKeys([])
    public var array = false
    
    public func apply(to array: inout JSONArray) {
        for i in 0..<array.count {
            if var object = array[i].object {
                apply(to: &object)
                array[i] = object
            }
        }
    }
    
    public func apply(to object: inout JSONObject) {
        switch mode {
        case .whitelistKeys(let fields):
            for key in object.keys where !fields.contains(key) {
                object[key] = nil
            }
        case .blacklistKeys(let fields):
            let keys = object.keys
            for key in fields where keys.contains(key) {
                object[key] = nil
            }
        }
    }
}

internal let allocator = ByteBufferAllocator()

extension HTTPMessageContainer {
    public func decodeJSONObject() -> Future<JSONObject> {
        return self.http.body.consumeData(on: self).map { data in
            var buffer = allocator.buffer(capacity: data.count)
            buffer.write(bytes: data)
            return try JSONObject(buffer: buffer)
        }
    }
    
    public func setJSON(to object: JSONObject) {
        http.body = HTTPBody(buffer: object.jsonBuffer)
    }
    
    public func modifyingJSONObject(_ run: @escaping (inout JSONObject) throws -> ()) -> Future<Self> {
        return self.decodeJSONObject().map { object in
            var object = object
            try run(&object)
            self.setJSON(to: object)
            return self
        }
    }
}