#ifndef __COMMS_HPP__
#define __COMMS_HPP__

#include <vector>

#include "enclave_api.h"

// 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(const uint8_t *data, uint32_t len, bool encrypt=true);

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

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

// 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.
uint8_t* default_in_msg_get_buf(NodeCommState &commst,
    uint32_t tot_enc_chunk_size);

// An in_msg_received handler when we don't actually expect a message
// from a given node at a given time.
void unknown_in_msg_received(NodeCommState &nodest,
    uint8_t *data, uint32_t plaintext_len, uint32_t);

// The enclave-to-enclave communication protocol is as follows.  It
// probably could just be attested TLS in a production environment, but
// we're not implementing remote attestation at this time.  This means
// that the list of other enclaves' public keys are currently just
// blindly trusted, so add a remote attestation step to validate them if
// you want to deploy this for real.
//
// The protocol starts with a Sign-and-MAC (SIGMA) handshake, in the
// pre-specified peer setting.  The client is the lower-numbered node,
// and the server is the higher-numbered node.  The protocol is:
//
// Message 1 C -> S: g^x
// Message 2 S -> C: g^y, Sig_S(MAC_{H_1a(g^{xy})}(g^y, g^x, Pub_S, Pub_C)
// Message 3 C -> S: Sig_C(MAC_{H_1b(g^{xy})}(g^x, g^y, Pub_C, Pub_S)
//
// where Pub_C and Pub_S are the long-term signature keys of C and S.
//
// After the handshake, the client-to-server AES-GCM key is set to
// H_2a(g^{xy}) and the server-to-client AES-GCM key is set to
// H_2b(g^{xy}).  H_na(x) and H_nb(x) are the first 128 bits and the
// last 128 bits of SHA256(n || x) respectively.
//
// After the handshake, data is sent in logical messages, which are
// divided into chunks of size at most FRAME_SIZE - SGX_AESGCM_MAC_SIZE
// bytes of plaintext, which will expand to at most FRAME_SIZE bytes of
// ciphertext.  The IV for the first chunk in each direction is
// 0x01 0x00 0x00 ... 0x00 (remember they use different keys in the two
// directions), and each chunk increments the IV in a little-endian
// manner.  The MAC tag of SGX_AESGCM_MAC_SIZE bytes is at the end of
// the chunk.

bool comms_init_nodestate(const EnclaveAPINodeConfig *apinodeconfigs,
    nodenum_t num_nodes, nodenum_t my_node_num);

#endif