/*
 * Copyright 2020, 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 Logging
import NIOCore
import NIOHPACK
import NIOHTTP1
import NIOHTTP2
import NIOTLS

/// Configures a server pipeline for gRPC with the appropriate handlers depending on the HTTP
/// version used for transport.
///
/// If TLS is enabled then the handler listens for an 'TLSUserEvent.handshakeCompleted' event and
/// configures the pipeline appropriately for the protocol negotiated via ALPN. If TLS is not
/// configured then the HTTP version is determined by parsing the inbound byte stream.
final class GRPCServerPipelineConfigurator: ChannelInboundHandler, RemovableChannelHandler {
  internal typealias InboundIn = ByteBuffer
  internal typealias InboundOut = ByteBuffer

  /// The server configuration.
  private let configuration: Server.Configuration

  /// A buffer containing the buffered bytes.
  private var buffer: ByteBuffer?

  /// The current state.
  private var state: State

  private enum ALPN {
    /// ALPN is expected. It may or may not be required, however.
    case expected(required: Bool)

    /// ALPN was expected but not required and no protocol was negotiated in the handshake. We may
    /// now fall back to parsing bytes on the connection.
    case expectedButFallingBack

    /// ALPN is not expected; this is a cleartext connection.
    case notExpected
  }

  private enum State {
    /// The pipeline isn't configured yet.
    case notConfigured(alpn: ALPN)
    /// We're configuring the pipeline.
    case configuring
  }

  init(configuration: Server.Configuration) {
    if let tls = configuration.tlsConfiguration {
      self.state = .notConfigured(alpn: .expected(required: tls.requireALPN))
    } else {
      self.state = .notConfigured(alpn: .notExpected)
    }

    self.configuration = configuration
  }

  /// Makes a gRPC idle handler for the server..
  private func makeIdleHandler() -> GRPCIdleHandler {
    return .init(
      idleTimeout: self.configuration.connectionIdleTimeout,
      keepalive: self.configuration.connectionKeepalive,
      logger: self.configuration.logger
    )
  }

  /// Makes an HTTP/2 handler.
  private func makeHTTP2Handler() -> NIOHTTP2Handler {
    var configuration = NIOHTTP2Handler.ConnectionConfiguration()
    configuration.initialSettings = [
      HTTP2Setting(
        parameter: .maxConcurrentStreams,
        value: self.configuration.httpMaxConcurrentStreams
      ),
      HTTP2Setting(
        parameter: .maxHeaderListSize,
        value: HPACKDecoder.defaultMaxHeaderListSize
      ),
      HTTP2Setting(
        parameter: .maxFrameSize,
        value: self.configuration.httpMaxFrameSize
      ),
      HTTP2Setting(
        parameter: .initialWindowSize,
        value: self.configuration.httpTargetWindowSize
      ),
    ]
    configuration.maximumRecentlyResetStreams = self.configuration.httpMaxResetStreams
    return NIOHTTP2Handler(mode: .server, connectionConfiguration: configuration)
  }

  /// Makes an HTTP/2 multiplexer suitable handling gRPC requests.
  private func makeHTTP2Multiplexer(for channel: Channel) -> HTTP2StreamMultiplexer {
    return .init(
      mode: .server,
      channel: channel,
      targetWindowSize: self.configuration.httpTargetWindowSize
    ) { [logger = self.configuration.logger] stream in
      // Sync options were added to the HTTP/2 stream channel in 1.17.0 (we require at least this)
      // so this shouldn't be `nil`, but it's not a problem if it is.
      let http2StreamID = try? stream.syncOptions?.getOption(HTTP2StreamChannelOptions.streamID)
      let streamID =
        http2StreamID.map { streamID in
          return String(Int(streamID))
        } ?? "<unknown>"

      var logger = logger
      logger[metadataKey: MetadataKey.h2StreamID] = "\(streamID)"

      do {
        // TODO: provide user configuration for header normalization.
        let handler = self.makeHTTP2ToRawGRPCHandler(normalizeHeaders: true, logger: logger)
        try stream.pipeline.syncOperations.addHandler(handler)
        return stream.eventLoop.makeSucceededVoidFuture()
      } catch {
        return stream.eventLoop.makeFailedFuture(error)
      }
    }
  }

  /// Makes an HTTP/2 to raw gRPC server handler.
  private func makeHTTP2ToRawGRPCHandler(
    normalizeHeaders: Bool,
    logger: Logger
  ) -> HTTP2ToRawGRPCServerCodec {
    return HTTP2ToRawGRPCServerCodec(
      servicesByName: self.configuration.serviceProvidersByName,
      encoding: self.configuration.messageEncoding,
      errorDelegate: self.configuration.errorDelegate,
      normalizeHeaders: normalizeHeaders,
      maximumReceiveMessageLength: self.configuration.maximumReceiveMessageLength,
      logger: logger
    )
  }

  /// The pipeline finished configuring.
  private func configurationCompleted(result: Result<Void, Error>, context: ChannelHandlerContext) {
    switch result {
    case .success:
      context.pipeline.removeHandler(context: context, promise: nil)
    case let .failure(error):
      self.errorCaught(context: context, error: error)
    }
  }

  /// Configures the pipeline to handle gRPC requests on an HTTP/2 connection.
  private func configureHTTP2(context: ChannelHandlerContext) {
    // We're now configuring the pipeline.
    self.state = .configuring

    // We could use 'Channel.configureHTTP2Pipeline', but then we'd have to find the right handlers
    // to then insert our keepalive and idle handlers between. We can just add everything together.
    let result: Result<Void, Error>

    do {
      // This is only ever called as a result of reading a user inbound event or reading inbound so
      // we'll be on the right event loop and sync operations are fine.
      let sync = context.pipeline.syncOperations
      try sync.addHandler(self.makeHTTP2Handler())
      try sync.addHandler(self.makeIdleHandler())
      try sync.addHandler(self.makeHTTP2Multiplexer(for: context.channel))
      result = .success(())
    } catch {
      result = .failure(error)
    }

    self.configurationCompleted(result: result, context: context)
  }

  /// Configures the pipeline to handle gRPC-Web requests on an HTTP/1 connection.
  private func configureHTTP1(context: ChannelHandlerContext) {
    // We're now configuring the pipeline.
    self.state = .configuring

    let result: Result<Void, Error>
    do {
      // This is only ever called as a result of reading a user inbound event or reading inbound so
      // we'll be on the right event loop and sync operations are fine.
      let sync = context.pipeline.syncOperations
      try sync.configureHTTPServerPipeline(withErrorHandling: true)
      try sync.addHandler(WebCORSHandler(configuration: self.configuration.webCORS))
      let scheme = self.configuration.tlsConfiguration == nil ? "http" : "https"
      try sync.addHandler(GRPCWebToHTTP2ServerCodec(scheme: scheme))
      // There's no need to normalize headers for HTTP/1.
      try sync.addHandler(
        self.makeHTTP2ToRawGRPCHandler(normalizeHeaders: false, logger: self.configuration.logger)
      )
      result = .success(())
    } catch {
      result = .failure(error)
    }

    self.configurationCompleted(result: result, context: context)
  }

  /// Attempts to determine the HTTP version from the buffer and then configure the pipeline
  /// appropriately. Closes the connection if the HTTP version could not be determined.
  private func determineHTTPVersionAndConfigurePipeline(
    buffer: ByteBuffer,
    context: ChannelHandlerContext
  ) {
    switch HTTPVersionParser.determineHTTPVersion(buffer) {
    case .http2:
      self.configureHTTP2(context: context)
    case .http1:
      self.configureHTTP1(context: context)
    case .unknown:
      // Neither H2 nor H1 or the length limit has been exceeded.
      self.configuration.logger.error("Unable to determine http version, closing")
      context.close(mode: .all, promise: nil)
    case .notEnoughBytes:
      ()  // Try again with more bytes.
    }
  }

  /// Handles a 'TLSUserEvent.handshakeCompleted' event and configures the pipeline to handle gRPC
  /// requests.
  private func handleHandshakeCompletedEvent(
    _ event: TLSUserEvent,
    alpnIsRequired: Bool,
    context: ChannelHandlerContext
  ) {
    switch event {
    case let .handshakeCompleted(negotiatedProtocol):
      let tlsVersion = try? context.channel.getTLSVersionSync()
      self.configuration.logger.debug(
        "TLS handshake completed",
        metadata: [
          "alpn": "\(negotiatedProtocol ?? "nil")",
          "tls_version": "\(tlsVersion.map(String.init(describing:)) ?? "nil")",
        ]
      )

      switch negotiatedProtocol {
      case let .some(negotiated):
        if GRPCApplicationProtocolIdentifier.isHTTP2Like(negotiated) {
          self.configureHTTP2(context: context)
        } else if GRPCApplicationProtocolIdentifier.isHTTP1(negotiated) {
          self.configureHTTP1(context: context)
        } else {
          self.configuration.logger.warning("Unsupported ALPN identifier '\(negotiated)', closing")
          context.close(mode: .all, promise: nil)
        }

      case .none:
        if alpnIsRequired {
          self.configuration.logger.warning("No ALPN protocol negotiated, closing'")
          context.close(mode: .all, promise: nil)
        } else {
          self.configuration.logger.warning("No ALPN protocol negotiated'")
          // We're now falling back to parsing bytes.
          self.state = .notConfigured(alpn: .expectedButFallingBack)
          self.tryParsingBufferedData(context: context)
        }
      }

    case .shutdownCompleted:
      // We don't care about this here.
      ()
    }
  }

  /// Try to parse the buffered data to determine whether or not HTTP/2 or HTTP/1 should be used.
  private func tryParsingBufferedData(context: ChannelHandlerContext) {
    if let buffer = self.buffer {
      self.determineHTTPVersionAndConfigurePipeline(buffer: buffer, context: context)
    }
  }

  // MARK: - Channel Handler

  internal func errorCaught(context: ChannelHandlerContext, error: Error) {
    if let delegate = self.configuration.errorDelegate {
      let baseError: Error

      if let errorWithContext = error as? GRPCError.WithContext {
        baseError = errorWithContext.error
      } else {
        baseError = error
      }

      delegate.observeLibraryError(baseError)
    }

    context.close(mode: .all, promise: nil)
  }

  internal func userInboundEventTriggered(context: ChannelHandlerContext, event: Any) {
    switch self.state {
    case let .notConfigured(alpn: .expected(required)):
      if let event = event as? TLSUserEvent {
        self.handleHandshakeCompletedEvent(event, alpnIsRequired: required, context: context)
      }

    case .notConfigured(alpn: .expectedButFallingBack),
      .notConfigured(alpn: .notExpected),
      .configuring:
      ()
    }

    context.fireUserInboundEventTriggered(event)
  }

  internal func channelRead(context: ChannelHandlerContext, data: NIOAny) {
    var buffer = self.unwrapInboundIn(data)
    self.buffer.setOrWriteBuffer(&buffer)

    switch self.state {
    case .notConfigured(alpn: .notExpected),
      .notConfigured(alpn: .expectedButFallingBack):
      // If ALPN isn't expected, or we didn't negotiate via ALPN and we don't require it then we
      // can try parsing the data we just buffered.
      self.tryParsingBufferedData(context: context)

    case .notConfigured(alpn: .expected),
      .configuring:
      // We expect ALPN or we're being configured, just buffer the data, we'll forward it later.
      ()
    }

    // Don't forward the reads: we'll do so when we have configured the pipeline.
  }

  internal func removeHandler(
    context: ChannelHandlerContext,
    removalToken: ChannelHandlerContext.RemovalToken
  ) {
    // Forward any buffered reads.
    if let buffer = self.buffer {
      self.buffer = nil
      context.fireChannelRead(self.wrapInboundOut(buffer))
    }
    context.leavePipeline(removalToken: removalToken)
  }
}

// MARK: - HTTP Version Parser

struct HTTPVersionParser {
  /// HTTP/2 connection preface bytes. See RFC 7540 § 5.3.
  private static let http2ClientMagic = [
    UInt8(ascii: "P"),
    UInt8(ascii: "R"),
    UInt8(ascii: "I"),
    UInt8(ascii: " "),
    UInt8(ascii: "*"),
    UInt8(ascii: " "),
    UInt8(ascii: "H"),
    UInt8(ascii: "T"),
    UInt8(ascii: "T"),
    UInt8(ascii: "P"),
    UInt8(ascii: "/"),
    UInt8(ascii: "2"),
    UInt8(ascii: "."),
    UInt8(ascii: "0"),
    UInt8(ascii: "\r"),
    UInt8(ascii: "\n"),
    UInt8(ascii: "\r"),
    UInt8(ascii: "\n"),
    UInt8(ascii: "S"),
    UInt8(ascii: "M"),
    UInt8(ascii: "\r"),
    UInt8(ascii: "\n"),
    UInt8(ascii: "\r"),
    UInt8(ascii: "\n"),
  ]

  /// Determines whether the bytes in the `ByteBuffer` are prefixed with the HTTP/2 client
  /// connection preface.
  static func prefixedWithHTTP2ConnectionPreface(_ buffer: ByteBuffer) -> SubParseResult {
    let view = buffer.readableBytesView

    guard view.count >= HTTPVersionParser.http2ClientMagic.count else {
      // Not enough bytes.
      return .notEnoughBytes
    }

    let slice = view[view.startIndex ..< view.startIndex.advanced(by: self.http2ClientMagic.count)]
    return slice.elementsEqual(HTTPVersionParser.http2ClientMagic) ? .accepted : .rejected
  }

  enum ParseResult: Hashable {
    case http1
    case http2
    case unknown
    case notEnoughBytes
  }

  enum SubParseResult: Hashable {
    case accepted
    case rejected
    case notEnoughBytes
  }

  private static let maxLengthToCheck = 1024

  static func determineHTTPVersion(_ buffer: ByteBuffer) -> ParseResult {
    switch Self.prefixedWithHTTP2ConnectionPreface(buffer) {
    case .accepted:
      return .http2

    case .notEnoughBytes:
      switch Self.prefixedWithHTTP1RequestLine(buffer) {
      case .accepted:
        // Not enough bytes to check H2, but enough to confirm H1.
        return .http1
      case .notEnoughBytes:
        // Not enough bytes to check H2 or H1.
        return .notEnoughBytes
      case .rejected:
        // Not enough bytes to check H2 and definitely not H1.
        return .notEnoughBytes
      }

    case .rejected:
      switch Self.prefixedWithHTTP1RequestLine(buffer) {
      case .accepted:
        // Not H2, but H1 is confirmed.
        return .http1
      case .notEnoughBytes:
        // Not H2, but not enough bytes to reject H1 yet.
        return .notEnoughBytes
      case .rejected:
        // Not H2 or H1.
        return .unknown
      }
    }
  }

  private static let http1_1 = [
    UInt8(ascii: "H"),
    UInt8(ascii: "T"),
    UInt8(ascii: "T"),
    UInt8(ascii: "P"),
    UInt8(ascii: "/"),
    UInt8(ascii: "1"),
    UInt8(ascii: "."),
    UInt8(ascii: "1"),
  ]

  /// Determines whether the bytes in the `ByteBuffer` are prefixed with an HTTP/1.1 request line.
  static func prefixedWithHTTP1RequestLine(_ buffer: ByteBuffer) -> SubParseResult {
    var readableBytesView = buffer.readableBytesView

    // We don't need to validate the request line, only determine whether we think it's an HTTP1
    // request line. Another handler will parse it properly.

    // From RFC 2616 § 5.1:
    //   Request-Line = Method SP Request-URI SP HTTP-Version CRLF

    // Get through the first space.
    guard readableBytesView.dropPrefix(through: UInt8(ascii: " ")) != nil else {
      let tooLong = buffer.readableBytes > Self.maxLengthToCheck
      return tooLong ? .rejected : .notEnoughBytes
    }

    // Get through the second space.
    guard readableBytesView.dropPrefix(through: UInt8(ascii: " ")) != nil else {
      let tooLong = buffer.readableBytes > Self.maxLengthToCheck
      return tooLong ? .rejected : .notEnoughBytes
    }

    // +2 for \r\n
    guard readableBytesView.count >= (Self.http1_1.count + 2) else {
      return .notEnoughBytes
    }

    guard let version = readableBytesView.dropPrefix(through: UInt8(ascii: "\r")),
      readableBytesView.first == UInt8(ascii: "\n")
    else {
      // If we didn't drop the prefix OR we did and the next byte wasn't '\n', then we had enough
      // bytes but the '\r\n' wasn't present: reject this as being HTTP1.
      return .rejected
    }

    return version.elementsEqual(Self.http1_1) ? .accepted : .rejected
  }
}

extension Collection where Self == Self.SubSequence, Self.Element: Equatable {
  /// Drops the prefix off the collection up to and including the first `separator`
  /// only if that separator appears in the collection.
  ///
  /// Returns the prefix up to but not including the separator if it was found, nil otherwise.
  mutating func dropPrefix(through separator: Element) -> SubSequence? {
    if self.isEmpty {
      return nil
    }

    guard let separatorIndex = self.firstIndex(of: separator) else {
      return nil
    }

    let prefix = self[..<separatorIndex]
    self = self[self.index(after: separatorIndex)...]
    return prefix
  }
}