teems/Enclave/storage.cpp

#include "utils.hpp"
#include "config.hpp"
#include "ORExpand.hpp"
#include "sort.hpp"
#include "storage.hpp"
#include "client.hpp"

#define PROFILE_STORAGE

StgClient *clients;
uint8_t *epoch_tokens;
uint8_t *epoch_msgbundles;

static struct {
    uint32_t max_users;
    uint32_t my_storage_node_id;
    // A local storage buffer, used when we need to do non-in-place
    // sorts of the messages that have arrived
    MsgBuffer stg_buf;
    // The destination vector for ORExpand
    std::vector<uint32_t> dest;
} storage_state;

bool storage_generateClientKeys(uint32_t num_clients, uint32_t my_stg_no) {
    uint16_t num_priv_channels = g_teems_config.m_priv_in;
    uint16_t msg_size = g_teems_config.msg_size;
    uint32_t pt_msgbundle_size = num_priv_channels * msg_size;

    clients = new StgClient[num_clients];

    for(uint32_t i =0; i < num_clients; i++) {
        uint32_t mid = storage_state.my_storage_node_id + i;
        clients[i].my_id = mid;
        clients[i].priv_friends = new clientid_t[g_teems_config.m_priv_in];
        // Initialize this client's private channel friends as themself
        for(int j =0; j <g_teems_config.m_priv_in; j++) {
            (clients[i].priv_friends)[j] = mid;
        }
    }

    uint32_t num_stg_nodes = g_teems_config.num_storage_nodes;
    uint32_t c_simid = my_stg_no;

    //printf("In Ingestion::genCK, num_clients = %d\n", num_clients);
    for (uint32_t i=0; i<num_clients; i++) {
        const sgx_aes_gcm_128bit_key_t *pESK = &(g_teems_config.ESK);
        unsigned char zeroes[SGX_AESGCM_KEY_SIZE];
        unsigned char iv[SGX_AESGCM_IV_SIZE];
        sgx_aes_gcm_128bit_tag_t tag;
        memset(zeroes, 0, SGX_AESGCM_KEY_SIZE);
        memset(iv, 0, SGX_AESGCM_IV_SIZE);

        memcpy(iv, (uint8_t*) (&c_simid), sizeof(c_simid));
        memcpy(iv + sizeof(c_simid), "STG", sizeof("STG"));

        sgx_status_t ret = SGX_SUCCESS;
        ret = sgx_rijndael128GCM_encrypt(pESK, zeroes, SGX_AESGCM_KEY_SIZE,
            (uint8_t*) (clients[i].key), iv, SGX_AESGCM_IV_SIZE, NULL, 0, &tag);
        if(ret!=SGX_SUCCESS) {
            printf("stg_generateClientKeys FAIL\n");
            return false;
        }

        /*
        if(c_simid % 10 == 0) {
            printf("Storage: c_simid = %d, Key:", c_simid);
            for (int k = 0; k<SGX_AESGCM_KEY_SIZE; k++) {
                printf("%x", (clients[i].key)[k]);
            }
            printf("\n");
        }
        */
        c_simid+=num_stg_nodes;

    }

    return true;
}

// route_init will call this function; no one else should call it
// explicitly.  The parameter is the number of messages that can fit in
// the storage-side MsgBuffer.  Returns true on success, false on
// failure.
bool storage_init(uint32_t max_users, uint32_t msg_buf_size)
{
    storage_state.max_users = max_users;
    storage_state.stg_buf.alloc(msg_buf_size);
    storage_state.dest.resize(msg_buf_size);
    uint32_t my_storage_node_id = 0;
    uint32_t my_stg_pos = 0;
    for (nodenum_t i=0; i<g_teems_config.num_nodes; ++i) {
        if (g_teems_config.roles[i] & ROLE_STORAGE) {
            if (i == g_teems_config.my_node_num) {
                storage_state.my_storage_node_id = my_storage_node_id << DEST_UID_BITS;
                my_stg_pos = my_storage_node_id;
            } else {
                ++my_storage_node_id;
            }
        }
    }

    printf("my_stg_pos = %d\n", my_stg_pos);
    storage_generateClientKeys(max_users, my_stg_pos);
    // sendClientTokens();

    return true;
}

// Handle the messages received by a storage node.  Pass a _locked_
// MsgBuffer.  This function will itself reset and unlock it when it's
// done with it.
void storage_received(MsgBuffer &storage_buf)
{
    uint16_t msg_size = g_teems_config.msg_size;
    nodenum_t my_node_num = g_teems_config.my_node_num;
    const uint8_t *msgs = storage_buf.buf;
    uint32_t num_msgs = storage_buf.inserted;
    uint32_t real = 0, padding = 0;
    uint32_t uid_mask = (1 << DEST_UID_BITS) - 1;
    uint32_t nid_mask = ~uid_mask;

#ifdef PROFILE_STORAGE
    unsigned long start_received = printf_with_rtclock("begin storage_received (%u)\n", storage_buf.inserted);
#endif

    // It's OK to test for errors in a way that's non-oblivous if
    // there's an error (but it should be oblivious if there are no
    // errors)
    for (uint32_t i=0; i<num_msgs; ++i) {
        uint32_t uid = *(const uint32_t*)(storage_buf.buf+(i*msg_size));
        bool ok = ((((uid & nid_mask) == storage_state.my_storage_node_id)
            & ((uid & uid_mask) < storage_state.max_users))
            | ((uid & uid_mask) == uid_mask));
        if (!ok) {
            printf("Received bad uid: %08x\n", uid);
            assert(ok);
        }
    }

    // Testing: report how many real and dummy messages arrived
    printf("Storage server received %u messages:\n", num_msgs);
    for (uint32_t i=0; i<num_msgs; ++i) {
        uint32_t dest_addr = *(const uint32_t*)msgs;
        nodenum_t dest_node =
            g_teems_config.storage_map[dest_addr >> DEST_UID_BITS];
        if (dest_node != my_node_num) {
            char hexbuf[2*msg_size + 1];
            for (uint32_t j=0;j<msg_size;++j) {
                snprintf(hexbuf+2*j, 3, "%02x", msgs[j]);
            }
            printf("Misrouted message: %s\n", hexbuf);
        } else if ((dest_addr & uid_mask) == uid_mask) {
            ++padding;
        } else {
            ++real;
        }
        msgs += msg_size;
    }
    printf("%u real, %u padding\n", real, padding);

    /*
    for (uint32_t i=0;i<num_msgs; ++i) {
        printf("%3d: %08x %08x\n", i,
        *(uint32_t*)(storage_buf.buf+(i*msg_size)),
        *(uint32_t*)(storage_buf.buf+(i*msg_size+4)));
    }
    */
    // Sort the received messages by userid into the
    // storage_state.stg_buf MsgBuffer.
#ifdef PROFILE_STORAGE
    unsigned long start_sort = printf_with_rtclock("begin oblivious sort (%u)\n", storage_buf.inserted);
#endif
    sort_mtobliv<UidKey>(g_teems_config.nthreads, storage_buf.buf,
        msg_size, storage_buf.inserted, storage_buf.bufsize,
        storage_state.stg_buf.buf);
#ifdef PROFILE_STORAGE
    printf_with_rtclock_diff(start_sort, "end oblivious sort (%u)\n", storage_buf.inserted);
#endif

    /*
    for (uint32_t i=0;i<num_msgs; ++i) {
        printf("%3d: %08x %08x\n", i,
        *(uint32_t*)(storage_state.stg_buf.buf+(i*msg_size)),
        *(uint32_t*)(storage_state.stg_buf.buf+(i*msg_size+4)));
    }
    */

#ifdef PROFILE_STORAGE
    unsigned long start_dest = printf_with_rtclock("begin setting dests (%u)\n", storage_state.stg_buf.bufsize);
#endif
    // Obliviously set the dest array
    uint32_t *dests = storage_state.dest.data();
    uint32_t stg_size = storage_state.stg_buf.bufsize;
    const uint8_t *buf = storage_state.stg_buf.buf;
    uint32_t m_priv_in = g_teems_config.m_priv_in;

    uint32_t uid = *(uint32_t*)(buf);
    uid &= uid_mask;
    // num_msgs is not a private value
    if (num_msgs > 0) {
        dests[0] = oselect_uint32_t(uid * m_priv_in, 0xffffffff,
            uid == uid_mask);
    }
    uint32_t prev_uid = uid;
    for (uint32_t i=1; i<num_msgs; ++i) {
        uid = *(uint32_t*)(buf + i*msg_size);
        uid &= uid_mask;
        uint32_t next = oselect_uint32_t(uid * m_priv_in, dests[i-1]+1,
            uid == prev_uid);
        dests[i] = oselect_uint32_t(next, 0xffffffff, uid == uid_mask);
        prev_uid = uid;
    }
    for (uint32_t i=num_msgs; i<stg_size; ++i) {
        dests[i] = 0xffffffff;
        *(uint32_t*)(buf + i*msg_size) = 0xffffffff;
    }
#ifdef PROFILE_STORAGE
    printf_with_rtclock_diff(start_dest, "end setting dests (%u)\n", stg_size);
#endif
    /*
    for (uint32_t i=0;i<stg_size; ++i) {
        printf("%3d: %08x %08x %u\n", i,
        *(uint32_t*)(storage_state.stg_buf.buf+(i*msg_size)),
        *(uint32_t*)(storage_state.stg_buf.buf+(i*msg_size+4)),
        dests[i]);
    }
    */

#ifdef PROFILE_STORAGE
    unsigned long start_expand = printf_with_rtclock("begin ORExpand (%u)\n", stg_size);
#endif
    ORExpand_parallel<OSWAP_16X>(storage_state.stg_buf.buf, dests,
        msg_size, stg_size, g_teems_config.nthreads);
#ifdef PROFILE_STORAGE
    printf_with_rtclock_diff(start_expand, "end ORExpand (%u)\n", stg_size);
#endif
    /*
    for (uint32_t i=0;i<stg_size; ++i) {
        printf("%3d: %08x %08x %u\n", i,
        *(uint32_t*)(storage_state.stg_buf.buf+(i*msg_size)),
        *(uint32_t*)(storage_state.stg_buf.buf+(i*msg_size+4)),
        dests[i]);
    }
    */

    // You can do more processing after these lines, as long as they
    // don't touch storage_buf.  They _can_ touch the backing buffer
    // storage_state.stg_buf.
    storage_buf.reset();
    pthread_mutex_unlock(&storage_buf.mutex);

    // TODO: Send the storage_state.stg_buf to clients
    // Encrypt the storage_state.stg_buf into MSG_BUNDLES
    // encryptMsgBundles(num_clients, num_priv_channels, msg_size);
    // Uses: storage_state.stg_buf, clients, epoch_msgbundles

    // generateTokens();
    // Uses: clients, epoch_tokens,

    // ocall_process_msgbundles(epoch_msgbundles, epoch_tokens);
    // sendClientTokens();
    //

    // processStorage(msg_bundles) :
    //    a) Send out msg_bundles to clients with open sockets
    //    b) Store the other msg_bundles in "backend"

    storage_state.stg_buf.reset();

#ifdef PROFILE_STORAGE
    printf_with_rtclock_diff(start_received, "end storage_received (%u)\n", storage_buf.inserted);
#endif
}


/*
    Given a local client identifier, generate the tokens for this client
    for their priv_friends for the next round.
    Populates the supplied tokens buffer with the correct tokens.
*/
bool generate_tokens(clientid_t lcid)
{
    uint32_t pt_tokens_size = (g_teems_config.m_priv_in * SGX_AESGCM_KEY_SIZE);
    uint32_t enc_tokens_size = (g_teems_config.m_priv_in * SGX_AESGCM_KEY_SIZE) +
        SGX_AESGCM_IV_SIZE + SGX_AESGCM_MAC_SIZE;

    const sgx_aes_gcm_128bit_key_t *pTSK = &(g_teems_config.TSK);
    unsigned char *tkn_iv_ptr = epoch_tokens + enc_tokens_size * lcid;
    unsigned char *tkn_ptr = tkn_iv_ptr + SGX_AESGCM_IV_SIZE;
    unsigned char *tkn_tag = tkn_ptr + pt_tokens_size;

    // We construct the plaintext underlying the token in the correct location
    // for this client in epoch_tokens
    // The tokens get stored in token_body
    // Later we encrypt token_body with the client's storage key and overwrite
    // the correct location in epoch_tokens
    unsigned char token_body[pt_tokens_size];

    memset(token_body, 0, pt_tokens_size);
    memset(tkn_iv_ptr, 0, SGX_AESGCM_IV_SIZE);
    // IV = client_id | epoch_no
    memcpy(tkn_iv_ptr, (uint8_t*) (&(clients[lcid].my_id)), sizeof(clientid_t));
    // TODO: Add epoch to IV
    //memcpy(tkn_iv_ptr + sizeof(clientid_t), epoch_no, sizeof(epoch_no));

    unsigned char *ptr = tkn_ptr;
    for(int i = 0; i<g_teems_config.m_priv_in; i++)
    {
        memcpy(ptr, (uint8_t*) (&(clients[lcid].my_id)), sizeof(clientid_t));
        memcpy(ptr + sizeof(clientid_t), (uint8_t*) (&(clients[lcid].priv_friends[i])), sizeof(clientid_t));
        ptr+=SGX_AESGCM_KEY_SIZE;
    }

    sgx_status_t ret = SGX_SUCCESS;
    ret = sgx_rijndael128GCM_encrypt(pTSK, tkn_ptr, pt_tokens_size,
        (uint8_t*) token_body, tkn_iv_ptr, SGX_AESGCM_IV_SIZE, NULL, 0,
        (sgx_aes_gcm_128bit_tag_t*) tkn_tag);
    if(ret!=SGX_SUCCESS) {
        printf("generate_tokens: Creating token FAIL\n");
        return false;
    }

    /*
    if(lcid == 0) {
        printf("Checking generated token_body:");
        for(uint32_t i = 0; i < pt_tokens_size; i++) {
            printf("%x", token_body[i]);
        }
        printf("\n");
    }
    */

    unsigned char *cl_iv = clients[lcid].iv;
    ret = (sgx_rijndael128GCM_encrypt(&(clients[lcid].key), token_body, pt_tokens_size,
        (uint8_t*) tkn_ptr, cl_iv, SGX_AESGCM_IV_SIZE, NULL, 0,
        (sgx_aes_gcm_128bit_tag_t*) tkn_tag));
    if(ret!=SGX_SUCCESS) {
        printf("generate_tokens: Encrypting token FAIL\n");
        return false;
    }

    memcpy(tkn_iv_ptr, cl_iv, SGX_AESGCM_IV_SIZE);
    // Update IV
    uint64_t *iv_ctr = (uint64_t*) cl_iv;
    (*iv_ctr)+=1;

    /*
    if(lcid == 0) {
        printf("Encrypted client token:");
        for(uint32_t i = 0; i < pt_tokens_size; i++) {
            printf("%x", tkn_ptr[i]);
        }
        printf("\n");
    }
    */

    return true;
}


bool ecall_storage_authenticate(clientid_t cid, unsigned char *auth_message)
{
    bool ret = false;
    uint32_t lcid = cid / g_teems_config.num_storage_nodes;
    const sgx_aes_gcm_128bit_key_t *ckey = &(clients[lcid].key);

    ret = authenticateClient(auth_message, ckey);

    if(!ret) {
        printf("Storage authentication FAIL\n");
    }

    // When clients authenticate:
    // 1) Send back mailbox archive
    // 2) Send tokens for current epoch

    ret &= generate_tokens(lcid);
    /*
    uint32_t enc_token_bundle_size = SGX_AESGCM_KEY_SIZE + SGX_AESGCM_IV_SIZE
    unsigned char *token_ptr = lcid *
    // sendTokens()

    */
    /*
    if(ret) {
        printf("ecall_storage_auth : ret = SUCCESS\n");
    }
    else {
        printf("ecall_storage_auth : ret = FAIL\n");
    }
    */
    return ret;
}


void ecall_supply_storage_buffers(unsigned char *msgbundles,
    uint32_t msgbundles_size, unsigned char *tokens, uint32_t tokens_size)
{
    epoch_msgbundles = msgbundles;
    epoch_tokens = tokens;
    printf("Storage buffers were set");
}