/*
 * Copyright 2022-present 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.
 */

#include "test-mongocrypt.h"

/* Test authentication with the "gcp" KMS provider. */

static void _test_createdatakey_with_credentials(_mongocrypt_tester_t *tester) {
    mongocrypt_t *crypt;
    mongocrypt_ctx_t *ctx;
    mongocrypt_kms_ctx_t *kms;
    const char *kek = "{"
                      "'provider': 'gcp',"
                      "'projectId': 'test-projectId',"
                      "'location': 'test-location',"
                      "'keyRing': 'test-keyRing',"
                      "'keyName': 'test-keyName'"
                      "}";

    crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_DEFAULT);
    ctx = mongocrypt_ctx_new(crypt);
    ASSERT_OK(mongocrypt_ctx_setopt_key_encryption_key(ctx, TEST_BSON(kek)), ctx);
    ASSERT_OK(mongocrypt_ctx_datakey_init(ctx), ctx);

    ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_NEED_KMS);
    /* Assert first CTX_NEED_KMS state requests access token. */
    {
        kms = mongocrypt_ctx_next_kms_ctx(ctx);
        ASSERT_OR_PRINT_MSG(kms, "expected KMS context, got NULL");

        const char *endpoint;
        mongocrypt_kms_ctx_endpoint(kms, &endpoint);
        ASSERT_STREQUAL("oauth2.googleapis.com:443", endpoint);

        /* Satisfy request. */
        ASSERT_OK(mongocrypt_kms_ctx_feed(kms, TEST_FILE("./test/data/gcp-auth/oauth-response.txt")), kms);
        ASSERT_CMPINT((int)mongocrypt_kms_ctx_bytes_needed(kms), ==, 0);

        kms = mongocrypt_ctx_next_kms_ctx(ctx);
        ASSERT_OR_PRINT_MSG(NULL == kms, "expected NULL KMS context, got non-NULL");
        ASSERT_OK(mongocrypt_ctx_kms_done(ctx), ctx);
    }

    /* Assert second CTX_NEED_KMS state requests encryption. */
    {
        kms = mongocrypt_ctx_next_kms_ctx(ctx);
        ASSERT_OR_PRINT_MSG(kms, "expected KMS context, got NULL");

        const char *endpoint;
        mongocrypt_kms_ctx_endpoint(kms, &endpoint);
        ASSERT_STREQUAL("cloudkms.googleapis.com:443", endpoint);

        /* Satisfy request. */
        ASSERT_OK(mongocrypt_kms_ctx_feed(kms, TEST_FILE("./test/data/gcp-auth/encrypt-response.txt")), kms);
        ASSERT_CMPINT((int)mongocrypt_kms_ctx_bytes_needed(kms), ==, 0);

        kms = mongocrypt_ctx_next_kms_ctx(ctx);
        ASSERT_OR_PRINT_MSG(NULL == kms, "expected NULL KMS context, got non-NULL");
        ASSERT_OK(mongocrypt_ctx_kms_done(ctx), ctx);
    }

    ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_READY);

    mongocrypt_ctx_destroy(ctx);
    mongocrypt_destroy(crypt);
}

static void _test_encrypt_with_credentials(_mongocrypt_tester_t *tester) {
    mongocrypt_t *crypt;
    mongocrypt_ctx_t *ctx;
    mongocrypt_kms_ctx_t *kms;
    mongocrypt_binary_t *uuid;
    const char *uuid_data = "\x61\x61\x61\x61\x61\x61\x61\x61\x61\x61\x61\x61\x61\x61\x61\x61";

    crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_DEFAULT);
    ctx = mongocrypt_ctx_new(crypt);
    uuid = mongocrypt_binary_new_from_data((uint8_t *)uuid_data, UUID_LEN);
    ASSERT_OK(mongocrypt_ctx_setopt_key_id(ctx, uuid), ctx);
    ASSERT_OK(mongocrypt_ctx_setopt_algorithm(ctx, MONGOCRYPT_ALGORITHM_DETERMINISTIC_STR, -1), ctx);
    ASSERT_OK(mongocrypt_ctx_explicit_encrypt_init(ctx, TEST_BSON("{'v': 1}")), ctx);

    ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_NEED_MONGO_KEYS);
    {
        ASSERT_OK(mongocrypt_ctx_mongo_feed(ctx, TEST_FILE("./test/data/key-document-gcp.json")), ctx);
        ASSERT_OK(mongocrypt_ctx_mongo_done(ctx), ctx);
    }

    ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_NEED_KMS);
    /* Assert first CTX_NEED_KMS state requests access token. */
    {
        kms = mongocrypt_ctx_next_kms_ctx(ctx);
        ASSERT_OR_PRINT_MSG(kms, "expected KMS context, got NULL");

        const char *endpoint;
        mongocrypt_kms_ctx_endpoint(kms, &endpoint);
        ASSERT_STREQUAL("oauth2.googleapis.com:443", endpoint);

        /* Satisfy request. */
        ASSERT_OK(mongocrypt_kms_ctx_feed(kms, TEST_FILE("./test/data/gcp-auth/oauth-response.txt")), kms);
        ASSERT_CMPINT((int)mongocrypt_kms_ctx_bytes_needed(kms), ==, 0);

        kms = mongocrypt_ctx_next_kms_ctx(ctx);
        ASSERT_OR_PRINT_MSG(NULL == kms, "expected NULL KMS context, got non-NULL");
        ASSERT_OK(mongocrypt_ctx_kms_done(ctx), ctx);
    }

    /* Assert second CTX_NEED_KMS state requests decryption. */
    {
        kms = mongocrypt_ctx_next_kms_ctx(ctx);
        ASSERT_OR_PRINT_MSG(kms, "expected KMS context, got NULL");

        const char *endpoint;
        mongocrypt_kms_ctx_endpoint(kms, &endpoint);
        ASSERT_STREQUAL("cloudkms.googleapis.com:443", endpoint);

        /* Satisfy request. */
        ASSERT_OK(mongocrypt_kms_ctx_feed(kms, TEST_FILE("./test/data/gcp-auth/decrypt-response.txt")), kms);
        ASSERT_CMPINT((int)mongocrypt_kms_ctx_bytes_needed(kms), ==, 0);

        kms = mongocrypt_ctx_next_kms_ctx(ctx);
        ASSERT_OR_PRINT_MSG(NULL == kms, "expected NULL KMS context, got non-NULL");
        ASSERT_OK(mongocrypt_ctx_kms_done(ctx), ctx);
    }

    ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_READY);

    mongocrypt_binary_destroy(uuid);
    mongocrypt_ctx_destroy(ctx);
    mongocrypt_destroy(crypt);
}

static void _test_createdatakey_with_accesstoken(_mongocrypt_tester_t *tester) {
    mongocrypt_t *crypt;
    mongocrypt_ctx_t *ctx;
    mongocrypt_kms_ctx_t *kms;
    const char *kek = "{"
                      "'provider': 'gcp',"
                      "'projectId': 'test-projectId',"
                      "'location': 'test-location',"
                      "'keyRing': 'test-keyRing',"
                      "'keyName': 'test-keyName'"
                      "}";

    crypt = mongocrypt_new();
    ASSERT_OK(mongocrypt_setopt_kms_providers(crypt, TEST_BSON("{'gcp': {}}")), crypt);
    mongocrypt_setopt_use_need_kms_credentials_state(crypt);
    ASSERT_OK(mongocrypt_init(crypt), crypt);
    ctx = mongocrypt_ctx_new(crypt);
    ASSERT_OK(mongocrypt_ctx_setopt_key_encryption_key(ctx, TEST_BSON(kek)), ctx);
    ASSERT_OK(mongocrypt_ctx_datakey_init(ctx), ctx);

    ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_NEED_KMS_CREDENTIALS);
    { ASSERT_OK(mongocrypt_ctx_provide_kms_providers(ctx, TEST_BSON("{'gcp': { 'accessToken': 'foobar' } }")), ctx); }

    /* Assert first CTX_NEED_KMS state requests encryption. */
    {
        kms = mongocrypt_ctx_next_kms_ctx(ctx);
        ASSERT_OR_PRINT_MSG(kms, "expected KMS context, got NULL");

        const char *endpoint;
        mongocrypt_kms_ctx_endpoint(kms, &endpoint);
        ASSERT_STREQUAL("cloudkms.googleapis.com:443", endpoint);

        /* Satisfy request. */
        ASSERT_OK(mongocrypt_kms_ctx_feed(kms, TEST_FILE("./test/data/gcp-auth/encrypt-response.txt")), kms);
        ASSERT_CMPINT((int)mongocrypt_kms_ctx_bytes_needed(kms), ==, 0);

        kms = mongocrypt_ctx_next_kms_ctx(ctx);
        ASSERT_OR_PRINT_MSG(NULL == kms, "expected NULL KMS context, got non-NULL");
        ASSERT_OK(mongocrypt_ctx_kms_done(ctx), ctx);
    }

    ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_READY);

    mongocrypt_ctx_destroy(ctx);
    mongocrypt_destroy(crypt);
}

static void _test_encrypt_with_accesstoken(_mongocrypt_tester_t *tester) {
    mongocrypt_t *crypt;
    mongocrypt_ctx_t *ctx;
    mongocrypt_kms_ctx_t *kms;
    mongocrypt_binary_t *uuid;
    const char *uuid_data = "\x61\x61\x61\x61\x61\x61\x61\x61\x61\x61\x61\x61\x61\x61\x61\x61";

    crypt = mongocrypt_new();
    ASSERT_OK(mongocrypt_setopt_kms_providers(crypt, TEST_BSON("{'gcp': {}}")), crypt);
    mongocrypt_setopt_use_need_kms_credentials_state(crypt);
    ASSERT_OK(mongocrypt_init(crypt), crypt);
    ctx = mongocrypt_ctx_new(crypt);
    uuid = mongocrypt_binary_new_from_data((uint8_t *)uuid_data, UUID_LEN);
    ASSERT_OK(mongocrypt_ctx_setopt_key_id(ctx, uuid), ctx);
    ASSERT_OK(mongocrypt_ctx_setopt_algorithm(ctx, MONGOCRYPT_ALGORITHM_DETERMINISTIC_STR, -1), ctx);
    ASSERT_OK(mongocrypt_ctx_explicit_encrypt_init(ctx, TEST_BSON("{'v': 1}")), ctx);

    ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_NEED_KMS_CREDENTIALS);
    { ASSERT_OK(mongocrypt_ctx_provide_kms_providers(ctx, TEST_BSON("{'gcp': { 'accessToken': 'foobar' } }")), ctx); }

    ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_NEED_MONGO_KEYS);
    {
        ASSERT_OK(mongocrypt_ctx_mongo_feed(ctx, TEST_FILE("./test/data/key-document-gcp.json")), ctx);
        ASSERT_OK(mongocrypt_ctx_mongo_done(ctx), ctx);
    }

    ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_NEED_KMS);

    /* Assert first CTX_NEED_KMS state requests decryption. */
    {
        kms = mongocrypt_ctx_next_kms_ctx(ctx);
        ASSERT_OR_PRINT_MSG(kms, "expected KMS context, got NULL");

        const char *endpoint;
        mongocrypt_kms_ctx_endpoint(kms, &endpoint);
        ASSERT_STREQUAL("cloudkms.googleapis.com:443", endpoint);

        /* Satisfy request. */
        ASSERT_OK(mongocrypt_kms_ctx_feed(kms, TEST_FILE("./test/data/gcp-auth/decrypt-response.txt")), kms);
        ASSERT_CMPINT((int)mongocrypt_kms_ctx_bytes_needed(kms), ==, 0);

        kms = mongocrypt_ctx_next_kms_ctx(ctx);
        ASSERT_OR_PRINT_MSG(NULL == kms, "expected NULL KMS context, got non-NULL");
        ASSERT_OK(mongocrypt_ctx_kms_done(ctx), ctx);
    }

    ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_READY);

    mongocrypt_binary_destroy(uuid);
    mongocrypt_ctx_destroy(ctx);
    mongocrypt_destroy(crypt);
}

void _mongocrypt_tester_install_gcp_auth(_mongocrypt_tester_t *tester) {
    INSTALL_TEST(_test_createdatakey_with_credentials);
    INSTALL_TEST(_test_encrypt_with_credentials);
    INSTALL_TEST(_test_createdatakey_with_accesstoken);
    INSTALL_TEST(_test_encrypt_with_accesstoken);
}