teems/Enclave/config.cpp

#include "Enclave_t.h"
#include "comms.hpp"
#include "config.hpp"
#include "utils.hpp"
#include "route.hpp"
#include "ingest.hpp"

#define CEILDIV(x,y) (((x)+(y)-1)/(y))

Config g_teems_config;

int generateMasterKeys(sgx_aes_gcm_128bit_key_t master_secret,
    sgx_aes_gcm_128bit_key_t &ESK, sgx_aes_gcm_128bit_key_t &TSK)
{
    unsigned char zeroes[SGX_AESGCM_KEY_SIZE];
    unsigned char iv[SGX_AESGCM_IV_SIZE];
    sgx_aes_gcm_128bit_tag_t mac;
    memset(zeroes, 0, SGX_AESGCM_KEY_SIZE);
    memset(iv, 0, SGX_AESGCM_IV_SIZE);
    memcpy(iv, "Encryption", sizeof("Encryption"));
    sgx_status_t ret = SGX_SUCCESS;

    ret = sgx_rijndael128GCM_encrypt((const sgx_aes_gcm_128bit_key_t *)
        (master_secret), zeroes, SGX_AESGCM_KEY_SIZE,
        (uint8_t*) ESK, iv, SGX_AESGCM_IV_SIZE, NULL, 0, &mac);
    if(ret!=SGX_SUCCESS) {
        return -1;
    }

    printf("Encryption Master Key: ");
    for(int i=0;i<SGX_AESGCM_KEY_SIZE;i++) {
        printf("%x", ESK[i]);
    }
    printf("\n\n");

    memset(iv, 0, SGX_AESGCM_IV_SIZE);
    memcpy(iv, "Token", sizeof("Token"));
    ret = sgx_rijndael128GCM_encrypt((const sgx_aes_gcm_128bit_key_t *)
        (master_secret), zeroes, SGX_AESGCM_KEY_SIZE,
        (uint8_t*) TSK, iv, SGX_AESGCM_IV_SIZE, NULL, 0, &mac);
    if(ret!=SGX_SUCCESS) {
        return -1;
    }

    printf("Token Master Key: ");
    for(int i=0;i<SGX_AESGCM_KEY_SIZE;i++) {
        printf("%x", TSK[i]);
    }
    printf("\n");

    return 1;
}


bool ecall_config_load(threadid_t nthreads, bool private_routing,
    EnclaveAPIParams *apiparams,
    EnclaveAPINodeConfig *apinodeconfigs,
    nodenum_t num_nodes, nodenum_t my_node_num)
{
    g_teems_config.nthreads = nthreads;
    g_teems_config.num_nodes = num_nodes;
    g_teems_config.num_ingestion_nodes = 0;
    g_teems_config.num_routing_nodes = 0;
    g_teems_config.num_storage_nodes = 0;
    g_teems_config.my_node_num = my_node_num;
    g_teems_config.user_count = apiparams->user_count;
    g_teems_config.msg_size = apiparams->msg_size;
    g_teems_config.m_priv_out = apiparams->m_priv_out;
    g_teems_config.m_priv_in = apiparams->m_priv_in;
    g_teems_config.m_pub_out = apiparams->m_pub_out;
    g_teems_config.m_pub_in = apiparams->m_pub_in;
    g_teems_config.private_routing = private_routing;
    memcpy(g_teems_config.master_secret, apiparams->master_secret, SGX_AESGCM_KEY_SIZE);

    // Temporary vectors to store node numbers for nodes of different
    // types, where the node numbers are smaller than our own node
    // number
    std::vector<nodenum_t> ing_smaller, rte_smaller, str_smaller;
    uint16_t cumul_weight = 0;
    g_teems_config.weights.clear();
    g_teems_config.ingestion_nodes.clear();
    g_teems_config.routing_nodes.clear();
    g_teems_config.storage_nodes.clear();
    for (nodenum_t i=0; i<num_nodes; ++i) {
        NodeWeight nw;
        nw.startweight = cumul_weight;
        nw.weight = apinodeconfigs[i].weight;
        if (apinodeconfigs[i].roles & ROLE_INGESTION) {
            g_teems_config.num_ingestion_nodes += 1;
            if (i < my_node_num) {
                ing_smaller.push_back(i);
            } else {
                g_teems_config.ingestion_nodes.push_back(i);
            }
        }
        if (apinodeconfigs[i].roles & ROLE_ROUTING) {
            g_teems_config.num_routing_nodes += 1;
            if (i < my_node_num) {
                rte_smaller.push_back(i);
            } else {
                g_teems_config.routing_nodes.push_back(i);
            }
        }
        if (apinodeconfigs[i].roles & ROLE_STORAGE) {
            g_teems_config.num_storage_nodes += 1;
            if (i < my_node_num) {
                str_smaller.push_back(i);
            } else {
                g_teems_config.storage_nodes.push_back(i);
            }
        }
        cumul_weight += nw.weight;
        g_teems_config.weights.push_back(nw);
        if (i == my_node_num) {
            g_teems_config.my_weight = nw.weight;
        }
    }
    g_teems_config.tot_weight = cumul_weight;
    // Concatenate the *_smaller vectors to the ends of the
    // g_teems_config.*_nodes vectors.  This way, each node has a list
    // of nodes of each role starting with itself and "looping around".
    // This should make the communication pattern have less of a
    // bottleneck.
    g_teems_config.ingestion_nodes.insert(
        g_teems_config.ingestion_nodes.end(),
        ing_smaller.begin(),
        ing_smaller.end());
    g_teems_config.routing_nodes.insert(
        g_teems_config.routing_nodes.end(),
        rte_smaller.begin(),
        rte_smaller.end());
    g_teems_config.storage_nodes.insert(
        g_teems_config.storage_nodes.end(),
        str_smaller.begin(),
        str_smaller.end());

    // Initialize the threadpool and the pseudorandom bytes pools
    threadpool_init(nthreads);
    PRB_pool_init(nthreads);

    if(apinodeconfigs[my_node_num].roles & ROLE_INGESTION) {
        sgx_aes_gcm_128bit_key_t ESK, TSK;
        generateMasterKeys(g_teems_config.master_secret, ESK, TSK);

        uint32_t num_clients_total = g_teems_config.user_count;
        uint32_t num_ing_nodes = g_teems_config.num_ingestion_nodes;
        uint32_t clients_per_server = CEILDIV(num_clients_total, num_ing_nodes);
        uint32_t ing_with_additional = num_clients_total % num_ing_nodes;
        uint32_t num_smaller_ing = (uint32_t) ing_smaller.size();
        uint32_t num_clients_this_ing = clients_per_server;
        num_clients_this_ing += (num_smaller_ing < ing_with_additional)? 1: 0;
        uint32_t client_start = num_smaller_ing * clients_per_server;
        if (ing_with_additional > 0) {
            if(ing_with_additional > num_smaller_ing) {
                client_start+= num_smaller_ing;
            } else {
                client_start+= ing_with_additional;
            }
        }
        g_ing.initialize(num_clients_this_ing, client_start, ESK);

        /*
        // Check that the keys generated in Enclave match the ones generated
        // in the client application
        for(uint32_t i=0; i<g_ing.getClientNum(); i++) {
            printf("Client Key %d: ", i + g_ing.getClientStart());
            sgx_aes_gcm_128bit_key_t key;
            unsigned char* ckey = (g_ing.getClientKeys())[i];
            memcpy(key, ckey, SGX_AESGCM_KEY_SIZE);
            for(int j = 0; j<SGX_AESGCM_KEY_SIZE; j++) {
                printf("%x", key[j]);
            }
            printf("\n");
        }
        */

    }

    if (!route_init()) {
        return false;
    }
    return comms_init_nodestate(apinodeconfigs, num_nodes, my_node_num);
}

void ecall_close()
{
    PRB_pool_shutdown();
    threadpool_shutdown();
}