Browse Source

Ingestion servers handle client connections asynchronously correctly (using async_read now). Clients send CLIENTS_AUTHENTICATE message first, then several CLIENT_MESSAGE_BUNDLE messages

Sajin Sasy 1 year ago
parent
commit
1aeb41385b
8 changed files with 206 additions and 49 deletions
  1. 103 37
      App/net.cpp
  2. 5 4
      App/net.hpp
  3. 42 5
      Client/clients.cpp
  4. 2 0
      Enclave/Enclave.edl
  5. 41 2
      Enclave/ingest.cpp
  6. 2 0
      Enclave/ingest.hpp
  7. 9 1
      Untrusted/Untrusted.cpp
  8. 2 0
      Untrusted/Untrusted.hpp

+ 103 - 37
App/net.cpp

@@ -172,59 +172,128 @@ void NodeIO::recv_commands(
 }
 
 /*
-  Handler for received client messages.
-
+    Receive clients dropped off messages, i.e. a CLIENT_MESSAGE_BUNDLE
 */
-void NetIO::handle_async_clients(std::shared_ptr<tcp::socket> csocket,
-    const boost::system::error_code& error, size_t auth_size,
-    size_t msgbundle_size)
+
+void NetIO::receive_msgbundle(tcp::socket* csocket)
 {
-    if(!error) {
-#ifdef VERBOSE_NET
-        printf("Accept handler success\n");
-#endif
-        // Read header (1 uint64_t) from the socket and extract the client ID
-        size_t header;
-        clientid_t cid;
-        boost::asio::read(*csocket,
-             boost::asio::buffer(&header, sizeof(uint64_t)));
+    // Read header (1 uint64_t) from the socket and extract the client ID
+    size_t *header = new size_t;
+    boost::asio::async_read(*csocket, boost::asio::buffer(header, sizeof(size_t)),
+        [this, csocket, header]
+        (boost::system::error_code ec, std::size_t) {
+        if (ec) {
+            if(ec == boost::asio::error::eof) {
+                // Client connection terminated so we delete this socket
+                delete(csocket);
+            }
+            else {
+                printf("Error %s\n", ec.message().c_str());
+            }
+            return;
+        }
 
-        if((header & 0xff) == CLIENT_AUTHENTICATE) {
+        if((*header & 0xff) == CLIENT_MESSAGE_BUNDLE) {
+            clientid_t cid = (clientid_t)(*header >> 8);
             // Read the authentication token
-            boost::asio::read(*csocket,
-               boost::asio::buffer(&header, auth_size));
-
-        } else if ((header & 0xff) == CLIENT_MESSAGE_BUNDLE) {
+            delete(header);
             unsigned char *msgbundle = (unsigned char*) malloc(msgbundle_size);
-            cid = (clientid_t)(header >> 8);
 
-            // Read the message_bundle
-            boost::asio::read(*csocket,
-               boost::asio::buffer(msgbundle, msgbundle_size));
+            boost::asio::async_read(*csocket, boost::asio::buffer(msgbundle, msgbundle_size),
+                [this, csocket, msgbundle, cid]
+                (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);
+                    }
+                    else {
+                        printf("Error %s\n", ecc.message().c_str());
+                    }
+                    return;
+                }
+
+                //Ingest the message_bundle
+                bool ret = ecall_ingest_msgbundle(cid, msgbundle, apiparams.m_priv_out);
+                free(msgbundle);
+            });
 
-            //Ingest the message_bundle
-            bool ret = ecall_ingest_msgbundle(cid, msgbundle, apiparams.m_priv_out);
-            free(msgbundle);
+        // Continue to async receive client message bundles
+        receive_msgbundle(csocket);
         }
+    });
+}
 
-        start_accept(auth_size, msgbundle_size);
-    } else {
+/*
+    Handle new client connections.
+    New clients always send a CLIENT_AUTHENTICATE message, and then followed
+    by their CLIENT_MESSAGE_BUNDLE every epoch.
+*/
+void NetIO::authenticate_new_client(tcp::socket* csocket,
+    const boost::system::error_code& error)
+{
+    if(error) {
         printf("Accept handler failed\n");
+        return;
     }
+
+#ifdef VERBOSE_NET
+        printf("Accept handler success\n");
+#endif
+
+    // Read header (1 uint64_t) from the socket and extract the client ID
+    size_t *header = new size_t;
+    boost::asio::async_read(*csocket, boost::asio::buffer(header, sizeof(size_t)),
+        [this, csocket, header]
+        (boost::system::error_code ec, std::size_t) {
+        if (ec) {
+            if(ec == boost::asio::error::eof) {
+                // Client connection terminated so we delete this socket
+                delete(csocket);
+            } else {
+                printf("Error %s\n", ec.message().c_str());
+            }
+            return;
+        }
+        else {
+            if((*header & 0xff) == CLIENT_AUTHENTICATE) {
+                clientid_t cid = (clientid_t)(*header >> 8);
+                // Read the authentication token
+                unsigned char* auth_string = (unsigned char*) malloc(auth_size);
+                delete(header);
+
+                boost::asio::async_read(*csocket, boost::asio::buffer(auth_string, auth_size),
+                    [this, csocket, auth_string, cid]
+                    (boost::system::error_code ecn, std::size_t) {
+                    if (ecn) {
+                        printf("Error %s\n", ecn.message().c_str());
+                        return;
+                    }
+
+                    bool ret = ecall_authenticate(cid, auth_string);
+                    free(auth_string);
+                });
+
+            // Receive client message bundles on this socket
+            receive_msgbundle(csocket);
+            }
+        }
+    });
+    start_accept();
 }
 
 /*
-  Asynchronously accept client connections
+  Asynchronously accept new client connections
 */
-void NetIO::start_accept(size_t auth_size, size_t msgbundle_size)
+void NetIO::start_accept()
 {
-    std::shared_ptr<tcp::socket> csocket(new tcp::socket(io_context()));
+    tcp::socket *csocket = new tcp::socket(io_context());
 #ifdef VERBOSE_NET
     std::cout << "Accepting on " << myconf.clistenhost << ":" << myconf.clistenport << "\n";
 #endif
     client_acceptor->async_accept(*csocket,
-        boost::bind(&NetIO::handle_async_clients, this, csocket,
-        boost::asio::placeholders::error, auth_size, msgbundle_size));
+        boost::bind(&NetIO::authenticate_new_client, this, csocket,
+        boost::asio::placeholders::error));
 }
 
 
@@ -304,14 +373,11 @@ NetIO::NetIO(boost::asio::io_context &io_context, const Config &config)
                 resolver.resolve(this->myconf.clistenhost,
                 this->myconf.clistenport)->endpoint()));
 
-        size_t auth_size, msgbundle_size;
         auth_size = SGX_AESGCM_MAC_SIZE;
         msgbundle_size = SGX_AESGCM_IV_SIZE +
             (apiparams.m_priv_out * apiparams.msg_size) + SGX_AESGCM_MAC_SIZE;
-        start_accept(auth_size, msgbundle_size);
+        start_accept();
     }
-
-
 }
 
 

+ 5 - 4
App/net.hpp

@@ -131,10 +131,11 @@ class NetIO {
     std::deque<std::optional<NodeIO>> nodeios;
     std::shared_ptr<tcp::acceptor> client_acceptor;
 
-    void handle_async_clients(std::shared_ptr<tcp::socket> socket,
-        const boost::system::error_code& error, size_t auth_size,
-        size_t msgbundle_size);
-    void start_accept(size_t auth_size, size_t msgbundle_size);
+    void receive_msgbundle(tcp::socket* socket);
+    void authenticate_new_client(tcp::socket* socket,
+        const boost::system::error_code& error);
+    void start_accept();
+    size_t auth_size, msgbundle_size;
 
 public:
     NetIO(boost::asio::io_context &io_context, const Config &config);

+ 42 - 5
Client/clients.cpp

@@ -416,6 +416,42 @@ void Client::sendMessageBundle(uint16_t priv_out, uint16_t msg_size,
         boost::asio::buffer(enc_msgbundle, enc_bundle_size));
 }
 
+int Client::sendAuthMessage()
+{
+    uint32_t auth_size = sizeof(uint64_t) + SGX_AESGCM_KEY_SIZE;
+    unsigned char *auth_string = (unsigned char*) malloc(auth_size);
+    unsigned char *as_ptr = auth_string;
+    uint64_t header = (id << 8) + CLIENT_AUTHENTICATE;
+    memcpy(as_ptr, &header, sizeof(header));
+    as_ptr+=sizeof(header);
+
+    unsigned char zeroes[SGX_AESGCM_KEY_SIZE];
+    unsigned char tag[SGX_AESGCM_MAC_SIZE];
+    memset(iv, 0, SGX_AESGCM_IV_SIZE);
+    memset(zeroes, 0, SGX_AESGCM_KEY_SIZE);
+    memset(tag, 0, SGX_AESGCM_KEY_SIZE);
+
+    if (sizeof(zeroes) != gcm_encrypt(zeroes, SGX_AESGCM_KEY_SIZE, NULL, 0, key,
+            iv, SGX_AESGCM_IV_SIZE, as_ptr, tag)) {
+        printf("generateClientEncryptionKey failed\n");
+        return -1;
+    }
+
+    // Update IV
+    uint64_t *iv_ctr = (uint64_t*) iv;
+    (*iv_ctr)+=1;
+
+    printf("Client %d auth_string: \n", id);
+    for(int i=0; i<auth_size; i++) {
+        printf("%x", auth_string[i]);
+    }
+    printf("\n");
+
+    boost::asio::write(*ingestion_sock,
+        boost::asio::buffer(auth_string, auth_size));
+
+    return 1;
+}
 
 void generateClients(boost::asio::io_context &io_context,
     uint32_t cstart, uint32_t cstop, Client* &clients,
@@ -436,7 +472,7 @@ void generateClients(boost::asio::io_context &io_context,
         clients[i].initClient(i, client_key);
 
         clients[i].initializeSocket(io_context, ingestion_nodes[ing_node_this_client]);
-        //clients[i].sendAuthMessage();
+        clients[i].sendAuthMessage();
 
         /*
         // Test that the keys generated match those generated within
@@ -479,6 +515,7 @@ void sendMessageBundles(uint32_t cstart, uint32_t cstop, Client* &clients,
     4) Repeat from 1)
 */
 
+
 int main(int argc, char **argv)
 {
     // Unbuffer stdout
@@ -550,8 +587,8 @@ int main(int argc, char **argv)
     }
 
     // Multithreaded client message bundle generation and send
-    //uint32_t epoch = 0;
-    //while(epoch < 3) {
+    uint32_t epoch = 0;
+    while(epoch < 3) {
         for(int i=0; i<nthreads; i++) {
             uint32_t cstart, cstop;
             cstart = i * clients_per_thread;
@@ -562,7 +599,7 @@ int main(int argc, char **argv)
         for(int i=0; i<nthreads; i++) {
             threads[i].join();
         }
-        //epoch++;
-    //}
+        epoch++;
+    }
     delete [] clients;
 }

+ 2 - 0
Enclave/Enclave.edl

@@ -49,6 +49,8 @@ enclave {
             [user_check] uint8_t *msgbundle,
             uint32_t num_msgs);
 
+        public bool ecall_authenticate(clientid_t cid,
+            [user_check] uint8_t *auth_string);
     };
 
     untrusted {

+ 41 - 2
Enclave/ingest.cpp

@@ -45,13 +45,19 @@ void displayEncMessageBundle(unsigned char *bundle, uint16_t priv_out, uint16_t
 }
 
 bool ecall_ingest_msgbundle(clientid_t cid, unsigned char *msgbundle,
-    uint32_t num_msgs) {
-
+    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_string)
+{
+    bool ret;
+    ret = g_ing.authenticate(cid, auth_string);
+    return ret;
+}
 
 void Ingestion::initialize(uint32_t cnum, uint32_t cstart, sgx_aes_gcm_128bit_key_t &ESK) {
     clients.num = cnum;
@@ -65,6 +71,39 @@ void Ingestion::initialize(uint32_t cnum, uint32_t cstart, sgx_aes_gcm_128bit_ke
     buffer.alloc(max_buffer_size);
 }
 
+bool Ingestion::authenticate(clientid_t cid, unsigned char *auth_string)
+{
+    int auth_success = 0;
+    // Fetch corresponding client key
+    clientid_t lcid = cid - g_ing.clients.start;
+    sgx_aes_gcm_128bit_key_t ckey;
+    memcpy(ckey, (g_ing.clients).keys[lcid], SGX_AESGCM_KEY_SIZE);
+
+    unsigned char computed_auth[SGX_AESGCM_KEY_SIZE];
+    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);
+    sgx_status_t ret = SGX_SUCCESS;
+
+    ret = sgx_rijndael128GCM_encrypt(&ckey, zeroes, SGX_AESGCM_KEY_SIZE,
+        computed_auth, iv, SGX_AESGCM_IV_SIZE, NULL, 0, &mac);
+    if(ret!=SGX_SUCCESS) {
+        return -1;
+    }
+
+    auth_success = memcmp(auth_string, computed_auth, SGX_AESGCM_KEY_SIZE);
+
+    if(auth_success == 0) {
+        printf("authentication SUCCESS\n");
+        return true;
+    } else {
+        printf("authentication FAIL\n");
+        return false;
+    }
+}
+
 bool Ingestion::processMsgBundle(clientid_t cid, unsigned char *msgbundle,
     uint32_t num_msgs) {
 

+ 2 - 0
Enclave/ingest.hpp

@@ -36,6 +36,8 @@ public:
         return (clients.keys);
     }
 
+    bool authenticate(clientid_t cid, unsigned char *auth_string);
+
     bool processMsgBundle(clientid_t cid, unsigned char *msgbundle,
         uint32_t num_msgs);
 };

+ 9 - 1
Untrusted/Untrusted.cpp

@@ -306,8 +306,16 @@ void ocall_routing_round_complete(void *cbpointer, uint32_t round_num)
 }
 
 bool ecall_ingest_msgbundle(clientid_t cid, unsigned char *msgbundle,
-    uint32_t num_msgs) {
+    uint32_t num_msgs)
+{
     bool ret;
     ecall_ingest_msgbundle(global_eid, &ret, cid, msgbundle, num_msgs);
     return ret;
 }
+
+bool ecall_authenticate(clientid_t cid, unsigned char *auth_string)
+{
+    bool ret;
+    ecall_authenticate(global_eid, &ret, cid, auth_string);
+    return ret;
+}

+ 2 - 0
Untrusted/Untrusted.hpp

@@ -43,4 +43,6 @@ void ecall_routing_proceed(std::function<void(uint32_t)> cb);
 bool ecall_ingest_msgbundle(clientid_t cid, unsigned char *msgbundle,
     uint32_t num_msgs);
 
+bool ecall_authenticate(clientid_t cid, unsigned char *auth_string);
+
 #endif