// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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 "common/encryption_util.h"

#include <gen_cpp/cloud.pb.h>
#include <glog/logging.h>
#include <math.h>
#include <openssl/err.h>
#include <openssl/evp.h>
#include <openssl/ossl_typ.h>
#include <sys/types.h>

#include <cstring>
#include <memory>
#include <random>
#include <string>
#include <string_view>
#include <thread>
#include <unordered_map>

#include "common/config.h"
#include "common/kms.h"
#include "common/logging.h"
#include "common/util.h"
#include "cpp/sync_point.h"
#include "meta-store/keys.h"
#include "meta-store/txn_kv.h"
#include "meta-store/txn_kv_error.h"

namespace doris::cloud {
namespace config {
extern std::string encryption_key;
}; // namespace config

enum class EncryptionMode {
    AES_128_ECB,
    AES_256_ECB,
    AES_128_CBC,
    AES_256_CBC,
    AES_128_CFB,
    AES_256_CFB,
    AES_128_CFB1,
    AES_256_CFB1,
    AES_128_CFB8,
    AES_256_CFB8,
    AES_128_CFB128,
    AES_256_CFB128,
    AES_128_CTR,
    AES_256_CTR,
    AES_128_OFB,
    AES_256_OFB,
};

enum EncryptionState { AES_SUCCESS = 0, AES_BAD_DATA = -1 };

class EncryptionUtil {
public:
    static int encrypt(EncryptionMode mode, const unsigned char* source, uint32_t source_length,
                       const unsigned char* key, uint32_t key_length, const char* iv_str,
                       int iv_input_length, bool padding, unsigned char* encrypt);

    static int decrypt(EncryptionMode mode, const unsigned char* encrypt, uint32_t encrypt_length,
                       const unsigned char* key, uint32_t key_length, const char* iv_str,
                       int iv_input_length, bool padding, unsigned char* decrypt_content);
};

// aes encrypt/dencrypt
static const int ENCRYPTION_MAX_KEY_LENGTH = 256;

const EVP_CIPHER* get_evp_type(const EncryptionMode mode) {
    switch (mode) {
    case EncryptionMode::AES_128_ECB:
        return EVP_aes_128_ecb();
    case EncryptionMode::AES_128_CBC:
        return EVP_aes_128_cbc();
    case EncryptionMode::AES_128_CFB:
        return EVP_aes_128_cfb();
    case EncryptionMode::AES_128_CFB1:
        return EVP_aes_128_cfb1();
    case EncryptionMode::AES_128_CFB8:
        return EVP_aes_128_cfb8();
    case EncryptionMode::AES_128_CFB128:
        return EVP_aes_128_cfb128();
    case EncryptionMode::AES_128_CTR:
        return EVP_aes_128_ctr();
    case EncryptionMode::AES_128_OFB:
        return EVP_aes_128_ofb();
    case EncryptionMode::AES_256_ECB:
        return EVP_aes_256_ecb();
    case EncryptionMode::AES_256_CBC:
        return EVP_aes_256_cbc();
    case EncryptionMode::AES_256_CFB:
        return EVP_aes_256_cfb();
    case EncryptionMode::AES_256_CFB1:
        return EVP_aes_256_cfb1();
    case EncryptionMode::AES_256_CFB8:
        return EVP_aes_256_cfb8();
    case EncryptionMode::AES_256_CFB128:
        return EVP_aes_256_cfb128();
    case EncryptionMode::AES_256_CTR:
        return EVP_aes_256_ctr();
    case EncryptionMode::AES_256_OFB:
        return EVP_aes_256_ofb();
    default:
        return nullptr;
    }
}

static uint mode_key_sizes[] = {
        128 /* AES_128_ECB */,  256 /* AES_256_ECB */,    128 /* AES_128_CBC */,
        256 /* AES_256_CBC */,  128 /* AES_128_CFB */,    256 /* AES_256_CFB */,
        128 /* AES_128_CFB1 */, 256 /* AES_256_CFB1 */,   128 /* AES_128_CFB8 */,
        256 /* AES_256_CFB8 */, 128 /* AES_128_CFB128 */, 256 /* AES_256_CFB128 */,
        128 /* AES_128_CTR */,  256 /* AES_256_CTR */,    128 /* AES_128_OFB */,
        256 /* AES_256_OFB */,
};

static void create_key(const unsigned char* origin_key, uint32_t key_length, uint8_t* encrypt_key,
                       EncryptionMode mode) {
    const uint key_size = mode_key_sizes[int(mode)] / 8;
    uint8_t* origin_key_end = ((uint8_t*)origin_key) + key_length; /* origin key boundary*/

    uint8_t* encrypt_key_end; /* encrypt key boundary */
    encrypt_key_end = encrypt_key + key_size;

    std::memset(encrypt_key, 0, key_size); /* initialize key  */

    uint8_t* ptr;        /* Start of the encrypt key*/
    uint8_t* origin_ptr; /* Start of the origin key */
    for (ptr = encrypt_key, origin_ptr = (uint8_t*)origin_key; origin_ptr < origin_key_end;
         ptr++, origin_ptr++) {
        if (ptr == encrypt_key_end) {
            /* loop over origin key until we used all key */
            ptr = encrypt_key;
        }
        *ptr ^= *origin_ptr;
    }
}

static int do_encrypt(EVP_CIPHER_CTX* cipher_ctx, const EVP_CIPHER* cipher,
                      const unsigned char* source, uint32_t source_length,
                      const unsigned char* encrypt_key, const unsigned char* iv, bool padding,
                      unsigned char* encrypt, int* length_ptr) {
    int ret = EVP_EncryptInit(cipher_ctx, cipher, encrypt_key, iv);
    if (ret == 0) {
        return ret;
    }
    ret = EVP_CIPHER_CTX_set_padding(cipher_ctx, padding);
    if (ret == 0) {
        return ret;
    }
    int u_len = 0;

    ret = EVP_EncryptUpdate(cipher_ctx, encrypt, &u_len, source, source_length);
    if (ret == 0) {
        return ret;
    }
    int f_len = 0;
    ret = EVP_EncryptFinal(cipher_ctx, encrypt + u_len, &f_len);
    *length_ptr = u_len + f_len;
    return ret;
}

int EncryptionUtil::encrypt(EncryptionMode mode, const unsigned char* source,
                            uint32_t source_length, const unsigned char* key, uint32_t key_length,
                            const char* iv_str, int iv_input_length, bool padding,
                            unsigned char* encrypt) {
    const EVP_CIPHER* cipher = get_evp_type(mode);
    /* The encrypt key to be used for encryption */
    unsigned char encrypt_key[ENCRYPTION_MAX_KEY_LENGTH / 8];
    create_key(key, key_length, encrypt_key, mode);

    int iv_length = EVP_CIPHER_iv_length(cipher);
    if (cipher == nullptr || (iv_length > 0 && !iv_str)) {
        return AES_BAD_DATA;
    }
    char* init_vec = nullptr;
    std::string iv_default("SELECTDBCLOUD___");

    if (iv_str) {
        init_vec = iv_default.data();
        memcpy(init_vec, iv_str, std::min(iv_input_length, EVP_MAX_IV_LENGTH));
        init_vec[iv_length] = '\0';
    }
    EVP_CIPHER_CTX* cipher_ctx = EVP_CIPHER_CTX_new();
    EVP_CIPHER_CTX_reset(cipher_ctx);
    int length = 0;
    int ret = do_encrypt(cipher_ctx, cipher, source, source_length, encrypt_key,
                         reinterpret_cast<unsigned char*>(init_vec), padding, encrypt, &length);
    EVP_CIPHER_CTX_free(cipher_ctx);
    if (ret == 0) {
        ERR_clear_error();
        return AES_BAD_DATA;
    } else {
        return length;
    }
}

static int do_decrypt(EVP_CIPHER_CTX* cipher_ctx, const EVP_CIPHER* cipher,
                      const unsigned char* encrypt, uint32_t encrypt_length,
                      const unsigned char* encrypt_key, const unsigned char* iv, bool padding,
                      unsigned char* decrypt_content, int* length_ptr) {
    int ret = EVP_DecryptInit(cipher_ctx, cipher, encrypt_key, iv);
    if (ret == 0) {
        return ret;
    }
    ret = EVP_CIPHER_CTX_set_padding(cipher_ctx, padding);
    if (ret == 0) {
        return ret;
    }
    int u_len = 0;
    ret = EVP_DecryptUpdate(cipher_ctx, decrypt_content, &u_len, encrypt, encrypt_length);
    if (ret == 0) {
        return ret;
    }
    int f_len = 0;
    ret = EVP_DecryptFinal_ex(cipher_ctx, decrypt_content + u_len, &f_len);
    *length_ptr = u_len + f_len;
    return ret;
}

int EncryptionUtil::decrypt(EncryptionMode mode, const unsigned char* encrypt,
                            uint32_t encrypt_length, const unsigned char* key, uint32_t key_length,
                            const char* iv_str, int iv_input_length, bool padding,
                            unsigned char* decrypt_content) {
    const EVP_CIPHER* cipher = get_evp_type(mode);

    /* The encrypt key to be used for decryption */
    unsigned char encrypt_key[ENCRYPTION_MAX_KEY_LENGTH / 8];
    create_key(key, key_length, encrypt_key, mode);

    int iv_length = EVP_CIPHER_iv_length(cipher);
    if (cipher == nullptr || (iv_length > 0 && !iv_str)) {
        return AES_BAD_DATA;
    }
    char* init_vec = nullptr;
    std::string iv_default("SELECTDBCLOUD___");

    if (iv_str) {
        init_vec = iv_default.data();
        memcpy(init_vec, iv_str, std::min(iv_input_length, EVP_MAX_IV_LENGTH));
        init_vec[iv_length] = '\0';
    }
    EVP_CIPHER_CTX* cipher_ctx = EVP_CIPHER_CTX_new();
    EVP_CIPHER_CTX_reset(cipher_ctx);
    int length = 0;
    int ret = do_decrypt(cipher_ctx, cipher, encrypt, encrypt_length, encrypt_key,
                         reinterpret_cast<unsigned char*>(init_vec), padding, decrypt_content,
                         &length);
    EVP_CIPHER_CTX_free(cipher_ctx);
    if (ret > 0) {
        return length;
    } else {
        ERR_clear_error();
        return AES_BAD_DATA;
    }
}

// base64 endcode/decode

static char encoding_table[] = {'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M',
                                'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z',
                                'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm',
                                'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z',
                                '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '+', '/'};

static const char base64_pad = '=';

static short decoding_table[256] = {
        -2, -2, -2, -2, -2, -2, -2, -2, -2, -1, -1, -2, -2, -1, -2, -2, -2, -2, -2, -2, -2, -2,
        -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -1, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, 62,
        -2, -2, -2, 63, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, -2, -2, -2, -2, -2, -2, -2, 0,
        1,  2,  3,  4,  5,  6,  7,  8,  9,  10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22,
        23, 24, 25, -2, -2, -2, -2, -2, -2, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38,
        39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, -2, -2, -2, -2, -2, -2, -2, -2, -2,
        -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2,
        -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2,
        -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2,
        -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2,
        -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2,
        -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2};

static int mod_table[] = {0, 2, 1};

size_t base64_encode(const unsigned char* data, size_t length, unsigned char* encoded_data) {
    size_t output_length = (size_t)(4.0 * ceil((double)length / 3.0));

    if (encoded_data == nullptr) {
        return 0;
    }

    for (uint32_t i = 0, j = 0; i < length;) {
        uint32_t octet_a = i < length ? data[i++] : 0;
        uint32_t octet_b = i < length ? data[i++] : 0;
        uint32_t octet_c = i < length ? data[i++] : 0;
        uint32_t triple = (octet_a << 0x10) + (octet_b << 0x08) + octet_c;

        encoded_data[j++] = encoding_table[(triple >> 3 * 6) & 0x3F];
        encoded_data[j++] = encoding_table[(triple >> 2 * 6) & 0x3F];
        encoded_data[j++] = encoding_table[(triple >> 1 * 6) & 0x3F];
        encoded_data[j++] = encoding_table[(triple >> 0 * 6) & 0x3F];
    }

    for (int i = 0; i < mod_table[length % 3]; i++) {
        encoded_data[output_length - 1 - i] = '=';
    }

    return output_length;
}

size_t base64_decode(const char* data, size_t length, char* decoded_data) {
    const char* current = data;
    size_t ch = 0;
    size_t i = 0;
    size_t j = 0;
    size_t k = 0;

    // run through the whole string, converting as we go
    while ((ch = *current++) != '\0' && length-- > 0) {
        if (ch >= 256 || ch < 0) {
            return -1;
        }

        if (ch == base64_pad) {
            if (*current != '=' && (i % 4) == 1) {
                return -1;
            }
            continue;
        }

        ch = decoding_table[ch];
        // a space or some other separator character, we simply skip over
        if (ch == -1) {
            continue;
        } else if (ch == -2) {
            return -1;
        }

        switch (i % 4) {
        case 0:
            decoded_data[j] = ch << 2;
            break;
        case 1:
            decoded_data[j++] |= ch >> 4;
            decoded_data[j] = (ch & 0x0f) << 4;
            break;
        case 2:
            decoded_data[j++] |= ch >> 2;
            decoded_data[j] = (ch & 0x03) << 6;
            break;
        case 3:
            decoded_data[j++] |= ch;
            break;
        default:
            break;
        }

        i++;
    }

    k = j;
    /* mop things up if we ended on a boundary */
    if (ch == base64_pad) {
        switch (i % 4) {
        case 1:
            return 0;
        case 2:
            k++;
        case 3:
            decoded_data[k] = 0;
        default:
            break;
        }
    }

    decoded_data[j] = '\0';

    return j;
}

// encrypt/dencrypt with base64

static std::unordered_map<std::string, EncryptionMode> to_encryption_mode {
        {"AES_128_ECB", EncryptionMode::AES_128_ECB},
        {"AES_256_ECB", EncryptionMode::AES_256_ECB},
        {"AES_128_CBC", EncryptionMode::AES_128_CBC},
        {"AES_256_CBC", EncryptionMode::AES_256_CBC},
        {"AES_128_CFB", EncryptionMode::AES_128_CFB},
        {"AES_256_CFB", EncryptionMode::AES_256_CFB},
        {"AES_128_CFB1", EncryptionMode::AES_128_CFB1},
        {"AES_256_CFB1", EncryptionMode::AES_256_CFB1},
        {"AES_128_CFB8", EncryptionMode::AES_128_CFB8},
        {"AES_256_CFB8", EncryptionMode::AES_256_CFB8},
        {"AES_128_CFB128", EncryptionMode::AES_128_CFB128},
        {"AES_256_CFB128", EncryptionMode::AES_256_CFB128},
        {"AES_128_CTR", EncryptionMode::AES_128_CTR},
        {"AES_256_CTR", EncryptionMode::AES_256_CTR},
        {"AES_128_OFB", EncryptionMode::AES_128_OFB},
        {"AES_256_OFB", EncryptionMode::AES_256_OFB},
};

static inline int encrypt_to_base64_impl(std::string_view source, EncryptionMode mode,
                                         const std::string& key, std::string* encrypt) {
    /*
     * Buffer for ciphertext. Ensure the buffer is long enough for the
     * ciphertext which may be longer than the plaintext, depending on the
     * algorithm and mode.
     */
    int cipher_len = source.length() + 16;
    std::string cipher_text(cipher_len, '0');
    int cipher_text_len = EncryptionUtil::encrypt(
            mode, (unsigned char*)source.data(), source.length(), (unsigned char*)key.c_str(),
            key.length(), nullptr, 0, true, (unsigned char*)cipher_text.data());
    if (cipher_text_len < 0) {
        return -1;
    }

    int encoded_len = (size_t)(4.0 * ceil(cipher_text_len / 3.0));
    std::string encoded_text(encoded_len, '0');
    int encoded_text_len = base64_encode((unsigned char*)cipher_text.data(), cipher_text_len,
                                         (unsigned char*)encoded_text.data());
    if (encoded_text_len < 0) {
        return -1;
    }
    encrypt->assign((char*)encoded_text.data(), encoded_text_len);
    return 0;
}

static int encrypt_to_base64(std::string_view source, const std::string& encrypt_method,
                             const std::string& key, std::string* encrypt) {
    if (source.empty()) {
        *encrypt = "";
        return 0;
    }
    auto iter = to_encryption_mode.find(encrypt_method);
    if (iter == to_encryption_mode.end()) {
        return -1;
    }

    return encrypt_to_base64_impl(source, iter->second, key, encrypt);
}

static inline int decrypt_with_base64_impl(std::string_view encrypt, EncryptionMode mode,
                                           const std::string& key, std::string* source) {
    // base64
    std::unique_ptr<char[]> decoded_text(new char[encrypt.length()]);
    int decoded_text_len = base64_decode(encrypt.data(), encrypt.length(), decoded_text.get());
    if (decoded_text_len < 0) {
        return -1;
    }

    std::unique_ptr<char[]> plain_text(new char[decoded_text_len]);
    int plain_text_len = EncryptionUtil::decrypt(
            mode, (unsigned char*)decoded_text.get(), decoded_text_len, (unsigned char*)key.c_str(),
            key.length(), nullptr, 0, true, (unsigned char*)plain_text.get());
    if (plain_text_len < 0) {
        return -1;
    }
    source->assign(plain_text.get(), plain_text_len);
    return 0;
}

static int decrypt_with_base64(std::string_view encrypt, const std::string& encrypt_method,
                               const std::string& key, std::string* source) {
    if (encrypt.empty()) {
        *source = "";
        return 0;
    }
    auto iter = to_encryption_mode.find(encrypt_method);
    if (iter == to_encryption_mode.end()) {
        return -1;
    }
    return decrypt_with_base64_impl(encrypt, iter->second, key, source);
}

int encrypt_ak_sk(AkSkRef plain_ak_sk, const std::string& encryption_method,
                  const std::string& encryption_key, AkSkPair* cipher_ak_sk) {
    std::string encrypt_ak;
    std::string encrypt_sk;
    if (encrypt_to_base64(plain_ak_sk.second, encryption_method, encryption_key, &encrypt_sk) !=
        0) {
        *cipher_ak_sk = {"", ""};
        return -1;
    }
    *cipher_ak_sk = {std::string(plain_ak_sk.first), std::move(encrypt_sk)};
    return 0;
}

int decrypt_ak_sk(AkSkRef cipher_ak_sk, const std::string& encryption_method,
                  const std::string& encryption_key, AkSkPair* plain_ak_sk) {
    std::string ak;
    std::string sk;
    if (decrypt_with_base64(cipher_ak_sk.second, encryption_method, encryption_key, &sk) != 0) {
        *plain_ak_sk = {"", ""};
        return -1;
    }
    *plain_ak_sk = {std::string(cipher_ak_sk.first), std::move(sk)};
    return 0;
}

int decrypt_ak_sk_helper(std::string_view cipher_ak, std::string_view cipher_sk,
                         const EncryptionInfoPB& encryption_info, AkSkPair* plain_ak_sk_pair) {
    std::string key;
    int ret = get_encryption_key_for_ak_sk(encryption_info.key_id(), &key);
    { TEST_SYNC_POINT_CALLBACK("decrypt_ak_sk:get_encryption_key", &key, &ret); }
    if (ret != 0) {
        LOG(WARNING) << "failed to get encryption key version_id: " << encryption_info.key_id();
        return -1;
    }
    ret = decrypt_ak_sk({cipher_ak, cipher_sk}, encryption_info.encryption_method(), key,
                        plain_ak_sk_pair);
    if (ret != 0) {
        LOG(WARNING) << "failed to decrypt";
        return -1;
    }
    return 0;
}

/**
 * @brief Generates a random root key. If a root key already exists, returns immediately. 
 * 
 * @param txn_kv 
 * @param kms_client 
 * @param plaintext store the plaintext of the root key
 * @param encoded_ciphertext store the base64-encoded ciphertext of the root key.
 * @return int 0 for success to generate, 1 for not need to generate, -1 for failure.
 */
static int generate_random_root_key(TxnKv* txn_kv, KmsClient* kms_client, std::string* plaintext,
                                    std::string* encoded_ciphertext) {
    /**
    * 1. If KMS is enabled, use KMS to generate a new key.
    * 2. If KMS is not enabled, try using the encryption_key from the configuration, which must be in Base64 format.
    * 3. If no key is found in the configuration, generate a random key in memory.
    */
    std::string key = system_meta_service_encryption_key_info_key();
    std::string val;
    std::unique_ptr<Transaction> txn;
    TxnErrorCode err = txn_kv->create_txn(&txn);
    if (err != TxnErrorCode::TXN_OK) {
        LOG_WARNING("failed to create txn").tag("err", err);
        return -1;
    }
    err = txn->get(key, &val);
    if (err != TxnErrorCode::TXN_OK && err != TxnErrorCode::TXN_KEY_NOT_FOUND) {
        LOG_WARNING("failed to get key of encryption_key_info").tag("err", err);
        return -1;
    }

    if (err == TxnErrorCode::TXN_OK) {
        if (config::enable_kms && config::focus_add_kms_data_key) {
            EncryptionKeyInfoPB key_info;
            if (!key_info.ParseFromString(val)) {
                LOG_WARNING("failed to parse encryption_root_key");
                return -1;
            }
            for (const auto& item : key_info.items()) {
                if (item.has_kms_info()) {
                    return 1;
                }
            }
            LOG(INFO) << "focus to create kms data key";
        } else {
            LOG(INFO) << "not need to generate root key";
            return 1;
        }
    }

    // 1. use kms to generate a new key
    if (config::enable_kms) {
        if (kms_client == nullptr) {
            LOG_WARNING("no kms client");
            return -1;
        }
        std::string encoded_root_key_ciphertext;
        std::string encoded_root_key_plaintext;
        if (kms_client->generate_data_key(&encoded_root_key_ciphertext,
                                          &encoded_root_key_plaintext) != 0) {
            LOG_WARNING("failed to generate data key");
            return -1;
        }
        if (encoded_root_key_ciphertext.empty() || encoded_root_key_plaintext.empty()) {
            LOG_WARNING("empty data key generated");
            return -1;
        }

        // decode plaintext
        std::string root_key_plaintext(encoded_root_key_plaintext.length(), '0');
        int decoded_len =
                base64_decode(encoded_root_key_plaintext.c_str(),
                              encoded_root_key_plaintext.length(), root_key_plaintext.data());
        if (decoded_len < 0) {
            LOG_WARNING("failed to decode plaintext of kms");
            return -1;
        }
        root_key_plaintext.assign(root_key_plaintext.data(), decoded_len);

        *plaintext = std::move(root_key_plaintext);
        *encoded_ciphertext = std::move(encoded_root_key_ciphertext);
        return 0;
    }

    // 2. try using the encryption_key from the configuration
    if (!cloud::config::encryption_key.empty()) {
        std::string decoded_string(cloud::config::encryption_key.length(), '0');
        int decoded_text_len =
                base64_decode(cloud::config::encryption_key.c_str(),
                              cloud::config::encryption_key.length(), decoded_string.data());
        if (decoded_text_len < 0) {
            LOG_WARNING("fail to decode encryption_key in config");
            return -1;
        }
        decoded_string.assign(decoded_string.data(), decoded_text_len);
        *plaintext = std::move(decoded_string);
        *encoded_ciphertext = cloud::config::encryption_key;
        return 0;
    }

    // 3. otherwise, generate a random data key in memory
    std::mt19937 rnd(time(nullptr));
    std::uniform_int_distribution<short> dist(std::numeric_limits<char>::min(),
                                              std::numeric_limits<char>::max());
    std::string root_key_plaintext(32, '0');
    for (char& i : root_key_plaintext) {
        i = (char)dist(rnd);
    }

    // encode in base64
    int key_len = root_key_plaintext.length();
    int encoded_len = (size_t)(4.0 * ceil(key_len / 3.0));
    std::string encoded_text(encoded_len, '0');
    int encoded_text_len = base64_encode((unsigned char*)root_key_plaintext.data(), key_len,
                                         (unsigned char*)encoded_text.data());
    if (encoded_text_len < 0) {
        LOG_WARNING("failed to encode encryption_key");
        return -1;
    }
    std::string encoded_root_key_ciphertext;
    encoded_root_key_ciphertext.assign(encoded_text.data(), encoded_text_len);

    *plaintext = std::move(root_key_plaintext);
    *encoded_ciphertext = std::move(encoded_root_key_ciphertext);
    return 0;
}

// Todo: Does not need to be locked now, only generated when the process is initialized
std::map<int64_t, std::string> global_encryption_key_info_map; // key_id->encryption_key

static int get_current_root_keys(TxnKv* txn_kv, std::map<int64_t, std::string>* keys) {
    std::unique_ptr<KmsClient> kms_client;
    if (config::enable_kms) {
        if (config::kms_info_encryption_key.empty() || config::kms_info_encryption_method.empty() ||
            config::kms_ak.empty() || config::kms_sk.empty()) {
            LOG_WARNING("incorrect kms conf")
                    .tag("encryption_key", config::kms_info_encryption_key)
                    .tag("encryption_method", config::kms_info_encryption_method)
                    .tag("ak", config::kms_ak)
                    .tag("sk", config::kms_sk);
            return -1;
        }
        std::string decoded_encryption_key(config::kms_info_encryption_key.length(), '0');
        int decoded_key_len = cloud::base64_decode(config::kms_info_encryption_key.c_str(),
                                                   config::kms_info_encryption_key.length(),
                                                   decoded_encryption_key.data());
        decoded_encryption_key.assign(decoded_encryption_key.data(), decoded_key_len);
        AkSkPair out;
        if (decrypt_ak_sk({config::kms_ak, config::kms_sk}, config::kms_info_encryption_method,
                          decoded_encryption_key, &out) != 0) {
            LOG_WARNING("failed to decrypt kms info");
            return -1;
        }

        KmsConf conf {out.first,          out.second,      config::kms_endpoint,
                      config::kms_region, config::kms_cmk, config::kms_provider};

        auto ret = create_kms_client(std::move(conf), &kms_client);
        if (ret != 0) {
            LOG_WARNING("failed to create kms client").tag("ret", ret);
            return -1;
        }
        ret = kms_client->init();
        if (ret != 0) {
            LOG_WARNING("failed to init kms client").tag("ret", ret);
            return -1;
        }
    }

    // To avoid transaction timeouts, it is necessary to first generate a root key
    std::string root_key_plaintext;
    std::string encoded_root_key_ciphertext;
    int ret = generate_random_root_key(txn_kv, kms_client.get(), &root_key_plaintext,
                                       &encoded_root_key_ciphertext);
    if (ret == -1) {
        LOG_WARNING("failed to generate random root key");
        return -1;
    }

    while (true) {
        std::string key = system_meta_service_encryption_key_info_key();
        std::string val;
        std::unique_ptr<Transaction> txn;
        TxnErrorCode err = txn_kv->create_txn(&txn);
        if (err != TxnErrorCode::TXN_OK) {
            LOG_WARNING("failed to create txn").tag("ret", ret);
            return -1;
        }
        err = txn->get(key, &val);
        if (ret != 0 && ret != 1) {
            LOG_WARNING("failed to get key of encryption_key_info").tag("ret", ret);
            return -1;
        }

        bool need_to_focus_add_kms_data_key = true;
        EncryptionKeyInfoPB key_info;
        if (err == TxnErrorCode::TXN_OK) {
            if (!key_info.ParseFromString(val)) {
                LOG_WARNING("failed to parse encryption_root_key");
                return -1;
            }

            LOG_INFO("get server encryption_root_key").tag("key_info", proto_to_json(key_info));

            for (const auto& item : key_info.items()) {
                std::string encoded_root_key_plaintext;
                if (item.has_kms_info()) {
                    need_to_focus_add_kms_data_key = false;
                    // use kms to decrypt
                    if (kms_client == nullptr) {
                        LOG_WARNING("no kms client");
                        return -1;
                    }
                    if (item.kms_info().endpoint() != kms_client->conf().endpoint ||
                        item.kms_info().region() != kms_client->conf().region) {
                        LOG_WARNING("kms info is not match")
                                .tag("kms endpoint", kms_client->conf().endpoint)
                                .tag("kms region", kms_client->conf().region)
                                .tag("saved endpoint", item.kms_info().endpoint())
                                .tag("saved region", item.kms_info().region());
                        return -1;
                    }

                    auto ret = kms_client->decrypt(item.key(), &encoded_root_key_plaintext);
                    if (ret != 0) {
                        LOG_WARNING("failed to decrypt encryption_root_key");
                        return -1;
                    }
                } else {
                    encoded_root_key_plaintext = item.key(); // Todo: do not copy
                }

                std::string root_key_plaintext(encoded_root_key_plaintext.length(), '0');
                int decoded_text_len = base64_decode(encoded_root_key_plaintext.c_str(),
                                                     encoded_root_key_plaintext.length(),
                                                     root_key_plaintext.data());
                if (decoded_text_len < 0) {
                    LOG_WARNING("failed to decode encryption_root_key");
                    return -1;
                }
                root_key_plaintext.assign(root_key_plaintext.data(), decoded_text_len);
                keys->insert({item.key_id(), std::move(root_key_plaintext)});
            }
            if (config::enable_kms && config::focus_add_kms_data_key &&
                need_to_focus_add_kms_data_key) {
                // Todo: need to restart other ms to update global_encryption_key_info_map now
                LOG(INFO) << "focus to add kms data key";
            } else {
                return 0;
            }
        }

        // encryption_root_key not found, need to save a new root key into fdb
        if (root_key_plaintext.empty() || encoded_root_key_ciphertext.empty()) {
            LOG_WARNING("empty new root key");
            return -1;
        }

        int32_t new_key_id = key_info.items().size() + 1;
        auto* item = key_info.add_items();
        item->set_key_id(new_key_id);
        item->set_key(encoded_root_key_ciphertext);
        if (config::enable_kms) {
            item->mutable_kms_info()->set_endpoint(config::kms_endpoint);
            item->mutable_kms_info()->set_region(config::kms_region);
            item->mutable_kms_info()->set_cmk(config::kms_cmk);
        }

        val = key_info.SerializeAsString();
        if (val.empty()) {
            LOG_WARNING("failed to serialize");
            return -1;
        }
        txn->put(key, val);
        LOG_INFO("put server encryption_key")
                .tag("encryption_key", encoded_root_key_ciphertext)
                .tag("key_id", new_key_id);
        err = txn->commit();
        if (err == TxnErrorCode::TXN_CONFLICT) {
            LOG_WARNING("commit encryption_key is conflicted, retry it later");
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
            continue;
        } else if (err != TxnErrorCode::TXN_OK) {
            LOG_WARNING("failed to commit encryption_key");
            return -1;
        }
        keys->insert({new_key_id, std::move(root_key_plaintext)});
        return 0;
    }
    return 0;
}

int init_global_encryption_key_info_map(TxnKv* txn_kv) {
    if (get_current_root_keys(txn_kv, &global_encryption_key_info_map) != 0) {
        return -1;
    }
    DCHECK(!global_encryption_key_info_map.empty());
    return 0;
}

} // namespace doris::cloud