Unify the private and public routing code paths in epoch_process rather than duplicating code

Client/clients.cpp

@@ -887,44 +887,73 @@ void Client::epoch_process() {
     uint32_t pt_token_size = uint32_t(config.m_priv_out) * TOKEN_SIZE;
     uint32_t token_bundle_size = pt_token_size + SGX_AESGCM_IV_SIZE
         + SGX_AESGCM_MAC_SIZE;
-    unsigned char *enc_tokens = (unsigned char*) malloc (token_bundle_size);
+    unsigned char *enc_tokens = nullptr;
+    uint16_t num_in = config.m_pub_in;
 
-    if(private_routing) {
-        //Async read the encrypted tokens for this epoch
-        boost::asio::async_read(*storage_sock, boost::asio::buffer(enc_tokens, token_bundle_size),
-            [this, enc_tokens, token_bundle_size, pt_token_size]
-            (boost::system::error_code ec, std::size_t) {
-
-            if (ec) {
-                if(ec == boost::asio::error::eof) {
-                    delete storage_sock;
-                    storage_sock = nullptr;
-                } else {
-                    printf("Error %s\n", ec.message().c_str());
-                    printf("Client::epoch_process boost "
-                        "async_read_tokens failed\n");
-                }
-                free(enc_tokens);
-                return;
-            }
+    std::vector<boost::asio::mutable_buffer> toreceive;
 
-    #ifdef VERBOSE_CLIENT
-            if(sim_id == 0) {
-                printf("TEST: Client 0: Encrypted token bundle received:\n");
-                for(uint32_t i = 0; i < token_bundle_size; i++) {
-                    printf("%x", enc_tokens[i]);
-                }
-                printf("\n");
+    if (private_routing) {
+        enc_tokens = (unsigned char*) malloc (token_bundle_size);
+        toreceive.push_back(boost::asio::buffer(enc_tokens,
+            token_bundle_size));
+        num_in = config.m_priv_in;
+    }
+
+    uint16_t msg_size = config.msg_size;
+    uint32_t recv_pt_mailbox_size = ptMailboxSize(num_in, msg_size);
+    uint32_t recv_enc_mailbox_size = encMailboxSize(num_in, msg_size);
+    unsigned char *recv_pt_mailbox =
+        (unsigned char*) malloc (recv_pt_mailbox_size);
+    unsigned char *recv_enc_mailbox =
+        (unsigned char*) malloc (recv_enc_mailbox_size);
+
+    toreceive.push_back(boost::asio::buffer(recv_enc_mailbox,
+        recv_enc_mailbox_size));
+
+    // Async read the encrypted tokens (for private routing only) and
+    // encrypted mailbox (both private and public routing) for this
+    // epoch
+    boost::asio::async_read(*storage_sock, toreceive,
+        [this, enc_tokens, token_bundle_size, pt_token_size,
+        recv_pt_mailbox, recv_enc_mailbox]
+        (boost::system::error_code ec, std::size_t) {
+
+        if (ec) {
+            if(ec == boost::asio::error::eof) {
+                delete storage_sock;
+                storage_sock = nullptr;
+            } else {
+                printf("Error %s\n", ec.message().c_str());
+                printf("Client::epoch_process boost "
+                    "async_read_tokens failed\n");
             }
-    #endif
+            free(enc_tokens);
+            free(recv_pt_mailbox);
+            free(recv_enc_mailbox);
+            return;
+        }
+
+#ifdef TRACE_SOCKIO
+        recvlogger.log();
+#endif
 
+#ifdef VERBOSE_CLIENT
+        if(sim_id == 0) {
+            printf("TEST: Client 0: Encrypted token bundle received:\n");
+            for(uint32_t i = 0; i < token_bundle_size; i++) {
+                printf("%02x", enc_tokens[i]);
+            }
+            printf("\n");
+            printf("TEST: Client 0: Encrypted msgbundle received\n");
+        }
+#endif
 
+        if (private_routing) {
             // Decrypt the token bundle
             unsigned char *enc_tkn_ptr = enc_tokens + SGX_AESGCM_IV_SIZE;
             unsigned char *enc_tkn_tag = enc_tokens + SGX_AESGCM_IV_SIZE +
                 pt_token_size;
 
-
             int decrypted_bytes =  gcm_decrypt(enc_tkn_ptr, pt_token_size,
                     NULL, 0, enc_tkn_tag, (unsigned char*) &(this->stg_key),
                     enc_tokens, SGX_AESGCM_IV_SIZE,
@@ -941,108 +970,25 @@ void Client::epoch_process() {
             unsigned char *tkn_ptr = (unsigned char*) this->token_list;
             if(sim_id==0) {
                 printf("TEST: Client 0: Decrypted client tokens:\n");
-                for(int i = 0; i < 2 * SGX_AESGCM_KEY_SIZE; i++) {
-                    printf("%x", tkn_ptr[i]);
+                for(int i = 0; i < pt_token_size; i++) {
+                    printf("%02x", tkn_ptr[i]);
                 }
                 printf("\n");
             }
             */
+        }
 
-            // Async read the messages received in the last epoch
-            uint16_t priv_in = config.m_priv_in;
-            uint16_t msg_size = config.msg_size;
-            uint32_t recv_pt_mailbox_size = ptMailboxSize(priv_in, msg_size);
-            uint32_t recv_enc_mailbox_size = encMailboxSize(priv_in, msg_size);
-            unsigned char *recv_pt_mailbox =
-                (unsigned char*) malloc (recv_pt_mailbox_size);
-            unsigned char *recv_enc_mailbox =
-                (unsigned char*) malloc (recv_enc_mailbox_size);
-
-            boost::asio::async_read(*storage_sock,
-                boost::asio::buffer(recv_enc_mailbox, recv_enc_mailbox_size),
-                [this, recv_pt_mailbox, recv_enc_mailbox]
-                (boost::system::error_code ecc, std::size_t) {
-
-#ifdef TRACE_SOCKIO
-                recvlogger.log();
-#endif
-                if (ecc) {
-                    if(ecc == boost::asio::error::eof) {
-                        delete storage_sock;
-                        storage_sock = nullptr;
-                    } else {
-                        printf("Client: boost async_read failed for "
-                            "receiving msg_bundle\n");
-                        printf("Error %s\n", ecc.message().c_str());
-                    }
-                    free(recv_pt_mailbox);
-                    free(recv_enc_mailbox);
-                    return;
-                }
-
-    #ifdef VERBOSE_CLIENT
-                if(sim_id == 0) {
-                    printf("TEST: Client 0: Encrypted msgbundle received\n");
-                }
-    #endif
-
-                // Do whatever processing with the received messages here
-                // but for the benchmark, we just ignore the received
-                // messages
-
-                free(recv_enc_mailbox);
-                free(recv_pt_mailbox);
-
-                // Send this epoch's message bundle
-                sendMessageBundle();
-                epoch_process();
-            });
-        });
-    } else {
-        // Async read the messages received in the last epoch
-        uint16_t pub_in = config.m_pub_in;
-        uint16_t msg_size = config.msg_size;
-        uint32_t recv_pt_mailbox_size = ptMailboxSize(pub_in, msg_size);
-        uint32_t recv_enc_mailbox_size = encMailboxSize(pub_in, msg_size);
-        unsigned char *recv_pt_mailbox =
-            (unsigned char*) malloc (recv_pt_mailbox_size);
-        unsigned char *recv_enc_mailbox =
-            (unsigned char*) malloc (recv_enc_mailbox_size);
-
-        boost::asio::async_read(*storage_sock,
-            boost::asio::buffer(recv_enc_mailbox, recv_enc_mailbox_size),
-            [this, recv_pt_mailbox, recv_enc_mailbox]
-            (boost::system::error_code ecc, std::size_t) {
-
-#ifdef TRACE_SOCKIO
-            recvlogger.log();
-#endif
-            if (ecc) {
-                if(ecc == boost::asio::error::eof) {
-                    delete storage_sock;
-                    storage_sock = nullptr;
-                } else {
-                    printf("Client: boost async_read failed for "
-                        "receiving msg_bundle\n");
-                    printf("Error %s\n", ecc.message().c_str());
-                }
-                free(recv_pt_mailbox);
-                free(recv_enc_mailbox);
-                return;
-            }
-
-            // Do whatever processing with the received messages here
-            // but for the benchmark, we just ignore the received
-            // messages
+        // Do whatever processing with the received messages here
+        // but for the benchmark, we just ignore the received
+        // messages
 
-            free(recv_enc_mailbox);
-            free(recv_pt_mailbox);
+        free(recv_enc_mailbox);
+        free(recv_pt_mailbox);
 
-            // Send this epoch's message bundle
-            sendMessageBundle();
-            epoch_process();
-        });
-    }
+        // Send this epoch's message bundle
+        sendMessageBundle();
+        epoch_process();
+    });
 }
 
 void initializeClients(boost::asio::io_context &io_context, uint16_t nthreads)