/*

 * Copyright 2019, gRPC Authors All rights reserved.

 *

 * Licensed under the Apache License, Version 2.0 (the "License");

 * you may not use this file except in compliance with the License.

 * You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

import NIOCore

import NIOHTTP1

/// Handler that manages the CORS protocol for requests incoming from the browser.

internal final class WebCORSHandler {

  let configuration: Server.Configuration.CORS

  private var state: State = .idle

  private enum State: Equatable {

    /// Starting state.

    case idle

    /// CORS preflight request is in progress.

    case processingPreflightRequest

    /// "Real" request is in progress.

    case processingRequest(origin: String?)

  }

  init(configuration: Server.Configuration.CORS) {

    self.configuration = configuration

  }

}

extension WebCORSHandler: ChannelInboundHandler {

  typealias InboundIn = HTTPServerRequestPart

  typealias InboundOut = HTTPServerRequestPart

  typealias OutboundOut = HTTPServerResponsePart

  func channelRead(context: ChannelHandlerContext, data: NIOAny) {

    switch self.unwrapInboundIn(data) {

    case let .head(head):

      self.receivedRequestHead(context: context, head)

    case let .body(body):

      self.receivedRequestBody(context: context, body)

    case let .end(trailers):

      self.receivedRequestEnd(context: context, trailers)

    }

  }

  private func receivedRequestHead(context: ChannelHandlerContext, _ head: HTTPRequestHead) {

    if head.method == .OPTIONS,

      head.headers.contains(.accessControlRequestMethod),

      let origin = head.headers.first(name: "origin")

    {

      // If the request is OPTIONS with a access-control-request-method header it's a CORS

      // preflight request and is not propagated further.

      self.state = .processingPreflightRequest

      self.handlePreflightRequest(context: context, head: head, origin: origin)

    } else {

      self.state = .processingRequest(origin: head.headers.first(name: "origin"))

      context.fireChannelRead(self.wrapInboundOut(.head(head)))

    }

  }

  private func receivedRequestBody(context: ChannelHandlerContext, _ body: ByteBuffer) {

    // OPTIONS requests do not have a body, but still handle this case to be

    // cautious.

    if self.state == .processingPreflightRequest {

      return

    }

    context.fireChannelRead(self.wrapInboundOut(.body(body)))

  }

  private func receivedRequestEnd(context: ChannelHandlerContext, _ trailers: HTTPHeaders?) {

    if self.state == .processingPreflightRequest {

      // End of OPTIONS request; reset state and finish the response.

      self.state = .idle

      context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: nil)

    } else {

      context.fireChannelRead(self.wrapInboundOut(.end(trailers)))

    }

  }

  private func handlePreflightRequest(

    context: ChannelHandlerContext,

    head: HTTPRequestHead,

    origin: String

  ) {

    let responseHead: HTTPResponseHead

    if let allowedOrigin = self.configuration.allowedOrigins.header(origin) {

      var headers = HTTPHeaders()

      headers.reserveCapacity(4 + self.configuration.allowedHeaders.count)

      headers.add(name: .accessControlAllowOrigin, value: allowedOrigin)

      headers.add(name: .accessControlAllowMethods, value: "POST")

      for value in self.configuration.allowedHeaders {

        headers.add(name: .accessControlAllowHeaders, value: value)

      }

      if self.configuration.allowCredentialedRequests {

        headers.add(name: .accessControlAllowCredentials, value: "true")

      }

      if self.configuration.preflightCacheExpiration > 0 {

        headers.add(

          name: .accessControlMaxAge,

          value: "\(self.configuration.preflightCacheExpiration)"

        )

      }

      responseHead = HTTPResponseHead(version: head.version, status: .ok, headers: headers)

    } else {

      // Not allowed; respond with 403. This is okay in a pre-flight request.

      responseHead = HTTPResponseHead(version: head.version, status: .forbidden)

    }

    context.write(self.wrapOutboundOut(.head(responseHead)), promise: nil)

  }

}

extension WebCORSHandler: ChannelOutboundHandler {

  typealias OutboundIn = HTTPServerResponsePart

  func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise<Void>?) {

    let responsePart = self.unwrapOutboundIn(data)

    switch responsePart {

    case var .head(responseHead):

      switch self.state {

      case let .processingRequest(origin):

        self.prepareCORSResponseHead(&responseHead, origin: origin)

        context.write(self.wrapOutboundOut(.head(responseHead)), promise: promise)

      case .idle, .processingPreflightRequest:

        assertionFailure("Writing response head when no request is in progress")

        context.close(promise: nil)

      }

    case .body:

      context.write(data, promise: promise)

    case .end:

      self.state = .idle

      context.write(data, promise: promise)

    }

  }

  private func prepareCORSResponseHead(_ head: inout HTTPResponseHead, origin: String?) {

    guard let header = origin.flatMap({ self.configuration.allowedOrigins.header($0) }) else {

      // No origin or the origin is not allowed; don't treat it as a CORS request.

      return

    }

    head.headers.replaceOrAdd(name: .accessControlAllowOrigin, value: header)

    if self.configuration.allowCredentialedRequests {

      head.headers.add(name: .accessControlAllowCredentials, value: "true")

    }

    //! FIXME: Check whether we can let browsers keep connections alive. It's not possible

    // now as the channel has a state that can't be reused since the pipeline is modified to

    // inject the gRPC call handler.

    head.headers.replaceOrAdd(name: "Connection", value: "close")

  }

}

extension HTTPHeaders {

  fileprivate enum CORSHeader: String {

    case accessControlRequestMethod = "access-control-request-method"

    case accessControlRequestHeaders = "access-control-request-headers"

    case accessControlAllowOrigin = "access-control-allow-origin"

    case accessControlAllowMethods = "access-control-allow-methods"

    case accessControlAllowHeaders = "access-control-allow-headers"

    case accessControlAllowCredentials = "access-control-allow-credentials"

    case accessControlMaxAge = "access-control-max-age"

  }

  fileprivate func contains(_ name: CORSHeader) -> Bool {

    return self.contains(name: name.rawValue)

  }

  fileprivate mutating func add(name: CORSHeader, value: String) {

    self.add(name: name.rawValue, value: value)

  }

  fileprivate mutating func replaceOrAdd(name: CORSHeader, value: String) {

    self.replaceOrAdd(name: name.rawValue, value: value)

  }

}

extension Server.Configuration.CORS.AllowedOrigins {

  internal func header(_ origin: String) -> String? {

    switch self.wrapped {

    case .all:

      return "*"

    case .originBased:

      return origin

    case let .only(allowed):

      return allowed.contains(origin) ? origin : nil

    case let .custom(custom):

      return custom.check(origin: origin)

    }

  }

}