Storage servers send back token bundles and messages every epoch to clients that have authenticated

App/net.cpp

@@ -247,6 +247,7 @@ void NetIO::ing_authenticate_new_client(tcp::socket* csocket,
             if(ret) {
                 ing_receive_msgbundle(csocket, c_simid);
             } else{
+                printf("Client <-> Ingestion authentication failed\n");
                 delete(csocket);
             }
         }
@@ -299,38 +300,10 @@ void NetIO::stg_authenticate_new_client(tcp::socket* csocket,
             if(ret) {
                 uint32_t lcid = c_simid / num_stg_nodes;
                 client_sockets[lcid] = csocket;
-
-                //TODO: Send back this clients tokens for first epoch
-                unsigned char *tkn_ptr = epoch_tokens + (token_bundle_size * lcid);
-
-                /*
-                if(c_simid == 3) {
-                    printf("Just before async_write: Client tokens for c_simid 0:\n");
-                    unsigned char *tmp_tkn_ptr = tkn_ptr + SGX_AESGCM_IV_SIZE;
-                    for(uint32_t i = 0; i < 2 * SGX_AESGCM_KEY_SIZE; i++) {
-                        printf("%x", tmp_tkn_ptr[i]);
-                    }
-                    printf("\n");
-                }
-                */
-
-                boost::asio::async_write(*(client_sockets[lcid]),
-                    boost::asio::buffer(tkn_ptr, token_bundle_size),
-                    [this](boost::system::error_code ecc, std::size_t){
-
-                    if (ecc) {
-                        if(ecc == boost::asio::error::eof) {
-                            // Client connection terminated so we delete this socket
-                            // delete(csocket);
-                            printf("Client socket terminated!\n");
-                        } else {
-                            printf("Error %s\n", ecc.message().c_str());
-                        }
-                        return;
-                    }
-                });
-            } else {
-                printf("Net: Storage Authentication FAIL\n");
+            }
+            else{
+                printf("Client <-> Storage authentication failed\n");
+                delete (csocket);
             }
         }
     });
@@ -362,6 +335,70 @@ void NetIO::stg_start_accept()
         boost::asio::placeholders::error));
 }
 
+
+void NetIO::send_client_mailbox()
+{
+    // Send each client their tokens for the next epoch
+    for(uint32_t lcid = 0; lcid < num_clients_per_stg; lcid++)
+    {
+        unsigned char *tkn_ptr = epoch_tokens + lcid * token_bundle_size;
+        unsigned char *buf_ptr = epoch_msgbundles + lcid * msgbundle_size;
+
+        /*
+        if(lcid == 0) {
+            printf("In NetiO::send_client_mailbox. Token_bundle for lcid 0:\n");
+            for(int k=0; k<token_bundle_size; k++) {
+                printf("%x", tkn_ptr[k]);
+            }
+            printf("\n");
+        }
+        */
+
+        if(client_sockets[lcid]!=nullptr) {
+            boost::asio::async_write(*(client_sockets[lcid]),
+                boost::asio::buffer(tkn_ptr, token_bundle_size),
+                [this, lcid, buf_ptr](boost::system::error_code ec, std::size_t){
+
+                //printf("NetIO::send_client_mailbox, Client %d tokens was sent\n", lcid);
+                if (ec) {
+                    if(ec == boost::asio::error::eof) {
+                        // Client connection terminated so we delete this socket
+                        // delete(csocket);
+                        printf("Client socket terminated!\n");
+                    } else {
+                        printf("Error %s\n", ec.message().c_str());
+                    }
+                    return;
+                }
+
+                boost::asio::async_write(*(client_sockets[lcid]),
+                    boost::asio::buffer(buf_ptr, msgbundle_size),
+                    [this, lcid](boost::system::error_code ecc, std::size_t){
+
+                    //printf("NetIO::send_client_mailbox, Client %d messages was sent\n", lcid);
+                    if (ecc) {
+                        if(ecc == boost::asio::error::eof) {
+                            // Client connection terminated so we delete this socket
+                            // delete(csocket);
+                            printf("Client socket terminated!\n");
+                        } else {
+                            printf("Error %s\n", ecc.message().c_str());
+                        }
+                        return;
+                    }
+                });
+            });
+        }
+        /*
+        else {
+            printf("Client did not have a socket!\n");
+        }
+        */
+    }
+
+
+}
+
 NetIO::NetIO(boost::asio::io_context &io_context, const Config &config)
     : context(io_context), conf(config),
       myconf(config.nodes[config.my_node_num])
@@ -434,8 +471,9 @@ NetIO::NetIO(boost::asio::io_context &io_context, const Config &config)
 
 
     auth_size = sizeof(clientid_t) + sizeof(unsigned long) + SGX_AESGCM_KEY_SIZE;
-    msgbundle_size = SGX_AESGCM_IV_SIZE +
-        (apiparams.m_priv_out * apiparams.msg_size) + SGX_AESGCM_MAC_SIZE;
+    msgbundle_size = SGX_AESGCM_IV_SIZE
+        + (apiparams.m_priv_out * (apiparams.msg_size + TOKEN_SIZE))
+        + SGX_AESGCM_MAC_SIZE;
 
 
     if(myconf.roles & ROLE_STORAGE) {
@@ -450,13 +488,16 @@ NetIO::NetIO(boost::asio::io_context &io_context, const Config &config)
                 num_stg_nodes++;
             }
         }
-        uint32_t num_clients_per_stg = CEILDIV(num_users, num_stg_nodes);
-        client_sockets.resize(num_clients_per_stg);
-        uint16_t num_priv_channels = config.m_priv_in;
-        uint16_t msg_size = config.msg_size;
-        uint32_t epoch_msgbundles_size = num_clients_per_stg *
-            (SGX_AESGCM_IV_SIZE + msg_size * num_priv_channels + SGX_AESGCM_MAC_SIZE);
-        token_bundle_size = ((num_priv_channels * SGX_AESGCM_KEY_SIZE) + SGX_AESGCM_IV_SIZE + SGX_AESGCM_MAC_SIZE);
+        num_clients_per_stg = CEILDIV(num_users, num_stg_nodes);
+
+        for(uint32_t i = 0; i<num_clients_per_stg; i++) {
+            client_sockets.emplace_back(nullptr);
+        }
+
+        uint16_t num_priv_channels = config.m_priv_out;
+        uint32_t epoch_msgbundles_size = num_clients_per_stg * msgbundle_size;
+        token_bundle_size = ((num_priv_channels * SGX_AESGCM_KEY_SIZE)
+            + SGX_AESGCM_IV_SIZE + SGX_AESGCM_MAC_SIZE);
         uint32_t epoch_tokens_size = num_clients_per_stg * token_bundle_size;
         epoch_msgbundles = (unsigned char *) malloc(epoch_msgbundles_size);
         epoch_tokens = (unsigned char *) malloc (epoch_tokens_size);

App/net.hpp

@@ -141,6 +141,7 @@ class NetIO {
     void ing_start_accept();
 
     std::vector<tcp::socket*> client_sockets;
+    uint32_t num_clients_per_stg;
     unsigned char *epoch_tokens;
     unsigned char *epoch_msgbundles;
     uint16_t num_stg_nodes;
@@ -159,6 +160,7 @@ public:
         return nodeios[node_num].value();
     }
     const Config &config() { return conf; }
+    const NodeConfig &myconfig() { return myconf; }
     boost::asio::io_context &io_context() { return context; }
     // Call recv_commands with these arguments on each of the nodes (not
     // including ourselves)
@@ -166,6 +168,8 @@ public:
         std::function<void(boost::system::error_code)> error_cb,
         std::function<void(uint32_t)> epoch_cb);
 
+    void send_client_mailbox();
+
     // Close all the sockets
     void close();
 };

App/start.cpp

@@ -6,7 +6,7 @@
 #include "start.hpp"
 
 class Epoch {
-    boost::asio::io_context &io_context;
+    NetIO &netio;
     uint32_t epoch_num;
     std::mutex m;
     std::condition_variable cv;
@@ -15,7 +15,7 @@ class Epoch {
     void round_cb(uint32_t round_num) {
         if (round_num) {
             printf("Round %u complete\n", round_num);
-            boost::asio::post(io_context, [this]{
+            boost::asio::post(netio.io_context(), [this]{
                 proceed();
             });
         } else {
@@ -24,13 +24,21 @@ class Epoch {
                 std::lock_guard lk(m);
                 epoch_complete = true;
             }
+
+            const NodeConfig &my_conf = netio.myconfig();
+            if(my_conf.roles & ROLE_STORAGE) {
+                boost::asio::post(netio.io_context(), [this]{
+                    netio.send_client_mailbox();
+                });
+            }
+
             cv.notify_one();
         }
     }
 
 public:
-    Epoch(boost::asio::io_context &context, uint32_t ep_num):
-        io_context(context), epoch_num(ep_num),
+    Epoch(NetIO &netio_obj, uint32_t ep_num):
+        netio(netio_obj), epoch_num(ep_num),
         epoch_complete(false) {}
 
     void proceed() {
@@ -104,7 +112,7 @@ static void epoch(NetIO &netio, char **args) {
         }
     }
 
-    Epoch epoch(netio.io_context(), epoch_num);
+    Epoch epoch(netio, epoch_num);
     epoch.proceed();
     epoch.wait();
     // Launch threads to refill the precomputed Waksman networks we
@@ -122,7 +130,7 @@ static void epoch(NetIO &netio, char **args) {
 static void epoch_clients(NetIO &netio) {
 
     static uint32_t epoch_num = 1;
-    Epoch epoch(netio.io_context(), epoch_num);
+    Epoch epoch(netio, epoch_num);
     epoch.proceed();
     epoch.wait();
     // Launch threads to refill the precomputed Waksman networks we
@@ -155,13 +163,7 @@ static void route_clients_test(NetIO &netio)
     }
 
     // Run epoch
-    for (int i=0; i<3; ++i) {
-        if(i==0) {
-            // Clients setup sockets, authenticate, and get token_bundles
-            // from their storage servers during this interval.
-            usleep(EPOCH_INTERVAL);
-        }
-
+    for (int i=1; i<10; ++i) {
         struct timespec tp;
         clock_gettime(CLOCK_REALTIME_COARSE, &tp);
         unsigned long start = tp.tv_sec * 1000000 + tp.tv_nsec/1000;

Enclave/client.hpp

@@ -10,6 +10,9 @@ struct IngClient {
 
 struct StgClient{
     sgx_aes_gcm_128bit_key_t key;
+    clientid_t my_id;
+    clientid_t *priv_friends;
+    uint8_t iv[SGX_AESGCM_IV_SIZE];
 };
 
 

Enclave/config.cpp

@@ -7,6 +7,9 @@
 
 #define CEILDIV(x,y) (((x)+(y)-1)/(y))
 
+unsigned long ingestion_epoch;
+unsigned long storage_epoch;
+
 Config g_teems_config;
 
 int generateMasterKeys(sgx_aes_gcm_128bit_key_t master_secret,
@@ -168,6 +171,8 @@ bool ecall_config_load(threadid_t nthreads, bool private_routing,
         }
     }
 
+    ingestion_epoch = 0;
+    storage_epoch = 0;
 
     if (!route_init()) {
         return false;

Enclave/config.hpp

@@ -45,5 +45,8 @@ struct Config {
 
 extern Config g_teems_config;
 
+extern unsigned long ingestion_epoch;
+extern unsigned long storage_epoch;
+
 
 #endif

Enclave/ingest.cpp

@@ -88,7 +88,7 @@ bool Ingestion::processMsgBundle(clientid_t cid, unsigned char *msgbundle,
     msgbundle += SGX_AESGCM_IV_SIZE;
 
     uint16_t msg_size = g_teems_config.msg_size;
-    uint32_t msgbundle_size = num_msgs * 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);
 
@@ -97,31 +97,48 @@ bool Ingestion::processMsgBundle(clientid_t cid, unsigned char *msgbundle,
     if(ret!=SGX_SUCCESS) {
         printf("Ingestion::processMsgBundle FAIL\n");
         printf("Error code: %d", (uint32_t) ret);
+        free(dec_msgbundle);
+        return false;
     }
 
-    // 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) {
+    // 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);
-        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);
 
-    memmove(msg_queue.buf + head * msg_size,
-        dec_msgbundle, num_msgs * msg_size);
+        /*
+        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);
+        pthread_mutex_lock(&msg_queue.mutex);
+        msg_queue.inserted += num_msgs;
+        pthread_mutex_unlock(&msg_queue.mutex);
+
+        free(dec_msgbundle);
+        return true;
+    }
 
-    free(dec_msgbundle);
-    return true;
+    return false;
 }
 
 void Ingestion::generateClientKeys(sgx_aes_gcm_128bit_key_t &ESK)

Enclave/storage.cpp

@@ -31,9 +31,9 @@ bool storage_generateClientKeys(uint32_t num_clients, uint32_t my_stg_no) {
     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];
+        clients[i].priv_friends = new clientid_t[g_teems_config.m_priv_out];
         // Initialize this client's private channel friends as themself
-        for(int j =0; j <g_teems_config.m_priv_in; j++) {
+        for(int j =0; j <g_teems_config.m_priv_out; j++) {
             (clients[i].priv_friends)[j] = mid;
         }
     }
@@ -41,7 +41,6 @@ bool storage_generateClientKeys(uint32_t num_clients, uint32_t my_stg_no) {
     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];
@@ -77,6 +76,162 @@ bool storage_generateClientKeys(uint32_t num_clients, uint32_t my_stg_no) {
     return true;
 }
 
+bool generate_all_tokens()
+{
+    uint32_t pt_tokens_size = (g_teems_config.m_priv_out * SGX_AESGCM_KEY_SIZE);
+    uint32_t enc_tokens_size = (g_teems_config.m_priv_out * SGX_AESGCM_KEY_SIZE) +
+        SGX_AESGCM_IV_SIZE + SGX_AESGCM_MAC_SIZE;
+    unsigned char token_body[pt_tokens_size];
+
+    const sgx_aes_gcm_128bit_key_t *pTSK = &(g_teems_config.TSK);
+    for(uint32_t lcid=0; lcid<storage_state.max_users; lcid++) {
+
+        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 [S|R|Epoch] underlying the token in
+        // the correct location for this client in the epoch_tokens buffer
+        // The tokens ( i.e., CMAC over S|R|Epoch) is stored in token_body
+        // Then encrypt token_body with the client's storage key and overwrite
+        // the correct location in epoch_tokens
+
+        memset(token_body, 0, pt_tokens_size);
+        memset(tkn_iv_ptr, 0, SGX_AESGCM_IV_SIZE);
+        // IV = epoch_no, for encrypting the token bundle
+        memcpy(tkn_iv_ptr, &storage_epoch, sizeof(storage_epoch));
+
+        unsigned char *ptr = tkn_ptr;
+        for(int i = 0; i<g_teems_config.m_priv_out; i++)
+        {
+            memcpy(ptr, (&(clients[lcid].my_id)), sizeof(clientid_t));
+            memcpy(ptr + sizeof(clientid_t), (&(clients[lcid].priv_friends[i])), sizeof(clientid_t));
+            memcpy(ptr + 2 * sizeof(clientid_t), &storage_epoch, sizeof(storage_epoch));
+            ptr+=SGX_AESGCM_KEY_SIZE;
+        }
+
+        sgx_status_t ret = SGX_SUCCESS;
+        ret = sgx_rijndael128_cmac_msg(pTSK, tkn_ptr, pt_tokens_size,
+            (sgx_cmac_128bit_tag_t*) &token_body);
+        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 bundle:");
+            for(uint32_t i = 0; i < enc_tokens_size; i++) {
+                printf("%x", tkn_iv_ptr[i]);
+            }
+            printf("\n");
+        }
+        */
+    }
+
+    return true;
+}
+
+
+/* processMsgs
+   - Take all the messages in storage_state.stg_buf
+   - Encrypt them all with their corresponding client key and IV and store into
+      epoch_msgbundles
+   - ecall_send_msgbundle();
+*/
+bool processMsgs() {
+    unsigned char *epoch_buf_ptr = epoch_msgbundles;
+    unsigned char *stg_buf_ptr = storage_state.stg_buf.buf;
+    uint32_t msg_bundle_size = g_teems_config.m_priv_in * g_teems_config.msg_size;
+    uint32_t enc_msg_bundle_size = msg_bundle_size + SGX_AESGCM_IV_SIZE + SGX_AESGCM_MAC_SIZE;
+    sgx_status_t ret = SGX_SUCCESS;
+    unsigned char *epoch_buf_ct_ptr = epoch_buf_ptr + SGX_AESGCM_IV_SIZE;
+    unsigned char *epoch_buf_tag_ptr = epoch_buf_ct_ptr + msg_bundle_size;
+
+    for(uint32_t lcid = 0; lcid <storage_state.max_users; lcid++) {
+        memcpy(epoch_buf_ptr, clients[lcid].iv, SGX_AESGCM_IV_SIZE);
+        ret = sgx_rijndael128GCM_encrypt(&(clients[lcid].key), stg_buf_ptr, msg_bundle_size,
+            (uint8_t*) epoch_buf_ct_ptr, epoch_buf_ptr, SGX_AESGCM_IV_SIZE, NULL, 0,
+            (sgx_aes_gcm_128bit_tag_t*) epoch_buf_tag_ptr);
+        if(ret!=SGX_SUCCESS) {
+            printf("processMsgs: Encrypting msgs FAIL\n");
+            return false;
+        }
+
+        // Update IV
+        uint64_t *iv_ctr = (uint64_t*) clients[lcid].iv;
+        (*iv_ctr)+=1;
+
+        if(lcid==0) {
+            printf("\n\nMessage for lcid 0, S, R = %d, %d\n\n\n", *((uint32_t*) stg_buf_ptr),
+                *((uint32_t*) (stg_buf_ptr + 4)));
+        }
+
+        stg_buf_ptr+=msg_bundle_size;
+        epoch_buf_ptr+=enc_msg_bundle_size;
+        epoch_buf_ct_ptr+=enc_msg_bundle_size;
+        epoch_buf_tag_ptr+=enc_msg_bundle_size;
+    }
+
+    return true;
+}
+
+bool generateDummies() {
+    unsigned char *epoch_buf_ptr = epoch_msgbundles;
+    uint32_t msg_bundle_size = g_teems_config.m_priv_in * g_teems_config.msg_size;
+    uint32_t enc_msg_bundle_size = msg_bundle_size + SGX_AESGCM_IV_SIZE + SGX_AESGCM_MAC_SIZE;
+    sgx_status_t ret = SGX_SUCCESS;
+    unsigned char *epoch_buf_ct_ptr = epoch_buf_ptr + SGX_AESGCM_IV_SIZE;
+    unsigned char *epoch_buf_tag_ptr = epoch_buf_ct_ptr + msg_bundle_size;
+    unsigned char *zeroes = new unsigned char[msg_bundle_size];
+    memset(zeroes, 0, msg_bundle_size);
+
+    for(uint32_t lcid = 0; lcid <storage_state.max_users; lcid++) {
+        memcpy(epoch_buf_ptr, clients[lcid].iv, SGX_AESGCM_IV_SIZE);
+        ret = sgx_rijndael128GCM_encrypt(&(clients[lcid].key), zeroes, msg_bundle_size,
+            (uint8_t*) epoch_buf_ct_ptr, epoch_buf_ptr, SGX_AESGCM_IV_SIZE, NULL, 0,
+            (sgx_aes_gcm_128bit_tag_t*) epoch_buf_tag_ptr);
+        if(ret!=SGX_SUCCESS) {
+            printf("processMsgs: Encrypting msgs FAIL\n");
+            return false;
+        }
+
+        // Update IV
+        uint64_t *iv_ctr = (uint64_t*) clients[lcid].iv;
+        (*iv_ctr)+=1;
+
+        epoch_buf_ptr+=enc_msg_bundle_size;
+        epoch_buf_ct_ptr+=enc_msg_bundle_size;
+        epoch_buf_tag_ptr+=enc_msg_bundle_size;
+    }
+
+    delete zeroes;
+    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
@@ -249,21 +404,16 @@ void storage_received(MsgBuffer &storage_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,
+    generate_all_tokens();
 
-    // ocall_process_msgbundles(epoch_msgbundles, epoch_tokens);
-    // sendClientTokens();
-    //
+    uint32_t num_expected_msgs = g_teems_config.m_priv_in * storage_state.max_users;
+    if(num_msgs!= 0) {
+        processMsgs();
+    } else {
+        generateDummies();
+    }
 
-    // processStorage(msg_bundles) :
-    //    a) Send out msg_bundles to clients with open sockets
-    //    b) Store the other msg_bundles in "backend"
+    storage_epoch++;
 
     storage_state.stg_buf.reset();
 
@@ -272,92 +422,6 @@ void storage_received(MsgBuffer &storage_buf)
 #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;
@@ -369,26 +433,6 @@ bool ecall_storage_authenticate(clientid_t cid, unsigned char *auth_message)
     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;
 }
 
@@ -398,5 +442,4 @@ void ecall_supply_storage_buffers(unsigned char *msgbundles,
 {
     epoch_msgbundles = msgbundles;
     epoch_tokens = tokens;
-    printf("Storage buffers were set");
 }