teems/Enclave/ingest.cpp

#include <pthread.h>
#include "Enclave_t.h"
#include "utils.hpp"
#include "config.hpp"
#include "route.hpp"
#include "ingest.hpp"

Ingestion g_ing;

void displayMessage(unsigned char *msg, uint16_t msg_size) {
    clientid_t sid, rid;
    unsigned char *ptr = msg;
    sid = *((clientid_t*) ptr);
    ptr+=sizeof(sid);
    rid = *((clientid_t*) ptr);
    ptr+=sizeof(rid);
    printf("Sender ID: %d, Receiver ID: %d, Token: N/A\n", sid, rid );
    printf("Message: ");
    for(int j = 0; j<msg_size - sizeof(sid)*2; j++) {
        printf("%x", (*ptr));
        ptr++;
    }
    printf("\n");
}

void displayEncMessageBundle(unsigned char *bundle, uint16_t priv_out, uint16_t msg_size) {
    unsigned char *ptr = bundle;

    printf("IV: ");
    for(int i=0; i<SGX_AESGCM_IV_SIZE; i++) {
        printf("%x", ptr[i]);
    }
    printf("\n");
    ptr+= SGX_AESGCM_IV_SIZE;

    for(int i=0; i<priv_out; i++) {
        displayMessage(ptr, msg_size);
        ptr+=msg_size;
    }

    printf("MAC: ");
    for(int i=0; i<SGX_AESGCM_MAC_SIZE; i++) {
        printf("%x", ptr[i]);
    }
    printf("\n");
}

bool ecall_ingest_msgbundle(clientid_t cid, unsigned char *msgbundle,
    uint32_t num_msgs)
{
    bool ret;
    ret = g_ing.processMsgBundle(cid, msgbundle, num_msgs);
    return ret;
}

bool ecall_authenticate(clientid_t cid, unsigned char *auth_message)
{
    bool ret;
    ret = g_ing.authenticate(cid, auth_message);
    return ret;
}

void Ingestion::initialize(uint32_t num, uint32_t start, sgx_aes_gcm_128bit_key_t &ESK) {
    cnum = num;
    cstart = start;
    clients = new IngClient[cnum];

    generateClientKeys(ESK);

    max_buffer_size = g_teems_config.m_priv_out * cnum;
    buffer = &(route_state.ingbuf);
}

bool Ingestion::authenticate(clientid_t cid, unsigned char *auth_message)
{
    uint32_t lcid = cid-cstart;
    const sgx_aes_gcm_128bit_key_t *ckey = &(clients[lcid].key);
    return(authenticateClient(auth_message, ckey));
}

bool Ingestion::processMsgBundle(clientid_t cid, unsigned char *msgbundle,
    uint32_t num_msgs) {

    // Fetch corresponding client key
    clientid_t lcid = cid - g_ing.cstart;
    sgx_aes_gcm_128bit_key_t *ckey = &(clients[lcid].key);
    unsigned char *iv = msgbundle;
    msgbundle += SGX_AESGCM_IV_SIZE;

    uint16_t msg_size = g_teems_config.msg_size;
    uint32_t msgbundle_size = num_msgs * (msg_size + TOKEN_SIZE);
    unsigned char *dec_msgbundle = (unsigned char *) malloc (msgbundle_size);
    sgx_aes_gcm_128bit_tag_t *tag = (sgx_aes_gcm_128bit_tag_t*) (msgbundle + msgbundle_size);

    sgx_status_t ret = sgx_rijndael128GCM_decrypt(ckey, msgbundle, msgbundle_size,
        dec_msgbundle, iv, SGX_AESGCM_IV_SIZE, NULL, 0, tag);
    if(ret!=SGX_SUCCESS) {
        printf("Ingestion::processMsgBundle FAIL\n");
        printf("Error code: %d", (uint32_t) ret);
        free(dec_msgbundle);
        return false;
    }

    // TODO: Verify the tokens from end of the msgbundle
    // before appending them into the MsgBuffer
    bool verified  = true;

    if(verified) {
        // Append msgbundle to g_ing.buffer;
        MsgBuffer &msg_queue = *(g_ing.buffer);

        pthread_mutex_lock(&msg_queue.mutex);
        uint32_t head = msg_queue.reserved;
        if (head + num_msgs > g_ing.max_buffer_size) {
            pthread_mutex_unlock(&msg_queue.mutex);
            printf("Ingestions: Max %u messages exceeded\n",
                g_ing.max_buffer_size);
            return false;
        }
        msg_queue.reserved += num_msgs;
        pthread_mutex_unlock(&msg_queue.mutex);

        /*
        if(lcid==0) {
            printf("\n\nIngestion: Message for lcid 0, S, R = %ld, %ld\n\n\n", *((uint32_t*) dec_msgbundle),
                *((uint32_t*) (dec_msgbundle + 4)));
        }
        */

        memmove(msg_queue.buf + head * msg_size,
            dec_msgbundle, num_msgs * msg_size);

        pthread_mutex_lock(&msg_queue.mutex);
        msg_queue.inserted += num_msgs;
        pthread_mutex_unlock(&msg_queue.mutex);

        free(dec_msgbundle);
        return true;
    }

    return false;
}

void Ingestion::generateClientKeys(sgx_aes_gcm_128bit_key_t &ESK)
{
    printf("In Ingestion::genCK, num_clients = %d, client_start = %d, client_end = %d\n",
        cnum, cstart, cnum + cstart);
    for(uint32_t i=0; i<cnum; i++)
    {
        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);

        uint32_t client_num = cstart + i;
        memcpy(iv, (uint8_t*) (&client_num), sizeof(client_num));

        sgx_status_t ret = SGX_SUCCESS;
        ret = sgx_rijndael128GCM_encrypt((const sgx_aes_gcm_128bit_key_t *) (ESK),
            zeroes, SGX_AESGCM_KEY_SIZE, (uint8_t*) (clients[i].key), iv,
            SGX_AESGCM_IV_SIZE, NULL, 0, &tag);
        if(ret!=SGX_SUCCESS) {
            printf("Ingestion::GCK FAIL\n");
        }

    }
}

sgx_aes_gcm_128bit_key_t* Ingestion::getClientKey(uint32_t lcid)
{
    return(&(clients[lcid].key));
}