#include <vector>
#include <functional>
#include <cstring>

#include "sgx_tcrypto.h"
#include "sgx_tseal.h"
#include "Enclave_t.h"
#include "utils.hpp"
#include "config.hpp"

// Our public and private identity keys
static sgx_ec256_private_t g_privkey;
static sgx_ec256_public_t g_pubkey;

// What step of the handshake are we on?
enum HandshakeStep {
    HANDSHAKE_NONE,
    HANDSHAKE_C_SENT_1,
    HANDSHAKE_S_SENT_2,
    HANDSHAKE_COMPLETE
};

// Communication state for a node
struct NodeCommState {
    sgx_ec256_public_t pubkey;
    nodenum_t node_num;
    HandshakeStep handshake_step;

    // Our DH keypair during the handshake
    sgx_ec256_private_t handshake_dh_privkey;
    sgx_ec256_public_t handshake_dh_pubkey;

    // The server keeps this state between handshake messages 1 and 3
    uint8_t handshake_cli_srv_mac[16];

    // The outgoing and incoming AES keys after the handshake
    sgx_aes_gcm_128bit_key_t out_aes_key, in_aes_key;

    // The outgoing and incoming IV counters
    uint8_t out_aes_iv[SGX_AESGCM_IV_SIZE];
    uint8_t in_aes_iv[SGX_AESGCM_IV_SIZE];

    // The GCM state for incrementally building each outgoing chunk
    sgx_aes_state_handle_t out_aes_gcm_state;

    // The current outgoing frame and the current offset into it
    uint8_t *frame;
    uint32_t frame_offset;

    // The current outgoing message ciphertext size and the offset into
    // it of the start of the current frame
    uint32_t msg_size;
    uint32_t msg_frame_offset;

    // The current outgoing message plaintext size, how many plaintext
    // bytes we've already processed with message_data, and how many
    // plaintext bytes remain for the current chunk
    uint32_t msg_plaintext_size;
    uint32_t msg_plaintext_processed;
    uint32_t msg_plaintext_chunk_remain;

    // The current incoming message ciphertext size and the offset into
    // it of all previous chunks of this message
    uint32_t in_msg_size;
    uint32_t in_msg_offset;
    // The current incoming message number of plaintext bytes processed
    uint32_t in_msg_plaintext_processed;
    // The internal buffer where we're storing the (decrypted) message
    uint8_t *in_msg_buf;

    // The function to call when a new incoming message header arrives.
    // This function should return a pointer to enough memory to hold
    // the (decrypted) chunks of the message.  Remember that the length
    // passed here is the total size of the _encrypted_ chunks.  This
    // function should not itself modify the in_msg_size, in_msg_offset,
    // or in_msg_buf members.  This function will usually allocate an
    // appropriate amount of memory and return the pointer to it, but
    // may do other things, like return a pointer to the middle of a
    // previously allocated region of memory.
    std::function<uint8_t*(NodeCommState&,uint32_t)> in_msg_get_buf;

    // The function to call after the last chunk of a message has been
    // received.  If in_msg_get_buf allocated memory, this function
    // should deallocate it.  in_msg_size, in_msg_offset,
    // in_msg_plaintext_processed, and in_msg_buf will already have been
    // reset when this function is called.  The uint32_t that is passed
    // are the total size of the _decrypted_ data and the original total
    // size of the _encrypted_ chunks that was passed to in_msg_get_buf.
    std::function<void(NodeCommState&,uint8_t*,uint32_t,uint32_t)>
        in_msg_received;

    NodeCommState(const sgx_ec256_public_t* conf_pubkey, nodenum_t i) :
            node_num(i), handshake_step(HANDSHAKE_NONE),
            out_aes_gcm_state(NULL), frame(NULL),
            frame_offset(0), msg_size(0), msg_frame_offset(0),
            msg_plaintext_size(0), msg_plaintext_processed(0),
            msg_plaintext_chunk_remain(0),
            in_msg_size(0), in_msg_offset(0),
            in_msg_plaintext_processed(0), in_msg_buf(NULL),
            in_msg_get_buf(NULL), in_msg_received(NULL) {
        memmove(&pubkey, conf_pubkey, sizeof(pubkey));
    }

    void message_start(uint32_t plaintext_len, bool encrypt=true);

    void message_data(uint8_t *data, uint32_t len, bool encrypt=true);

    // Start the handshake (as the client)
    void handshake_start();
};

// A typical default in_msg_get_buf handler.  It computes the maximum
// possible size of the decrypted data, allocates that much memory, and
// returns a pointer to it.
static uint8_t* default_in_msg_get_buf(NodeCommState &commst,
    uint32_t tot_enc_chunk_size)
{
    uint32_t max_plaintext_bytes = tot_enc_chunk_size;

    // If the handshake is complete, chunks will be encrypted and have a
    // MAC tag attached which will not correspond to plaintext bytes, so
    // we can trim them.

    if (commst.handshake_step == HANDSHAKE_COMPLETE) {
        // The minimum number of chunks needed to transmit this message
        uint32_t min_num_chunks =
            (tot_enc_chunk_size + (FRAME_SIZE-1)) / FRAME_SIZE;
        // The maximum number of plaintext bytes this message could contain
        max_plaintext_bytes = tot_enc_chunk_size -
            SGX_AESGCM_MAC_SIZE * min_num_chunks;
    }

    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 message of %u bytes from node %lu:\n",
        plaintext_len, nodest.node_num);
    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,
    uint8_t *data, uint32_t plaintext_len, uint32_t);
static void handshake_3_msg_received(NodeCommState &nodest,
    uint8_t *data, uint32_t plaintext_len, uint32_t);

// Receive (at the server) the first handshake message
static void handshake_1_msg_received(NodeCommState &nodest,
    uint8_t *data, uint32_t plaintext_len, uint32_t)
{
    /*
    printf("Received handshake_1 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)) {
        printf("Received handshake_1 message of incorrect size %u\n",
            plaintext_len);
        return;
    }
    sgx_ecc_state_handle_t ecc_handle;
    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_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;
    }

    printf("Valid public key received from node %hu\n", nodest.node_num);

    // Create our own DH key pair
    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_dh_privkey, &peer_dh_pubkey,
        &sharedsecret, ecc_handle);
    memset(&our_dh_privkey, 0, sizeof(our_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*)&our_dh_pubkey, sizeof(our_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(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*)&peer_dh_pubkey, sizeof(peer_dh_pubkey),
        hmac_handle);
    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);
    sgx_hmac256_update((uint8_t*)&g_pubkey, sizeof(g_pubkey),
        hmac_handle);
    sgx_hmac256_final(cli_srv_mac, 16, hmac_handle);
    sgx_hmac256_close(hmac_handle);

    // Sign the server-to-client MAC
    sgx_ec256_signature_t srv_cli_sig;
    sgx_ecdsa_sign(srv_cli_mac, 16, &g_privkey, &srv_cli_sig, ecc_handle);

    sgx_ecc256_close_context(ecc_handle);

    // Save the state we'll need to process handshake message 3
    memmove(&nodest.in_aes_key, h2, 16);
    memmove(&nodest.out_aes_key, ((uint8_t*)h2)+16, 16);
    memmove(&nodest.handshake_cli_srv_mac, cli_srv_mac, 16);

    // Get us ready to receive handshake message 3
    nodest.in_msg_get_buf = default_in_msg_get_buf;
    nodest.in_msg_received = handshake_3_msg_received;
    nodest.handshake_step = HANDSHAKE_S_SENT_2;

    // Send handshake message 2
    nodest.message_start(sizeof(our_dh_pubkey) + sizeof(srv_cli_sig),
        false);
    nodest.message_data((uint8_t*)&our_dh_pubkey, sizeof(our_dh_pubkey),
        false);
    nodest.message_data((uint8_t*)&srv_cli_sig, sizeof(srv_cli_sig),
        false);
}

// 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);
    memset(&nodest.out_aes_iv, 0, SGX_AESGCM_IV_SIZE);
    memset(&nodest.in_aes_iv, 0, SGX_AESGCM_IV_SIZE);
    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), false);
    nodest.message_data((uint8_t*)&cli_srv_sig, sizeof(cli_srv_sig),
        false);

    // Send a test message
    nodest.message_start(12);
    unsigned char buf[13];
    memmove(buf, "Hello, world", 13);
    nodest.message_data(buf, 12);
}

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");
    */

    if (plaintext_len != sizeof(sgx_ec256_signature_t)) {
        printf("Received handshake_3 message of incorrect size %u\n",
            plaintext_len);
        return;
    }
    sgx_ecc_state_handle_t ecc_handle;
    sgx_ec256_signature_t peer_sig;
    memmove(&peer_sig, data, sizeof(peer_sig));
    delete[] data;
    sgx_ecc256_open_context(&ecc_handle);

    // Verify the signature on the client-to-server MAC
    uint8_t result;
    if (sgx_ecdsa_verify(nodest.handshake_cli_srv_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);

    // Our side of the handshake is complete
    memset(&nodest.out_aes_iv, 0, SGX_AESGCM_IV_SIZE);
    memset(&nodest.in_aes_iv, 0, SGX_AESGCM_IV_SIZE);
    nodest.handshake_step = HANDSHAKE_COMPLETE;
    nodest.in_msg_get_buf = default_in_msg_get_buf;
    nodest.in_msg_received = default_in_msg_received;

    // Send a test message
    nodest.message_start(12);
    unsigned char buf[13];
    memmove(buf, "Hello, world", 13);
    nodest.message_data(buf, 12);
}

// Start a new outgoing message.  Pass the number of _plaintext_ bytes
// the message will be.
void NodeCommState::message_start(uint32_t plaintext_len, bool encrypt)
{
    uint32_t ciphertext_len = plaintext_len;

    // If the handshake is complete, add SGX_AESGCM_MAC_SIZE bytes for
    // every FRAME_SIZE-SGX_AESGCM_MAC_SIZE bytes of plaintext.
    if (encrypt) {
        uint32_t num_chunks = (plaintext_len +
            FRAME_SIZE - SGX_AESGCM_MAC_SIZE - 1) /
            (FRAME_SIZE - SGX_AESGCM_MAC_SIZE);
        ciphertext_len = plaintext_len +
            num_chunks * SGX_AESGCM_MAC_SIZE;
    }
    ocall_message(&frame, node_num, ciphertext_len);
    frame_offset = 0;
    msg_size = ciphertext_len;
    msg_frame_offset = 0;
    msg_plaintext_size = plaintext_len;
    msg_plaintext_processed = 0;
    if (plaintext_len < FRAME_SIZE - SGX_AESGCM_MAC_SIZE) {
        msg_plaintext_chunk_remain = plaintext_len;
    } else {
        msg_plaintext_chunk_remain = FRAME_SIZE - SGX_AESGCM_MAC_SIZE;
    }
    if (!frame) {
        printf("Received NULL back from ocall_message\n");
    }
    if (msg_plaintext_chunk_remain > 0) {
        if (encrypt) {
            *(size_t*)out_aes_iv += 1;
            sgx_aes_gcm128_enc_init(out_aes_key, out_aes_iv,
                SGX_AESGCM_IV_SIZE, NULL, 0, &out_aes_gcm_state);
        }
    }
}

// Process len bytes of plaintext data into the current message.
void NodeCommState::message_data(uint8_t *data, uint32_t len, bool encrypt)
{
    while (len > 0) {
        if (msg_plaintext_chunk_remain == 0) {
            printf("Attempt to queue too much message data\n");
            return;
        }
        uint32_t bytes_to_process = len;
        if (bytes_to_process > msg_plaintext_chunk_remain) {
            bytes_to_process = msg_plaintext_chunk_remain;
        }
        if (frame == NULL) {
            printf("frame is NULL when queueing message data\n");
            return;
        }
        if (encrypt) {
            // Encrypt the data
            sgx_aes_gcm128_enc_update(data, bytes_to_process,
                frame+frame_offset, out_aes_gcm_state);
        } else {
            // Just copy the plaintext data during the handshake
            memmove(frame+frame_offset, data, bytes_to_process);
        }
        frame_offset += bytes_to_process;
        msg_plaintext_processed += bytes_to_process;
        msg_plaintext_chunk_remain -= bytes_to_process;
        len -= bytes_to_process;
        data += bytes_to_process;
        if (msg_plaintext_chunk_remain == 0) {
            // Complete and send this chunk
            if (encrypt) {
                sgx_aes_gcm128_enc_get_mac(frame+frame_offset,
                    out_aes_gcm_state);
                frame_offset += SGX_AESGCM_MAC_SIZE;
            }
            uint8_t *nextframe = NULL;
            ocall_chunk(&nextframe, node_num, frame, frame_offset);
            frame = nextframe;
            msg_frame_offset += frame_offset;
            frame_offset = 0;
            msg_plaintext_chunk_remain =
                msg_plaintext_size - msg_plaintext_processed;
            if (msg_plaintext_chunk_remain >
                    FRAME_SIZE - SGX_AESGCM_MAC_SIZE) {
                msg_plaintext_chunk_remain =
                    FRAME_SIZE - SGX_AESGCM_MAC_SIZE;
            }
            if (encrypt) {
                sgx_aes_gcm_close(out_aes_gcm_state);
                if (msg_plaintext_chunk_remain > 0) {
                    *(size_t*)out_aes_iv += 1;
                    sgx_aes_gcm128_enc_init(out_aes_key, out_aes_iv,
                        SGX_AESGCM_IV_SIZE, NULL, 0, &out_aes_gcm_state);
                }
            }
        }
    }
}

// The communication states for all the nodes.  There's an entry for
// ourselves in here, but it is unused.
static std::vector<NodeCommState> commstates;
static nodenum_t tot_nodes, my_node_num;

// Generate a new identity signature key.  Output the public key and the
// sealed private key.  outsealedpriv must point to SEALEDPRIVKEY_SIZE =
// sizeof(sgx_sealed_data_t) + sizeof(sgx_ec256_private_t) + 18 bytes of
// memory.
void ecall_identity_key_new(sgx_ec256_public_t *outpub,
    sgx_sealed_data_t *outsealedpriv)
{
    sgx_ecc_state_handle_t ecc_handle;

    sgx_ecc256_open_context(&ecc_handle);

    sgx_ecc256_create_key_pair(&g_privkey, &g_pubkey, ecc_handle);
    memmove(outpub, &g_pubkey, sizeof(g_pubkey));

    sgx_ecc256_close_context(ecc_handle);

    sgx_seal_data(18, (const uint8_t*)"TEEMS Identity key",
        sizeof(g_privkey), (const uint8_t*)&g_privkey,
        SEALED_PRIVKEY_SIZE, outsealedpriv);
}

// Load an identity key from a sealed privkey.  Output the resulting
// public key.  insealedpriv must point to sizeof(sgx_sealed_data_t) +
// sizeof(sgx_ec256_private_t) bytes of memory.  Returns true for
// success, false for failure.
bool ecall_identity_key_load(sgx_ec256_public_t *outpub,
    const sgx_sealed_data_t *insealedpriv)
{
    sgx_ecc_state_handle_t ecc_handle;

    char aad[18];
    uint32_t aadsize = sizeof(aad);
    sgx_ec256_private_t privkey;
    uint32_t privkeysize = sizeof(privkey);
    sgx_status_t res = sgx_unseal_data(
        insealedpriv, (uint8_t*)aad, &aadsize,
        (uint8_t*)&privkey, &privkeysize);

    if (res || aadsize != sizeof(aad) || privkeysize != sizeof(privkey)
            || memcmp(aad, "TEEMS Identity key", sizeof(aad))) {
        return false;
    }

    sgx_ecc256_open_context(&ecc_handle);

    sgx_ec256_public_t pubkey;
    int valid;
    if (sgx_ecc256_calculate_pub_from_priv(&privkey, &pubkey) ||
            sgx_ecc256_check_point(&pubkey, ecc_handle, &valid) ||
            !valid) {
        sgx_ecc256_close_context(ecc_handle);
        return false;
    }

    sgx_ecc256_close_context(ecc_handle);

    memmove(&g_pubkey, &pubkey, sizeof(pubkey));
    memmove(&g_privkey, &privkey, sizeof(privkey));
    memmove(outpub, &pubkey, sizeof(pubkey));

    return true;
}

bool comms_init_nodestate(const EnclaveAPINodeConfig *apinodeconfigs,
    nodenum_t num_nodes, nodenum_t me)
{
    sgx_ecc_state_handle_t ecc_handle;
    sgx_ecc256_open_context(&ecc_handle);

    commstates.clear();
    tot_nodes = 0;
    commstates.reserve(num_nodes);
    for (nodenum_t i=0; i<num_nodes; ++i) {
        // Check that the pubkey is valid
        int valid;
        if (sgx_ecc256_check_point(&apinodeconfigs[i].pubkey,
                ecc_handle, &valid) ||
                !valid) {
            printf("Pubkey for node %hu invalid\n", i);
            commstates.clear();
            sgx_ecc256_close_context(ecc_handle);
            return false;
        }
        commstates.emplace_back(&apinodeconfigs[i].pubkey, i);
    }
    sgx_ecc256_close_context(ecc_handle);

    my_node_num = me;

    // Check that no one other than us has our pubkey (deals with
    // reflection attacks)
    for (nodenum_t i=0; i<num_nodes; ++i) {
        if (i == my_node_num) continue;
        if (!memcmp(&commstates[i].pubkey,
                &commstates[my_node_num].pubkey,
                sizeof(commstates[i].pubkey))) {
            printf("Pubkey %hu matches our own; possible reflection attack?\n",
                i);
            commstates.clear();
            return false;
        }
    }

    tot_nodes = num_nodes;

    // There will be an enclave-to-enclave channel between us and each
    // other node's enclave.  For the node numbers smaller than ours, we
    // will be the server for the handshake for that channel.  Prepare
    // to receive the first handshake message from those nodes'
    // enclaves.
    for (nodenum_t i=0; i<my_node_num; ++i) {
        commstates[i].in_msg_get_buf = default_in_msg_get_buf;
        commstates[i].in_msg_received = handshake_1_msg_received;
    }
    return true;
}

bool ecall_message(nodenum_t node_num, uint32_t message_len)
{
    if (node_num >= tot_nodes) {
        printf("Out-of-range node_num %hu received in ecall_message\n",
            node_num);
        return false;
    }
    NodeCommState &nodest = commstates[node_num];

    if (nodest.in_msg_size != nodest.in_msg_offset) {
        printf("Received ecall_message without completing previous message\n");
        return false;
    }
    if (!nodest.in_msg_get_buf) {
        printf("No message header handler registered\n");
        return false;
    }
    uint8_t *buf = nodest.in_msg_get_buf(nodest, message_len);
    if (!buf) {
        printf("Message header handler returned NULL\n");
        return false;
    }
    nodest.in_msg_size = message_len;
    nodest.in_msg_offset = 0;
    nodest.in_msg_plaintext_processed = 0;
    nodest.in_msg_buf = buf;
    return true;
}

bool ecall_chunk(nodenum_t node_num, const uint8_t *chunkdata,
    uint32_t chunklen)
{
    if (node_num >= tot_nodes) {
        printf("Out-of-range node_num %hu received in ecall_chunk\n",
            node_num);
        return false;
    }
    NodeCommState &nodest = commstates[node_num];

    if (nodest.in_msg_size == nodest.in_msg_offset) {
        printf("Received ecall_chunk after completing message\n");
        return false;
    }
    if (!nodest.in_msg_buf) {
        printf("No incoming message buffer allocated\n");
        return false;
    }
    if (!nodest.in_msg_received) {
        printf("No message received handler registered\n");
        return false;
    }
    if (nodest.in_msg_offset + chunklen > nodest.in_msg_size) {
        printf("Chunk larger than remaining message size\n");
        return false;
    }
    if (nodest.handshake_step == HANDSHAKE_COMPLETE) {
        // Decrypt the incoming data
        *(size_t*)(nodest.in_aes_iv) += 1;
        if (sgx_rijndael128GCM_decrypt(&nodest.in_aes_key, chunkdata,
                chunklen - SGX_AESGCM_MAC_SIZE,
                nodest.in_msg_buf + nodest.in_msg_plaintext_processed,
                nodest.in_aes_iv, SGX_AESGCM_IV_SIZE, NULL, 0,
                (const sgx_aes_gcm_128bit_tag_t *)
                    (chunkdata + chunklen - SGX_AESGCM_MAC_SIZE))) {
            printf("Decryption failed\n");
            return false;
        }
        nodest.in_msg_plaintext_processed +=
            chunklen - SGX_AESGCM_MAC_SIZE;
    } else {
        // Just copy the handshake data
        memmove(nodest.in_msg_buf + nodest.in_msg_plaintext_processed,
            chunkdata, chunklen);
        nodest.in_msg_plaintext_processed += chunklen;
    }
    nodest.in_msg_offset += chunklen;
    if (nodest.in_msg_offset == nodest.in_msg_size) {
        // This was the last chunk; handle the received message
        uint8_t* buf = nodest.in_msg_buf;
        uint32_t plaintext_processed = nodest.in_msg_plaintext_processed;
        uint32_t msg_size = nodest.in_msg_size;
        nodest.in_msg_buf = NULL;
        nodest.in_msg_size = 0;
        nodest.in_msg_offset = 0;
        nodest.in_msg_plaintext_processed = 0;
        nodest.in_msg_received(nodest, buf, plaintext_processed, msg_size);
    }
    return true;
}

// Start the handshake (as the client)
void NodeCommState::handshake_start()
{
    sgx_ecc_state_handle_t ecc_handle;

    sgx_ecc256_open_context(&ecc_handle);

    // Create a DH keypair
    sgx_ecc256_create_key_pair(&handshake_dh_privkey, &handshake_dh_pubkey,
        ecc_handle);

    sgx_ecc256_close_context(ecc_handle);

    // Get us ready to receive handshake message 2
    in_msg_get_buf = default_in_msg_get_buf;
    in_msg_received = handshake_2_msg_received;
    handshake_step = HANDSHAKE_C_SENT_1;

    // Send the public key as the first message
    message_start(sizeof(handshake_dh_pubkey), false);

    message_data((uint8_t*)&handshake_dh_pubkey,
        sizeof(handshake_dh_pubkey), false);
}

// Start all handshakes for which we are the client
bool ecall_comms_start()
{
    for (nodenum_t t = my_node_num+1; t<tot_nodes; ++t) {
        commstates[t].handshake_start();
    }
    return true;
}