/*
 * Copyright 2019-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 "mongocrypt-ciphertext-private.h"
#include "mongocrypt-crypto-private.h"
#include "mongocrypt-marking-private.h"
#include "mongocrypt-private.h"

#include "test-mongocrypt.h"

/* From BSON Binary subtype 6 specification:
struct fle_blob {
 uint8  fle_blob_subtype = (1 or 2);
 uint8  key_uuid[16];
 uint8  original_bson_type;
 uint8  ciphertext[ciphertext_length];
}
*/
static void _test_ciphertext_serialization(_mongocrypt_tester_t *tester) {
    _mongocrypt_ciphertext_t original, returned;
    _mongocrypt_buffer_t serialized;
    char *expected = "\x01\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0A\x0B\x0C"
                     "\x0D\x0E\x0F\x02\x00\x01";
    mongocrypt_status_t *status;

    _mongocrypt_ciphertext_init(&original);
    _mongocrypt_ciphertext_init(&returned);
    _mongocrypt_buffer_init(&serialized);
    status = mongocrypt_status_new();

    original.blob_subtype = 1;
    original.original_bson_type = 2;
    _mongocrypt_tester_fill_buffer(&original.data, 2);
    _mongocrypt_tester_fill_buffer(&original.key_id, 16);

    BSON_ASSERT(_mongocrypt_serialize_ciphertext(&original, &serialized));
    BSON_ASSERT(0 == memcmp(expected, serialized.data, serialized.len));

    /* Now parse it back, should get the same ciphertext. */
    BSON_ASSERT(_mongocrypt_ciphertext_parse_unowned(&serialized, &returned, status));
    BSON_ASSERT(original.blob_subtype == returned.blob_subtype);
    BSON_ASSERT(original.original_bson_type == returned.original_bson_type);
    BSON_ASSERT(0 == _mongocrypt_buffer_cmp(&original.key_id, &returned.key_id));
    BSON_ASSERT(0 == _mongocrypt_buffer_cmp(&original.data, &returned.data));

    _mongocrypt_ciphertext_cleanup(&original);
    _mongocrypt_ciphertext_cleanup(&returned);
    _mongocrypt_buffer_cleanup(&serialized);
    mongocrypt_status_destroy(status);
}

static void _test_malformed_ciphertext(_mongocrypt_tester_t *tester) {
    _mongocrypt_buffer_t serialized;
    _mongocrypt_ciphertext_t returned;
    mongocrypt_status_t *status;

    status = mongocrypt_status_new();
    /* the minimum size for a ciphertext is 19 bytes. */
    _mongocrypt_tester_fill_buffer(&serialized, 18);

    BSON_ASSERT(!_mongocrypt_ciphertext_parse_unowned(&serialized, &returned, status));
    BSON_ASSERT(0 == strcmp(mongocrypt_status_message(status, NULL), "malformed ciphertext, too small"));
    _mongocrypt_buffer_cleanup(&serialized);

    _mongocrypt_tester_fill_buffer(&serialized, 19);
    /* give a valid blob_subtype. */
    serialized.data[0] = 1;
    BSON_ASSERT(_mongocrypt_ciphertext_parse_unowned(&serialized, &returned, status));

    /* now an invalid blob_subtype. */
    serialized.data[0] = 16;
    BSON_ASSERT(!_mongocrypt_ciphertext_parse_unowned(&serialized, &returned, status));
    BSON_ASSERT(
        0 == strcmp(mongocrypt_status_message(status, NULL), "malformed ciphertext, expected blob subtype of 1 or 2"));

    _mongocrypt_ciphertext_cleanup(&returned);
    _mongocrypt_buffer_cleanup(&serialized);
    mongocrypt_status_destroy(status);
}

static void _test_ciphertext_algorithm(_mongocrypt_tester_t *tester) {
    mongocrypt_t *crypt;
    mongocrypt_ctx_t *ctx;
    mongocrypt_status_t *status;
    _mongocrypt_key_broker_t *kb;
    _mongocrypt_buffer_t zeros;
    _mongocrypt_ciphertext_t type1_valueA, type1_valueA_again, type2_valueA, type1_valueB;
    _mongocrypt_marking_t marking = {0};
    bson_iter_t a_iter, b_iter;
    bson_t *bson;
    bool res;

    status = mongocrypt_status_new();
    _mongocrypt_ciphertext_init(&type1_valueA);
    _mongocrypt_ciphertext_init(&type1_valueA_again);
    _mongocrypt_ciphertext_init(&type2_valueA);
    _mongocrypt_ciphertext_init(&type1_valueB);

    _mongocrypt_buffer_init(&zeros);
    _mongocrypt_buffer_resize(&zeros, MONGOCRYPT_IV_LEN);
    memset(zeros.data, 0, MONGOCRYPT_IV_LEN);

    crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_DEFAULT);
    /* use a mongocrypt_ctx_t to get a key broker */
    ctx = mongocrypt_ctx_new(crypt);
    ASSERT_OK(mongocrypt_ctx_encrypt_init(ctx, "test", -1, TEST_FILE("./test/example/cmd.json")), ctx);

    _mongocrypt_buffer_from_binary(&marking.u.fle1.key_id, TEST_BIN(16));
    marking.u.fle1.key_id.subtype = BSON_SUBTYPE_UUID;
    kb = &ctx->kb;
    _mongocrypt_key_broker_add_test_key(kb, &marking.u.fle1.key_id);

    bson = BCON_NEW("v", "a", "v", "b");
    bson_iter_init(&a_iter, bson);
    bson_iter_next(&a_iter);
    bson_iter_init(&b_iter, bson);
    bson_iter_next(&b_iter);
    bson_iter_next(&b_iter);

    /* Marking type = 1, plaintext = a */
    marking.u.fle1.algorithm = 1;
    memcpy(&marking.u.fle1.v_iter, &a_iter, sizeof(bson_iter_t));
    res = _mongocrypt_marking_to_ciphertext((void *)kb, &marking, &type1_valueA, status);
    ASSERT_OR_PRINT(res, status);

    /* Marking type = 1, plaintext = a */
    marking.u.fle1.algorithm = 1;
    memcpy(&marking.u.fle1.v_iter, &a_iter, sizeof(bson_iter_t));
    res = _mongocrypt_marking_to_ciphertext((void *)kb, &marking, &type1_valueA_again, status);
    ASSERT_OR_PRINT(res, status);

    /* Marking type = 2, plaintext = a */
    marking.u.fle1.algorithm = 2;
    memcpy(&marking.u.fle1.v_iter, &a_iter, sizeof(bson_iter_t));
    res = _mongocrypt_marking_to_ciphertext((void *)kb, &marking, &type2_valueA, status);
    ASSERT_OR_PRINT(res, status);

    /* Marking type = 1, plaintext = b */
    marking.u.fle1.algorithm = 1;
    memcpy(&marking.u.fle1.v_iter, &b_iter, sizeof(bson_iter_t));
    res = _mongocrypt_marking_to_ciphertext((void *)kb, &marking, &type1_valueB, status);
    ASSERT_OR_PRINT(res, status);

    /* Shorten all buffers to their IV length's */
    type1_valueA.data.len = MONGOCRYPT_IV_LEN;
    type1_valueA_again.data.len = MONGOCRYPT_IV_LEN;
    type2_valueA.data.len = MONGOCRYPT_IV_LEN;
    type1_valueB.data.len = MONGOCRYPT_IV_LEN;

    BSON_ASSERT(0 != _mongocrypt_buffer_cmp(&type1_valueA.data, &zeros));
    BSON_ASSERT(0 != _mongocrypt_buffer_cmp(&type1_valueA_again.data, &zeros));
    BSON_ASSERT(0 != _mongocrypt_buffer_cmp(&type2_valueA.data, &zeros));
    BSON_ASSERT(0 != _mongocrypt_buffer_cmp(&type1_valueB.data, &zeros));

    /* Type 1 IV should be repeatable for same plaintext. */
    BSON_ASSERT(0 == _mongocrypt_buffer_cmp(&type1_valueA.data, &type1_valueA_again.data));
    /* Type 1 IV should differ from type 2 random IV. */
    BSON_ASSERT(0 != _mongocrypt_buffer_cmp(&type1_valueA.data, &type2_valueA.data));
    /* Type 1 IV should differ if plaintext differs. */
    BSON_ASSERT(0 != _mongocrypt_buffer_cmp(&type1_valueA.data, &type1_valueB.data));

    _mongocrypt_ciphertext_cleanup(&type1_valueA);
    _mongocrypt_ciphertext_cleanup(&type1_valueA_again);
    _mongocrypt_ciphertext_cleanup(&type2_valueA);
    _mongocrypt_ciphertext_cleanup(&type1_valueB);
    _mongocrypt_buffer_cleanup(&zeros);
    _mongocrypt_marking_cleanup(&marking);
    mongocrypt_ctx_destroy(ctx);
    mongocrypt_destroy(crypt);
    bson_destroy(bson);
    mongocrypt_status_destroy(status);
}

static void _test_ciphertext_serialize_associated_data(_mongocrypt_tester_t *tester) {
    _mongocrypt_ciphertext_t ciphertext;
    _mongocrypt_buffer_t serialized;
    /* Expected associated data is:
     * One byte for blob subtype for deterministic encryption:
     * \x01
     * Followed by 16 byte UUID, a repeating 123 pattern:
     * \x01\x02\x03\x01\x02\x03\x01\x02\x03\x01\x02\x03\x01\x02\x03\x01
     * Followed by the BSON type for UTF8
     * \x02
     */
    char *expected = "\x01\x01\x02\x03\x01\x02\x03\x01\x02\x03\x01\x02\x03\x01"
                     "\x02\x03\x01\x02";

    _mongocrypt_ciphertext_init(&ciphertext);
    _mongocrypt_buffer_init(&serialized);

    /* Create a UUID */
    _mongocrypt_buffer_from_binary(&ciphertext.key_id, TEST_BIN(16));
    ciphertext.key_id.subtype = BSON_SUBTYPE_UUID;

    ciphertext.original_bson_type = BSON_TYPE_UTF8;
    ciphertext.blob_subtype = MC_SUBTYPE_FLE1DeterministicEncryptedValue;

    BSON_ASSERT(_mongocrypt_ciphertext_serialize_associated_data(&ciphertext, &serialized));
    BSON_ASSERT(serialized.len == 18);
    BSON_ASSERT(0 == memcmp(serialized.data, expected, strlen(expected)));

    _mongocrypt_ciphertext_cleanup(&ciphertext);
    _mongocrypt_buffer_cleanup(&serialized);
}

void _mongocrypt_tester_install_ciphertext(_mongocrypt_tester_t *tester) {
    INSTALL_TEST(_test_malformed_ciphertext);
    INSTALL_TEST(_test_ciphertext_serialization);
    INSTALL_TEST(_test_ciphertext_algorithm);
    INSTALL_TEST(_test_ciphertext_serialize_associated_data);
}
