Complete handshake message 2, start handshake message 3

Enclave/comms.cpp

@@ -1,6 +1,7 @@
 #include <vector>
 #include <functional>
 #include <cstring>
+#include <stdio.h>
 
 #include "sgx_tcrypto.h"
 #include "sgx_tseal.h"
@@ -27,11 +28,8 @@ struct NodeCommState {
     HandshakeStep handshake_step;
 
     // Our DH keypair during the handshake
-    sgx_ec256_private_t handshake_privkey;
-    sgx_ec256_public_t handshake_pubkey;
-
-    // The peer's DH public key during the handshake
-    sgx_ec256_public_t handshake_peer_pubkey;
+    sgx_ec256_private_t handshake_dh_privkey;
+    sgx_ec256_public_t handshake_dh_pubkey;
 
     // The outgoing and incoming AES keys after the handshake
     sgx_aes_gcm_128bit_key_t out_aes_key, in_aes_key;
@@ -133,6 +131,18 @@ static uint8_t* default_in_msg_get_buf(NodeCommState &commst,
     return new uint8_t[max_plaintext_bytes];
 }
 
+static void default_in_msg_received(NodeCommState &nodest,
+    uint8_t *data, uint32_t plaintext_len, uint32_t)
+{
+    printf("Received handshake_3 message of %u bytes:\n", plaintext_len);
+    for (uint32_t i=0;i<plaintext_len;++i) {
+        printf("%02x", data[i]);
+    }
+    printf("\n");
+
+    delete[] data;
+}
+
 static void handshake_1_msg_received(NodeCommState &nodest,
     uint8_t *data, uint32_t plaintext_len, uint32_t);
 static void handshake_2_msg_received(NodeCommState &nodest,
@@ -158,12 +168,13 @@ static void handshake_1_msg_received(NodeCommState &nodest,
         return;
     }
     sgx_ecc_state_handle_t ecc_handle;
-    sgx_ec256_public_t peer_pubkey;
-    memmove(&peer_pubkey, data, sizeof(peer_pubkey));
+    sgx_ec256_public_t peer_dh_pubkey;
+    memmove(&peer_dh_pubkey, data, sizeof(peer_dh_pubkey));
     delete[] data;
     sgx_ecc256_open_context(&ecc_handle);
     int valid;
-    if (sgx_ecc256_check_point(&peer_pubkey, ecc_handle, &valid) || !valid) {
+    if (sgx_ecc256_check_point(&peer_dh_pubkey, ecc_handle, &valid)
+            || !valid) {
         printf("Invalid public key received from node %hu\n",
             nodest.node_num);
         sgx_ecc256_close_context(ecc_handle);
@@ -173,15 +184,15 @@ static void handshake_1_msg_received(NodeCommState &nodest,
     printf("Valid public key received from node %hu\n", nodest.node_num);
 
     // Create our own DH key pair
-    sgx_ec256_public_t our_pubkey;
-    sgx_ec256_private_t our_privkey;
-    sgx_ecc256_create_key_pair(&our_privkey, &our_pubkey, ecc_handle);
+    sgx_ec256_public_t our_dh_pubkey;
+    sgx_ec256_private_t our_dh_privkey;
+    sgx_ecc256_create_key_pair(&our_dh_privkey, &our_dh_pubkey, ecc_handle);
 
     // Construct the shared secret
     sgx_ec256_dh_shared_t sharedsecret;
-    sgx_ecc256_compute_shared_dhkey(&our_privkey, &peer_pubkey,
+    sgx_ecc256_compute_shared_dhkey(&our_dh_privkey, &peer_dh_pubkey,
         &sharedsecret, ecc_handle);
-    memset(&our_privkey, 0, sizeof(our_privkey));
+    memset(&our_dh_privkey, 0, sizeof(our_dh_privkey));
 
     // Compute H1(sharedsecret) and H2(sharedsecret)
     sgx_sha_state_handle_t sha_handle;
@@ -205,9 +216,9 @@ static void handshake_1_msg_received(NodeCommState &nodest,
     sgx_hmac_state_handle_t hmac_handle;
     uint8_t srv_cli_mac[16];
     sgx_hmac256_init(h1, 16, &hmac_handle);
-    sgx_hmac256_update((uint8_t*)&our_pubkey, sizeof(our_pubkey),
+    sgx_hmac256_update((uint8_t*)&our_dh_pubkey, sizeof(our_dh_pubkey),
         hmac_handle);
-    sgx_hmac256_update((uint8_t*)&peer_pubkey, sizeof(peer_pubkey),
+    sgx_hmac256_update((uint8_t*)&peer_dh_pubkey, sizeof(peer_dh_pubkey),
         hmac_handle);
     sgx_hmac256_update((uint8_t*)&g_pubkey, sizeof(g_pubkey),
         hmac_handle);
@@ -219,9 +230,9 @@ static void handshake_1_msg_received(NodeCommState &nodest,
     // Compute the client-to-server MAC
     uint8_t cli_srv_mac[16];
     sgx_hmac256_init(((uint8_t*)h1)+16, 16, &hmac_handle);
-    sgx_hmac256_update((uint8_t*)&peer_pubkey, sizeof(peer_pubkey),
+    sgx_hmac256_update((uint8_t*)&peer_dh_pubkey, sizeof(peer_dh_pubkey),
         hmac_handle);
-    sgx_hmac256_update((uint8_t*)&our_pubkey, sizeof(our_pubkey),
+    sgx_hmac256_update((uint8_t*)&our_dh_pubkey, sizeof(our_dh_pubkey),
         hmac_handle);
     sgx_hmac256_update((uint8_t*)&nodest.pubkey, sizeof(nodest.pubkey),
         hmac_handle);
@@ -242,24 +253,139 @@ static void handshake_1_msg_received(NodeCommState &nodest,
     nodest.handshake_step = HANDSHAKE_S_SENT_2;
 
     // Send handshake message 2
-    nodest.message_start(sizeof(our_pubkey) + sizeof(srv_cli_sig));
-    nodest.message_data((uint8_t*)&our_pubkey, sizeof(our_pubkey));
+    nodest.message_start(sizeof(our_dh_pubkey) + sizeof(srv_cli_sig));
+    nodest.message_data((uint8_t*)&our_dh_pubkey, sizeof(our_dh_pubkey));
     nodest.message_data((uint8_t*)&srv_cli_sig, sizeof(srv_cli_sig));
 }
 
+// Receive (at the client) the secong handshake message
 static void handshake_2_msg_received(NodeCommState &nodest,
     uint8_t *data, uint32_t plaintext_len, uint32_t)
 {
+    /*
     printf("Received handshake_2 message of %u bytes:\n", plaintext_len);
     for (uint32_t i=0;i<plaintext_len;++i) {
         printf("%02x", data[i]);
     }
     printf("\n");
+    */
+
+    if (plaintext_len != sizeof(sgx_ec256_public_t) +
+            sizeof(sgx_ec256_signature_t)) {
+        printf("Received handshake_2 message of incorrect size %u\n",
+            plaintext_len);
+        return;
+    }
+    sgx_ecc_state_handle_t ecc_handle;
+    sgx_ec256_public_t peer_dh_pubkey;
+    sgx_ec256_signature_t peer_sig;
+    memmove(&peer_dh_pubkey, data, sizeof(peer_dh_pubkey));
+    memmove(&peer_sig, data+sizeof(peer_dh_pubkey), sizeof(peer_sig));
+    delete[] data;
+    sgx_ecc256_open_context(&ecc_handle);
+    int valid;
+    if (sgx_ecc256_check_point(&peer_dh_pubkey, ecc_handle, &valid)
+            || !valid) {
+        printf("Invalid public key received from node %hu\n",
+            nodest.node_num);
+        sgx_ecc256_close_context(ecc_handle);
+        return;
+    }
+
+    // Construct the shared secret
+    sgx_ec256_dh_shared_t sharedsecret;
+    sgx_ecc256_compute_shared_dhkey(&nodest.handshake_dh_privkey,
+        &peer_dh_pubkey, &sharedsecret, ecc_handle);
+    memset(&nodest.handshake_dh_privkey, 0,
+        sizeof(nodest.handshake_dh_privkey));
+
+    // Compute H1(sharedsecret) and H2(sharedsecret)
+    sgx_sha_state_handle_t sha_handle;
+    sgx_sha256_hash_t h1, h2;
+
+    sgx_sha256_init(&sha_handle);
+    sgx_sha256_update((const uint8_t*)"\x01", 1, sha_handle);
+    sgx_sha256_update((uint8_t*)&sharedsecret, sizeof(sharedsecret),
+        sha_handle);
+    sgx_sha256_get_hash(sha_handle, &h1);
+    sgx_sha256_close(sha_handle);
+
+    sgx_sha256_init(&sha_handle);
+    sgx_sha256_update((const uint8_t*)"\x02", 1, sha_handle);
+    sgx_sha256_update((uint8_t*)&sharedsecret, sizeof(sharedsecret),
+        sha_handle);
+    sgx_sha256_get_hash(sha_handle, &h2);
+    sgx_sha256_close(sha_handle);
+
+    // Compute the server-to-client MAC
+    sgx_hmac_state_handle_t hmac_handle;
+    uint8_t srv_cli_mac[16];
+    sgx_hmac256_init(h1, 16, &hmac_handle);
+    sgx_hmac256_update((uint8_t*)&peer_dh_pubkey, sizeof(peer_dh_pubkey),
+        hmac_handle);
+    sgx_hmac256_update((uint8_t*)&nodest.handshake_dh_pubkey,
+        sizeof(nodest.handshake_dh_pubkey), hmac_handle);
+    sgx_hmac256_update((uint8_t*)&nodest.pubkey, sizeof(nodest.pubkey),
+        hmac_handle);
+    sgx_hmac256_update((uint8_t*)&g_pubkey, sizeof(g_pubkey),
+        hmac_handle);
+    sgx_hmac256_final(srv_cli_mac, 16, hmac_handle);
+    sgx_hmac256_close(hmac_handle);
+
+    // Compute the client-to-server MAC
+    uint8_t cli_srv_mac[16];
+    sgx_hmac256_init(((uint8_t*)h1)+16, 16, &hmac_handle);
+    sgx_hmac256_update((uint8_t*)&nodest.handshake_dh_pubkey,
+        sizeof(nodest.handshake_dh_pubkey), hmac_handle);
+    sgx_hmac256_update((uint8_t*)&peer_dh_pubkey, sizeof(peer_dh_pubkey),
+        hmac_handle);
+    sgx_hmac256_update((uint8_t*)&g_pubkey, sizeof(g_pubkey),
+        hmac_handle);
+    sgx_hmac256_update((uint8_t*)&nodest.pubkey, sizeof(nodest.pubkey),
+        hmac_handle);
+    sgx_hmac256_final(cli_srv_mac, 16, hmac_handle);
+    sgx_hmac256_close(hmac_handle);
+
+    // Verify the signature on the server-to-client MAC
+    uint8_t result;
+    if (sgx_ecdsa_verify(srv_cli_mac, 16, &nodest.pubkey, &peer_sig,
+            &result, ecc_handle) || result != SGX_EC_VALID) {
+        printf("Invalid signature received from node %hu\n",
+            nodest.node_num);
+        sgx_ecc256_close_context(ecc_handle);
+        return;
+    }
+
+    printf("Valid signature received from node %hu\n", nodest.node_num);
+
+    // Sign the client-to-server MAC
+    sgx_ec256_signature_t cli_srv_sig;
+    sgx_ecdsa_sign(cli_srv_mac, 16, &g_privkey, &cli_srv_sig, ecc_handle);
+
+    sgx_ecc256_close_context(ecc_handle);
+
+    // Our side of the handshake is complete
+    memmove(&nodest.out_aes_key, h2, 16);
+    memmove(&nodest.in_aes_key, ((uint8_t*)h2)+16, 16);
+    nodest.handshake_step = HANDSHAKE_COMPLETE;
+    nodest.in_msg_get_buf = default_in_msg_get_buf;
+    nodest.in_msg_received = default_in_msg_received;
+
+    // Send handshake message 3
+    nodest.message_start(sizeof(cli_srv_sig));
+    nodest.message_data((uint8_t*)&cli_srv_sig, sizeof(cli_srv_sig));
 }
 
 static void handshake_3_msg_received(NodeCommState &nodest,
     uint8_t *data, uint32_t plaintext_len, uint32_t)
 {
+    printf("Received handshake_3 message of %u bytes:\n", plaintext_len);
+    for (uint32_t i=0;i<plaintext_len;++i) {
+        printf("%02x", data[i]);
+    }
+    printf("\n");
+
+    delete[] data;
 }
 
 // Start a new outgoing message.  Pass the number of _plaintext_ bytes
@@ -578,7 +704,7 @@ void NodeCommState::handshake_start()
     sgx_ecc256_open_context(&ecc_handle);
 
     // Create a DH keypair
-    sgx_ecc256_create_key_pair(&handshake_privkey, &handshake_pubkey,
+    sgx_ecc256_create_key_pair(&handshake_dh_privkey, &handshake_dh_pubkey,
         ecc_handle);
 
     sgx_ecc256_close_context(ecc_handle);
@@ -589,9 +715,9 @@ void NodeCommState::handshake_start()
     handshake_step = HANDSHAKE_C_SENT_1;
 
     // Send the public key as the first message
-    message_start(sizeof(handshake_pubkey));
+    message_start(sizeof(handshake_dh_pubkey));
 
-    message_data((uint8_t*)&handshake_pubkey, sizeof(handshake_pubkey));
+    message_data((uint8_t*)&handshake_dh_pubkey, sizeof(handshake_dh_pubkey));
 }
 
 // Start all handshakes for which we are the client