# frozen_string_literal: true

# Copyright (C) 2019-2020 MongoDB Inc.
#
# 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.

module Mongo
  module Crypt
    # A wrapper around mongocrypt_ctx_t, which manages the
    # state machine for encryption and decryption.
    #
    # This class is a superclass that defines shared methods
    # amongst contexts that are initialized for different purposes
    # (e.g. data key creation, encryption, explicit encryption, etc.)
    #
    # @api private
    class Context
      extend Forwardable

      def_delegators :@mongocrypt_handle, :kms_providers

      #  Create a new Context object
      #
      # @param [ Mongo::Crypt::Handle ] mongocrypt_handle A handle to libmongocrypt
      #   used to create a new context object.
      # @param [ ClientEncryption::IO ] io An instance of the IO class
      #   that implements driver I/O methods required to run the
      #   state machine.
      def initialize(mongocrypt_handle, io)
        @mongocrypt_handle = mongocrypt_handle
        # Ideally, this level of the API wouldn't be passing around pointer
        # references between objects, so this method signature is subject to change.

        # FFI::AutoPointer uses a custom release strategy to automatically free
        # the pointer once this object goes out of scope
        @ctx_p = FFI::AutoPointer.new(
          Binding.mongocrypt_ctx_new(@mongocrypt_handle.ref),
          Binding.method(:mongocrypt_ctx_destroy)
        )
        @encryption_io = io
        @cached_azure_token = nil
      end

      attr_reader :ctx_p

      # Returns the state of the mongocrypt_ctx_t
      #
      # @return [ Symbol ] The context state
      def state
        Binding.mongocrypt_ctx_state(@ctx_p)
      end

      # Runs the mongocrypt_ctx_t state machine and handles
      # all I/O on behalf of
      #
      # @param [ CsotTimeoutHolder ] timeout_holder CSOT timeouts for the
      #   operation the state.
      #
      # @return [ BSON::Document ] A BSON document representing the outcome
      #   of the state machine. Contents can differ depending on how the
      #   context was initialized..
      #
      # @raise [ Error::CryptError ] If the state machine enters the
      #   :error state
      #
      # This method is not currently unit tested. It is integration tested
      # in spec/integration/explicit_encryption_spec.rb
      def run_state_machine(timeout_holder)
        loop do
          timeout_ms = timeout_holder.remaining_timeout_ms!
          case state
          when :error
            Binding.check_ctx_status(self)
          when :ready
            # Finalize the state machine and return the result as a BSON::Document
            return Binding.ctx_finalize(self)
          when :done
            return nil
          when :need_mongo_keys
            provide_keys(timeout_ms)
          when :need_mongo_collinfo
            provide_collection_info(timeout_ms)
          when :need_mongo_collinfo_with_db
            provide_collection_info_with_db(timeout_ms)
          when :need_mongo_markings
            provide_markings(timeout_ms)
          when :need_kms
            feed_kms
          when :need_kms_credentials
            Binding.ctx_provide_kms_providers(
              self,
              retrieve_kms_credentials(timeout_holder).to_document
            )
          else
            raise Error::CryptError.new(
              "State #{state} is not supported by Mongo::Crypt::Context"
            )
          end
        end
      end

      # Which BSON mode to use when creating documents from the outcome of
      # the state machine. The default is :bson, which creates BSON::Document
      # with BSON types. Subclasses may override this method to
      # change this behavior.
      #
      # @return [ Symbol, nil ] The BSON mode.
      def bson_mode
        :bson
      end

      private

      def provide_markings(timeout_ms)
        cmd = Binding.ctx_mongo_op(self)

        result = @encryption_io.mark_command(cmd, db_name: @db_name, timeout_ms: timeout_ms)
        mongocrypt_feed(result)

        mongocrypt_done
      end

      def provide_collection_info(timeout_ms)
        feed_collection_info(@db_name, timeout_ms)
      end

      def provide_collection_info_with_db(timeout_ms)
        feed_collection_info(Binding.ctx_mongo_db(self), timeout_ms)
      end

      def feed_collection_info(db_name, timeout_ms)
        filter = Binding.ctx_mongo_op(self)

        @encryption_io.collection_info(db_name, filter, timeout_ms: timeout_ms).each do |result|
          mongocrypt_feed(result)
        end

        mongocrypt_done
      end

      def provide_keys(timeout_ms)
        filter = Binding.ctx_mongo_op(self)

        @encryption_io.find_keys(filter, timeout_ms: timeout_ms).each do |key|
          mongocrypt_feed(key) if key
        end

        mongocrypt_done
      end

      def feed_kms
        while (kms_context = Binding.ctx_next_kms_ctx(self))
          begin
            delay = Binding.kms_ctx_usleep(kms_context)
            sleep(delay / 1_000_000.0) unless delay.nil?
            provider = Binding.kms_ctx_get_kms_provider(kms_context)
            tls_options = @mongocrypt_handle.kms_tls_options(provider)
            @encryption_io.feed_kms(kms_context, tls_options)
          rescue Error::KmsError => e
            raise unless e.network_error?
            next if Binding.kms_ctx_fail(kms_context)

            raise_kms_retry_error(kms_context, e)
          end
        end
        Binding.ctx_kms_done(self)
      end

      # Raise a KmsError that wraps the KMS status message (which describes the
      # retry exhaustion) with the error from the last attempt.
      #
      # @param [ Mongo::Crypt::KmsContext ] kms_context
      # @param [ Mongo::Error::KmsError ] last_error The error from the last
      #   KMS request attempt.
      #
      # @raise [ Mongo::Error::KmsError ]
      def raise_kms_retry_error(kms_context, last_error)
        status = Binding.kms_ctx_status(kms_context)
        raise Error::KmsError.new(
          "#{status.message}, last attempt failed with: #{last_error.message}",
          code: status.code,
          network_error: true
        )
      end

      # Indicate that state machine is done feeding I/O responses back to libmongocrypt
      def mongocrypt_done
        Binding.mongocrypt_ctx_mongo_done(ctx_p)
      end

      # Feeds the result of a Mongo operation back to libmongocrypt.
      #
      # @param [ Hash ] doc BSON document to feed.
      #
      # @return [ BSON::Document ] BSON document containing the result.
      def mongocrypt_feed(doc)
        Binding.ctx_mongo_feed(self, doc)
      end

      # Retrieves KMS credentials for providers that are configured
      # for automatic credentials retrieval.
      #
      # @param [ CsotTimeoutHolder ] timeout_holder CSOT timeout.
      #
      # @return [ Crypt::KMS::Credentials ] Credentials for the configured
      #   KMS providers.
      def retrieve_kms_credentials(timeout_holder)
        providers = {}
        kms_providers.credentials_map.each do |identifier, creds|
          next unless creds.empty?

          case KMS.provider_base_type(identifier)
          when 'aws'
            begin
              aws_credentials = Mongo::Auth::Aws::CredentialsRetriever.new.credentials(timeout_holder)
            rescue Auth::Aws::CredentialsNotFound
              raise Error::CryptError.new(
                'Could not locate AWS credentials (checked environment variables, ECS and EC2 metadata)'
              )
            end
            providers[identifier] = aws_credentials.to_h
          when 'gcp'
            providers[identifier] = { access_token: gcp_access_token(timeout_holder) }
          when 'azure'
            providers[identifier] = { access_token: azure_access_token(timeout_holder) }
          end
        end
        KMS::Credentials.new(providers)
      end

      # Retrieves a GCP access token.
      #
      # @return [ String ] A GCP access token.
      #
      # @raise [ Error::CryptError ] If the GCP access token could not be
      def gcp_access_token(timeout_holder)
        KMS::GCP::CredentialsRetriever.fetch_access_token(timeout_holder)
      rescue KMS::CredentialsNotFound => e
        raise Error::CryptError.new(
          "Could not locate GCP credentials: #{e.class}: #{e.message}"
        )
      end

      # Returns an Azure access token, retrieving it if necessary.
      #
      # @return [ String ] An Azure access token.
      #
      # @raise [ Error::CryptError ] If the Azure access token could not be
      #   retrieved.
      def azure_access_token(timeout_holder)
        if @cached_azure_token.nil? || @cached_azure_token.expired?
          @cached_azure_token = KMS::Azure::CredentialsRetriever.fetch_access_token(timeout_holder: timeout_holder)
        end
        @cached_azure_token.access_token
      rescue KMS::CredentialsNotFound => e
        raise Error::CryptError.new(
          "Could not locate Azure credentials: #{e.class}: #{e.message}"
        )
      end
    end
  end
end
