/*
 * Copyright 2024, 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 Crypto
import Foundation
import GRPCCore
import GRPCNIOTransportHTTP2Posix
import NIOSSL
import SwiftASN1
import Testing
import X509

@Suite("HTTP/2 transport E2E tests with TLS enabled")
struct HTTP2TransportTLSEnabledTests {
  // - MARK: Tests

  @Test(
    "When using defaults, server does not perform client verification",
    arguments: [TransportKind.posix],
    [TransportKind.posix]
  )
  func testRPC_Defaults_OK(
    clientTransport: TransportKind,
    serverTransport: TransportKind
  ) async throws {
    let certificateKeyPairs = try SelfSignedCertificateKeyPairs()
    let clientConfig = self.makeDefaultTLSClientConfig(
      for: clientTransport,
      certificateKeyPairs: certificateKeyPairs
    )
    let serverConfig = self.makeDefaultTLSServerConfig(
      for: serverTransport,
      certificateKeyPairs: certificateKeyPairs
    )

    try await self.withClientAndServer(
      clientConfig: clientConfig,
      serverConfig: serverConfig
    ) { control in
      await #expect(throws: Never.self) {
        try await self.executeUnaryRPC(control: control)
      }
    }
  }

  @Test(
    "When using mTLS defaults, both client and server verify each others' certificates",
    arguments: [TransportKind.posix],
    [TransportKind.posix]
  )
  func testRPC_mTLS_OK(
    clientTransport: TransportKind,
    serverTransport: TransportKind
  ) async throws {
    let certificateKeyPairs = try SelfSignedCertificateKeyPairs()
    let clientConfig = self.makeMTLSClientConfig(
      for: clientTransport,
      certificateKeyPairs: certificateKeyPairs,
      serverHostname: "localhost"
    )
    let serverConfig = self.makeMTLSServerConfig(
      for: serverTransport,
      certificateKeyPairs: certificateKeyPairs,
      includeClientCertificateInTrustRoots: true
    )

    try await self.withClientAndServer(
      clientConfig: clientConfig,
      serverConfig: serverConfig
    ) { control in
      await #expect(throws: Never.self) {
        try await self.executeUnaryRPC(control: control)
      }
    }
  }

  @Test(
    "Error is surfaced when client fails server verification",
    arguments: [TransportKind.posix],
    [TransportKind.posix]
  )
  // Verification should fail because the custom hostname is missing on the client.
  func testClientFailsServerValidation(
    clientTransport: TransportKind,
    serverTransport: TransportKind
  ) async throws {
    let certificateKeyPairs = try SelfSignedCertificateKeyPairs()
    let clientConfig = self.makeMTLSClientConfig(
      for: clientTransport,
      certificateKeyPairs: certificateKeyPairs,
      serverHostname: "the-wrong-hostname"
    )

    let serverConfig = self.makeMTLSServerConfig(
      for: serverTransport,
      certificateKeyPairs: certificateKeyPairs,
      includeClientCertificateInTrustRoots: true
    )

    try await self.withClientAndServer(
      clientConfig: clientConfig,
      serverConfig: serverConfig
    ) { control in
      await #expect {
        try await self.executeUnaryRPC(control: control)
      } throws: { error in
        guard let rootError = error as? RPCError else {
          Issue.record("Should be an RPC error")
          return false
        }
        #expect(rootError.code == .unavailable)
        #expect(
          rootError.message
            == "The server accepted the TCP connection but closed the connection before completing the HTTP/2 connection preface."
        )

        guard
          let sslError = rootError.cause as? NIOSSLExtraError,
          case .failedToValidateHostname = sslError
        else {
          Issue.record(
            "Should be a NIOSSLExtraError.failedToValidateHostname error, but was: \(String(describing: rootError.cause))"
          )
          return false
        }

        return true
      }
    }
  }

  @Test(
    "Error is surfaced when server fails client verification",
    arguments: [TransportKind.posix],
    [TransportKind.posix]
  )
  // Verification should fail because the server does not have trust roots containing the client cert.
  func testServerFailsClientValidation(
    clientTransport: TransportKind,
    serverTransport: TransportKind
  ) async throws {
    let certificateKeyPairs = try SelfSignedCertificateKeyPairs()
    let clientConfig = self.makeMTLSClientConfig(
      for: clientTransport,
      certificateKeyPairs: certificateKeyPairs,
      serverHostname: "localhost"
    )
    let serverConfig = self.makeMTLSServerConfig(
      for: serverTransport,
      certificateKeyPairs: certificateKeyPairs,
      includeClientCertificateInTrustRoots: false
    )

    try await self.withClientAndServer(
      clientConfig: clientConfig,
      serverConfig: serverConfig
    ) { control in
      await #expect {
        try await self.executeUnaryRPC(control: control)
      } throws: { error in
        guard let rootError = error as? RPCError else {
          Issue.record("Should be an RPC error")
          return false
        }
        #expect(rootError.code == .unavailable)
        #expect(
          rootError.message
            == "The server accepted the TCP connection but closed the connection before completing the HTTP/2 connection preface."
        )

        guard
          let sslError = rootError.cause as? NIOSSL.BoringSSLError,
          case .sslError = sslError
        else {
          Issue.record(
            "Should be a NIOSSL.sslError error, but was: \(String(describing: rootError.cause))"
          )
          return false
        }

        return true
      }
    }
  }

  // - MARK: Test Utilities

  enum TransportKind: Sendable {
    case posix
  }

  struct Config<Transport, Security> {
    var security: Security
    var transport: Transport
  }

  enum ClientConfig {
    typealias Posix = Config<
      HTTP2ClientTransport.Posix.Config,
      HTTP2ClientTransport.Posix.TransportSecurity
    >
    case posix(Posix)
  }

  enum ServerConfig {
    typealias Posix = Config<
      HTTP2ServerTransport.Posix.Config,
      HTTP2ServerTransport.Posix.TransportSecurity
    >
    case posix(Posix)
  }

  private func makeDefaultPlaintextPosixClientConfig() -> ClientConfig.Posix {
    ClientConfig.Posix(
      security: .plaintext,
      transport: .defaults { config in
        config.backoff.initial = .milliseconds(100)
        config.backoff.multiplier = 1
        config.backoff.jitter = 0
      }
    )
  }

  private func makeDefaultTLSClientConfig(
    for transportSecurity: TransportKind,
    certificateKeyPairs: SelfSignedCertificateKeyPairs
  ) -> ClientConfig {
    switch transportSecurity {
    case .posix:
      var config = self.makeDefaultPlaintextPosixClientConfig()
      config.security = .tls {
        $0.trustRoots = .certificates([
          .bytes(certificateKeyPairs.server.certificate, format: .der)
        ])
      }
      config.transport.http2.authority = "localhost"
      return .posix(config)
    }
  }

  private func makeMTLSClientConfig(
    for transportKind: TransportKind,
    certificateKeyPairs: SelfSignedCertificateKeyPairs,
    serverHostname: String?
  ) -> ClientConfig {
    switch transportKind {
    case .posix:
      var config = self.makeDefaultPlaintextPosixClientConfig()
      config.security = .mTLS(
        certificateChain: [.bytes(certificateKeyPairs.client.certificate, format: .der)],
        privateKey: .bytes(certificateKeyPairs.client.key, format: .der)
      ) {
        $0.trustRoots = .certificates([
          .bytes(certificateKeyPairs.server.certificate, format: .der)
        ])
      }
      config.transport.http2.authority = serverHostname
      return .posix(config)
    }
  }

  private func makeDefaultPlaintextPosixServerConfig() -> ServerConfig.Posix {
    ServerConfig.Posix(security: .plaintext, transport: .defaults)
  }

  private func makeDefaultTLSServerConfig(
    for transportKind: TransportKind,
    certificateKeyPairs: SelfSignedCertificateKeyPairs
  ) -> ServerConfig {
    switch transportKind {
    case .posix:
      var config = self.makeDefaultPlaintextPosixServerConfig()
      config.security = .tls(
        certificateChain: [.bytes(certificateKeyPairs.server.certificate, format: .der)],
        privateKey: .bytes(certificateKeyPairs.server.key, format: .der)
      )
      return .posix(config)
    }
  }

  private func makeMTLSServerConfig(
    for transportKind: TransportKind,
    certificateKeyPairs: SelfSignedCertificateKeyPairs,
    includeClientCertificateInTrustRoots: Bool
  ) -> ServerConfig {
    switch transportKind {
    case .posix:
      var config = self.makeDefaultPlaintextPosixServerConfig()
      config.security = .mTLS(
        certificateChain: [.bytes(certificateKeyPairs.server.certificate, format: .der)],
        privateKey: .bytes(certificateKeyPairs.server.key, format: .der)
      ) {
        if includeClientCertificateInTrustRoots {
          $0.trustRoots = .certificates([
            .bytes(certificateKeyPairs.client.certificate, format: .der)
          ])
        }
      }
      return .posix(config)
    }
  }

  func withClientAndServer(
    clientConfig: ClientConfig,
    serverConfig: ServerConfig,
    _ test: (ControlClient) async throws -> Void
  ) async throws {
    try await withThrowingDiscardingTaskGroup { group in
      let server = self.makeServer(config: serverConfig)

      group.addTask {
        try await server.serve()
      }

      guard let address = try await server.listeningAddress?.ipv4 else {
        Issue.record("Unexpected address to connect to")
        return
      }
      let target: any ResolvableTarget = .ipv4(host: address.host, port: address.port)
      let client = try self.makeClient(config: clientConfig, target: target)

      group.addTask {
        try await client.run()
      }

      let control = ControlClient(wrapping: client)
      try await test(control)

      server.beginGracefulShutdown()
      client.beginGracefulShutdown()
    }
  }

  private func makeServer(config: ServerConfig) -> GRPCServer {
    let services = [ControlService()]

    switch config {
    case .posix(let config):
      let server = GRPCServer(
        transport: .http2NIOPosix(
          address: .ipv4(host: "127.0.0.1", port: 0),
          transportSecurity: config.security,
          config: config.transport
        ),
        services: services
      )

      return server
    }
  }

  private func makeClient(
    config: ClientConfig,
    target: any ResolvableTarget
  ) throws -> GRPCClient {
    let transport: any ClientTransport

    switch config {
    case .posix(let config):
      transport = try HTTP2ClientTransport.Posix(
        target: target,
        transportSecurity: config.security,
        config: config.transport,
        serviceConfig: ServiceConfig()
      )
    }

    return GRPCClient(transport: transport)
  }

  private func executeUnaryRPC(control: ControlClient) async throws {
    let input = ControlInput.with { $0.numberOfMessages = 1 }
    let request = ClientRequest(message: input)
    try await control.unary(request: request) { response in
      #expect(throws: Never.self) { try response.message }
    }
  }
}

struct SelfSignedCertificateKeyPairs {
  struct CertificateKeyPair {
    let certificate: [UInt8]
    let key: [UInt8]
  }

  let server: CertificateKeyPair
  let client: CertificateKeyPair

  init() throws {
    let server = try Self.makeSelfSignedDERCertificateAndPrivateKey(name: "Server Certificate")
    let client = try Self.makeSelfSignedDERCertificateAndPrivateKey(name: "Client Certificate")

    self.server = CertificateKeyPair(certificate: server.cert, key: server.key)
    self.client = CertificateKeyPair(certificate: client.cert, key: client.key)
  }

  private static func makeSelfSignedDERCertificateAndPrivateKey(
    name: String
  ) throws -> (cert: [UInt8], key: [UInt8]) {
    let swiftCryptoKey = P256.Signing.PrivateKey()
    let key = Certificate.PrivateKey(swiftCryptoKey)
    let subjectName = try DistinguishedName { CommonName(name) }
    let issuerName = subjectName
    let now = Date()
    let extensions = try Certificate.Extensions {
      Critical(
        BasicConstraints.isCertificateAuthority(maxPathLength: nil)
      )
      Critical(
        KeyUsage(digitalSignature: true, keyCertSign: true)
      )
      Critical(
        try ExtendedKeyUsage([.serverAuth, .clientAuth])
      )
      SubjectAlternativeNames([.dnsName("localhost")])
    }
    let certificate = try Certificate(
      version: .v3,
      serialNumber: Certificate.SerialNumber(),
      publicKey: key.publicKey,
      notValidBefore: now.addingTimeInterval(-60 * 60),
      notValidAfter: now.addingTimeInterval(60 * 60 * 24 * 365),
      issuer: issuerName,
      subject: subjectName,
      signatureAlgorithm: .ecdsaWithSHA256,
      extensions: extensions,
      issuerPrivateKey: key
    )

    var serializer = DER.Serializer()
    try serializer.serialize(certificate)

    let certBytes = serializer.serializedBytes
    let keyBytes = try key.serializeAsPEM().derBytes
    return (certBytes, keyBytes)
  }
}