PRSONA/prsona/src/serverMain.cpp

#include <atomic>
#include <chrono>
#include <iostream>
#include <fstream>
#include <cstring>
#include <cstdlib>
#include <vector>
#include <string>

#include "networkServer.hpp"

using namespace std;

atomic<size_t> epochNum(0);

// Initialize the classes we use
void initialize_prsona_classes()
{
    Scalar::init();
    PrsonaBase::init();
    PrsonaBase::set_client_malicious();
}

PrsonaServer *create_server_from_bgn_file(
    size_t numServers,
    struct synchronization_tool *sync,
    const char *filename)
{
    unique_lock<mutex> lck(sync->mtx);
    ifstream bgnFile(filename);
    BGN privateKey;
    bgnFile >> privateKey;

    return new PrsonaServer(numServers, privateKey);
}

Twistpoint update_generator_from_gen_file(
    Proof& pi,
    struct synchronization_tool *sync,
    const char *filename)
{
    unique_lock<mutex> lck(sync->mtx);
    ifstream genFile(filename);
    Twistpoint retval;

    genFile >> pi;
    genFile >> retval;

    return retval;
}

Twistpoint update_data_from_epoch_gen_file(
    vector<Proof>& pi,
    struct synchronization_tool *sync,
    const char *filename)
{
    unique_lock<mutex> lck(sync->mtx);
    ifstream epochFile(filename);
    
    Twistpoint retval;
    BinarySizeT sizeOfVector;
    pi.clear();

    epochFile >> sizeOfVector;
    for (size_t i = 0; i < sizeOfVector.val(); i++)
    {
        Proof currProof;
        epochFile >> currProof;

        pi.push_back(currProof);
    }
    epochFile >> retval;

    return retval;
}

static int epoch_websocket_data_handler(
    struct mg_connection *conn,
    int bits,
    char *data,
    size_t data_len,
    void *user_data)
{
    if ((bits & 0xf) == MG_WEBSOCKET_OPCODE_CONNECTION_CLOSE  || (bits & 0xf) == MG_WEBSOCKET_OPCODE_DATACOMPLETE)
        return false;
    if ((bits & 0xf) == MG_WEBSOCKET_OPCODE_DATACOMPLETE)
    {
        struct synchronization_tool *sync = (struct synchronization_tool *) user_data;
        
        unique_lock<mutex> lck(sync->mtx);
        sync->val++;

        return false;
    }

    if ((bits & 0xf) != MG_WEBSOCKET_OPCODE_BINARY && (bits & 0xf) != MG_WEBSOCKET_OPCODE_CONTINUATION)
    {
        std::cerr << "Unknown opcode: failing." << std::endl;
        return false;
    }

    struct synchronization_tool *sync = (struct synchronization_tool *) user_data;
    char *filename = (char *) mg_get_user_connection_data(conn);

    unique_lock<mutex> lck(sync->mtx);
    FILE *currFile = fopen(filename, "ab");
    fwrite(data, sizeof(char), data_len, currFile);
    fclose(currFile);

    return true;
}

static void epoch_websocket_close_handler(
    const struct mg_connection *conn,
    void *user_data)
{
    struct synchronization_tool *sync = (struct synchronization_tool *) user_data;

    unique_lock<mutex> lck(sync->mtx);
    sync->val2 = 0;
    sync->cv.notify_all();
}

Twistpoint get_generator(
    default_random_engine& rng,
    vector<Proof>& pi,
    PrsonaServer *prsonaServer,
    const vector<string>& serverIPs,
    const vector<int>& serverPorts,
    const string& selfIP,
    int selfPort,
    bool fresh)
{
    Twistpoint retval = PrsonaServer::EL_GAMAL_GENERATOR;
    pi.clear();

    if (fresh)
        retval = prsonaServer->add_curr_seed_to_generator(pi, retval);
    else
        retval = prsonaServer->add_rand_seed_to_generator(pi, retval);

    const char* which = (fresh ? GET_FRESH_GEN_URI : GET_BLIND_GEN_URI);

    for (size_t i = 0; i < serverIPs.size(); i++)
    {
        if (serverIPs[i] == selfIP && serverPorts[i] == selfPort)
            continue;

        struct synchronization_tool generatorSync;
        char *genFilename;

        bool flag = false;
        while (!flag)
        {
            struct mg_connection *conn =
                mg_connect_websocket_client(
                    serverIPs[i].c_str(),
                    serverPorts[i],
                    USE_SSL,
                    NULL,
                    0,
                    which,
                    "null",
                    file_websocket_data_handler,
                    file_websocket_close_handler,
                    &generatorSync);

            if (!conn)
            {
                cerr << "Couldn't get server " << i << "'s update on generator" << endl;
                continue;
            }

            stringstream buffer;
            string data;
            buffer << retval;
            data = buffer.str();

            mg_websocket_client_write(
                conn,
                MG_WEBSOCKET_OPCODE_BINARY,
                data.c_str(),
                data.length());

            unique_lock<mutex> lck(generatorSync.mtx);
            genFilename = set_temp_filename(rng, conn);
            generatorSync.val = 0;
            mg_websocket_client_write(
                conn,
                MG_WEBSOCKET_OPCODE_DATACOMPLETE,
                "",
                0);

            while (!generatorSync.val)
                generatorSync.cv.wait(lck);

            mg_close_connection(conn);

            flag = true;
        }

        Proof currProof;
        retval = update_generator_from_gen_file(currProof, &generatorSync, genFilename);
        pi.push_back(currProof);

        remove(genFilename);
        delete [] genFilename;
    }

    return retval;
}

void handout_generator(
    const vector<Proof>& pi,
    const Twistpoint& generator,
    PrsonaServer *prsonaServer,
    const vector<string>& serverIPs,
    const vector<int>& serverPorts,
    const string& selfIP,
    int selfPort,
    bool fresh)
{
    if (fresh)
        prsonaServer->initialize_fresh_generator(pi, generator);
    else
        prsonaServer->set_EG_blind_generator(pi, generator);

    stringstream buffer;
    string data;

    const char* which = (fresh ? GIVE_FRESH_GEN_URI : GIVE_BLIND_GEN_URI);
    
    BinarySizeT sizeOfVector(pi.size());
    buffer << sizeOfVector;
    for (size_t i = 0; i < sizeOfVector.val(); i++)
        buffer << pi[i];

    buffer << generator;
    data = buffer.str();

    for (size_t i = 0; i < serverIPs.size(); i++)
    {
        if (serverIPs[i] == selfIP && serverPorts[i] == selfPort)
            continue;

        bool flag = false;
        while (!flag)
        {
            struct mg_connection *conn =
                mg_connect_websocket_client(
                    serverIPs[i].c_str(),
                    serverPorts[i],
                    USE_SSL,
                    NULL,
                    0,
                    which,
                    "null",
                    empty_websocket_data_handler,
                    empty_websocket_close_handler,
                    NULL);

            if (!conn)
            {
                cerr << "Couldn't give " << (fresh ? "fresh" : "blind") << " generator to server " << i << endl;
                continue;
            }

            mg_websocket_client_write(
                conn,
                MG_WEBSOCKET_OPCODE_BINARY,
                data.c_str(),
                data.length());

            mg_websocket_client_write(
                conn,
                MG_WEBSOCKET_OPCODE_DATACOMPLETE,
                "",
                0);

            mg_close_connection(conn);

            flag = true;
        }
    }
}

Twistpoint initiate_epoch_updates(
    default_random_engine& rng,
    const string& recipient,
    int recipientPort,
    const string& data,
    vector<vector<Proof>>& generatorProofHolder,
    bool isBreakdown)
{
    Twistpoint retval;
    struct synchronization_tool epochSync;
    char * epochFilename;
    const char* which = (isBreakdown ? EPOCH_BREAK_DOWN_URI : EPOCH_BUILD_UP_URI);

    bool flag = false;
    while (!flag)
    {
        struct mg_connection *conn =
            mg_connect_websocket_client(
                recipient.c_str(),
                recipientPort,
                USE_SSL,
                NULL,
                0,
                which,
                "null",
                epoch_websocket_data_handler,
                epoch_websocket_close_handler,
                &epochSync);

        if (!conn)
        {
            std::cerr << "Trouble initiating epoch update with server at " << recipient << std::endl;
            continue;
        }
        
        unique_lock<mutex> lck(epochSync.mtx);
        epochFilename = set_temp_filename(rng, conn);
        epochSync.val = 0;
        epochSync.val2 = 1;

        mg_websocket_client_write(
            conn,
            MG_WEBSOCKET_OPCODE_BINARY,
            data.c_str(),
            data.length());

        mg_websocket_client_write(
            conn,
            MG_WEBSOCKET_OPCODE_DATACOMPLETE,
            "",
            0);

        while (epochSync.val2)
            epochSync.cv.wait(lck);

        if (!epochSync.val)
            flag = true;

        mg_close_connection(conn);
    }

    if (isBreakdown)
        return retval;

    vector<Proof> generatorProof;
    generatorProofHolder.clear();

    retval = update_data_from_epoch_gen_file(generatorProof, &epochSync, epochFilename);

    generatorProofHolder.push_back(generatorProof);

    remove(epochFilename);
    delete [] epochFilename;

    return retval;
}

vector<Proof> epoch_build_up(
    PrsonaServer *prsonaServer,
    default_random_engine& rng,
    const vector<string>& serverIPs,
    const vector<int>& serverPorts,
    const string& selfIP,
    int selfPort,
    Twistpoint& nextGenerator)
{
    std::vector<std::vector<std::vector<Proof>>> pi;
    std::vector<std::vector<std::vector<Twistpoint>>> permutationCommits;
    std::vector<std::vector<std::vector<Twistpoint>>> freshPseudonymCommits;
    std::vector<std::vector<std::vector<Twistpoint>>> freshPseudonymSeedCommits;
    std::vector<std::vector<std::vector<CurveBipoint>>> serverTallyCommits;
    std::vector<std::vector<std::vector<std::vector<TwistBipoint>>>> partwayVoteMatrixCommits;
    std::vector<std::vector<std::vector<std::vector<TwistBipoint>>>> finalVoteMatrixCommits;

    std::vector<std::vector<Proof>> generatorProofHolder(1);

    for (size_t i = 0; i < serverIPs.size(); i++)
    {
        if (serverIPs[i] == selfIP && serverPorts[i] == selfPort)
        {
            pi.clear();
            permutationCommits.clear();
            freshPseudonymCommits.clear();
            freshPseudonymSeedCommits.clear();
            serverTallyCommits.clear();
            partwayVoteMatrixCommits.clear();
            finalVoteMatrixCommits.clear();

            pi.push_back(generatorProofHolder);

            prsonaServer->build_up_midway_pseudonyms(
                pi,
                permutationCommits,
                freshPseudonymCommits,
                freshPseudonymSeedCommits,
                serverTallyCommits,
                partwayVoteMatrixCommits,
                finalVoteMatrixCommits,
                nextGenerator);

            vector<vector<Twistpoint>> currUserTallyMaskCommits;
            vector<vector<Twistpoint>> currUserTallyMessageCommits;
            vector<vector<Twistpoint>> currUserTallySeedCommits;

            string data = make_epoch_update_string(
                pi[1],
                permutationCommits[0],
                freshPseudonymCommits[0],
                freshPseudonymSeedCommits[0],
                serverTallyCommits[0],
                partwayVoteMatrixCommits[0],
                finalVoteMatrixCommits[0],
                currUserTallyMaskCommits,
                currUserTallyMessageCommits,
                currUserTallySeedCommits,
                nextGenerator,
                false);

            struct synchronization_tool epochSync;
            vector<struct mg_connection *> conns;
            unique_lock<mutex> lck(epochSync.mtx);
            epochSync.val = 1;
            for (size_t j = 0; j < serverIPs.size(); j++)
            {
                if (i == j)
                    continue;

                struct mg_connection *currConn = 
                    distribute_epoch_updates(
                        serverIPs[j],
                        serverPorts[j],
                        data,
                        &epochSync);

                conns.push_back(currConn);
            }

            while (epochSync.val < serverIPs.size())
                epochSync.cv.wait(lck);

            for (size_t j = 0; j < conns.size(); j++)
                mg_close_connection(conns[j]);

            generatorProofHolder = pi[0];
        }
        else
        {
            string data = make_epoch_initiator_string(
                generatorProofHolder[0],
                nextGenerator);
            nextGenerator = initiate_epoch_updates(
                rng,
                serverIPs[i],
                serverPorts[i],
                data,
                generatorProofHolder,
                false);
        }
    }

    return generatorProofHolder[0];
}

void epoch_break_down(
    PrsonaServer *prsonaServer,
    default_random_engine& rng,
    const vector<string>& serverIPs,
    const vector<int>& serverPorts,
    const string& selfIP,
    int selfPort,
    const vector<Proof>& generatorProof,
    const Twistpoint& nextGenerator)
{
    std::vector<std::vector<std::vector<Proof>>> pi;
    std::vector<std::vector<std::vector<Twistpoint>>> permutationCommits;
    std::vector<std::vector<std::vector<Twistpoint>>> freshPseudonymCommits;
    std::vector<std::vector<std::vector<Twistpoint>>> freshPseudonymSeedCommits;
    std::vector<std::vector<std::vector<CurveBipoint>>> serverTallyCommits;
    std::vector<std::vector<std::vector<std::vector<TwistBipoint>>>> partwayVoteMatrixCommits;
    std::vector<std::vector<std::vector<std::vector<TwistBipoint>>>> finalVoteMatrixCommits;
    std::vector<std::vector<std::vector<Twistpoint>>> userTallyMaskCommits;
    std::vector<std::vector<std::vector<Twistpoint>>> userTallyMessageCommits;
    std::vector<std::vector<std::vector<Twistpoint>>> userTallySeedCommits;

    for (size_t i = 0; i < serverIPs.size(); i++)
    {
        pi.clear();
        permutationCommits.clear();
        freshPseudonymCommits.clear();
        freshPseudonymSeedCommits.clear();
        serverTallyCommits.clear();
        partwayVoteMatrixCommits.clear();
        finalVoteMatrixCommits.clear();
        userTallyMaskCommits.clear();
        userTallyMessageCommits.clear();
        userTallySeedCommits.clear();

        if (serverIPs[i] == selfIP && serverPorts[i] == selfPort)
        {
            prsonaServer->break_down_midway_pseudonyms(
                generatorProof,
                pi,
                permutationCommits,
                freshPseudonymCommits,
                freshPseudonymSeedCommits,
                serverTallyCommits,
                partwayVoteMatrixCommits,
                finalVoteMatrixCommits,
                userTallyMaskCommits,
                userTallyMessageCommits,
                userTallySeedCommits,
                nextGenerator);

            string data = make_epoch_update_string(
                pi[0],
                permutationCommits[0],
                freshPseudonymCommits[0],
                freshPseudonymSeedCommits[0],
                serverTallyCommits[0],
                partwayVoteMatrixCommits[0],
                finalVoteMatrixCommits[0],
                userTallyMaskCommits[0],
                userTallyMessageCommits[0],
                userTallySeedCommits[0],
                nextGenerator,
                true);

            struct synchronization_tool epochSync;
            vector<struct mg_connection *> conns;
            epochSync.val = 1;
            for (size_t j = 0; j < serverIPs.size(); j++)
            {
                if (i == j)
                    continue;

                struct mg_connection *currConn = 
                    distribute_epoch_updates(
                        serverIPs[j],
                        serverPorts[j],
                        data,
                        &epochSync);

                conns.push_back(currConn);
            }

            unique_lock<mutex> lck(epochSync.mtx);
            while (epochSync.val < serverIPs.size())
                epochSync.cv.wait(lck);

            for (size_t j = 0; j < conns.size(); j++)
                mg_close_connection(conns[j]);
        }
        else
        {
            vector<vector<Proof>> unused;
            string data = make_epoch_initiator_string(
                generatorProof,
                nextGenerator);
            initiate_epoch_updates(
                rng,
                serverIPs[i],
                serverPorts[i],
                data,
                unused,
                true);
        }
    }
}

void tally_scores(
    PrsonaServer *prsonaServer,
    const vector<string>& serverIPs,
    const vector<int>& serverPorts,
    const string& selfIP,
    int selfPort,
    const Twistpoint& nextGenerator,
    std::vector<EGCiphertext>& userTallyScores,
    std::vector<CurveBipoint>& serverTallyScores)
{
    struct synchronization_tool tallySync;
    vector<struct mg_connection *> conns;

    unique_lock<mutex> lck(tallySync.mtx);
    tallySync.val = 1;
    for (size_t i = 0; i < serverIPs.size(); i++)
    {
        if (serverIPs[i] == selfIP && serverPorts[i] == selfPort)
            continue;
        
        struct mg_connection *currConn = NULL;
        while (!currConn)
        {
            currConn = mg_connect_websocket_client(
                serverIPs[i].c_str(),
                serverPorts[i],
                USE_SSL,
                NULL,
                0,
                GET_DECRYPTION_URI,
                "null",
                synchro_websocket_data_handler,
                synchro_websocket_close_handler,
                &tallySync);

            if (!currConn)
                cerr << "Trouble getting partial decryption from server at " << serverIPs[i] << ":" << serverPorts[i] << endl;
        }

        mg_websocket_client_write(
            currConn,
            MG_WEBSOCKET_OPCODE_DATACOMPLETE,
            "",
            0);

        conns.push_back(currConn);
    }

    while (tallySync.val < serverIPs.size())
        tallySync.cv.wait(lck);

    for (size_t i = 0; i < conns.size(); i++)
        mg_close_connection(conns[i]);

    std::vector<EGCiphertext> retval;
    std::vector<Twistpoint> currentPseudonyms = prsonaServer->get_current_pseudonyms();
    std::vector<Scalar> decryptedTalliedScores = prsonaServer->tally_scores();
    mpz_class maxScorePossibleThisRound =
        prsonaServer->get_max_possible_score().toInt() *
        PrsonaBase::get_max_allowed_vote();
    mpz_class topOfScoreRange =
        decryptedTalliedScores.size() * PrsonaBase::get_max_allowed_vote();

    userTallyScores.clear();
    serverTallyScores.clear();
    for (size_t i = 0; i < decryptedTalliedScores.size(); i++)
    {
        decryptedTalliedScores[i] =
            Scalar(
                (decryptedTalliedScores[i].toInt() * topOfScoreRange) /
                maxScorePossibleThisRound
            );

        EGCiphertext currCiphertext;
        userTallyScores.push_back(currCiphertext);
        CurveBipoint currServerScore;
        serverTallyScores.push_back(currServerScore);
        Scalar currMask;
        currMask.set_random();

        // Give the server the new weights,
        // to get passed around to the other servers
        prsonaServer->encrypt(
            serverTallyScores[i], decryptedTalliedScores[i]);

        userTallyScores[i].mask = currentPseudonyms[i] * currMask;
        userTallyScores[i].encryptedMessage =
            (nextGenerator * currMask) +
            (prsonaServer->get_blinding_generator() * decryptedTalliedScores[i]);
    }
}

void distribute_tallied_scores(
    PrsonaServer *prsonaServer,
    const vector<string>& serverIPs,
    const vector<int>& serverPorts,
    const string& selfIP,
    int selfPort,
    const Twistpoint& nextGenerator,
    const std::vector<EGCiphertext>& userTallyScores,
    const std::vector<CurveBipoint>& serverTallyScores)
{
    stringstream buffer;
    string data;
    
    BinarySizeT sizeOfVector(userTallyScores.size());
    buffer << sizeOfVector;
    for (size_t i = 0; i < sizeOfVector.val(); i++)
        buffer << userTallyScores[i];
    for (size_t i = 0; i < sizeOfVector.val(); i++)
        buffer << serverTallyScores[i];

    data = buffer.str();

    struct synchronization_tool tallySync;
    vector<struct mg_connection *> conns;

    unique_lock<mutex> lck(tallySync.mtx);
    prsonaServer->receive_tallied_scores(userTallyScores, serverTallyScores);
    tallySync.val = 1;
    for (size_t i = 0; i < serverIPs.size(); i++)
    {
        if (serverIPs[i] == selfIP && serverPorts[i] == selfPort)
            continue;
        
        struct mg_connection *currConn = NULL;
        while (!currConn)
        {
            currConn = mg_connect_websocket_client(
                serverIPs[i].c_str(),
                serverPorts[i],
                USE_SSL,
                NULL,
                0,
                GIVE_DECRYPTION_URI,
                "null",
                synchro_websocket_data_handler,
                synchro_websocket_close_handler,
                &tallySync);

            if (!currConn)
                cerr << "Trouble giving full re-encryption to server at " << serverIPs[i] << ":" << serverPorts[i] << endl;
        }

        mg_websocket_client_write(
            currConn,
            MG_WEBSOCKET_OPCODE_BINARY,
            data.c_str(),
            data.length());

        mg_websocket_client_write(
            currConn,
            MG_WEBSOCKET_OPCODE_DATACOMPLETE,
            "",
            0);

        conns.push_back(currConn);
    }

    while (tallySync.val < serverIPs.size())
        tallySync.cv.wait(lck);

    for (size_t i = 0; i < conns.size(); i++)
        mg_close_connection(conns[i]);
}

void epoch(
    mutex *updateMtx,
    atomic<size_t> *epochNum,
    PrsonaServer *prsonaServer,
    default_random_engine& rng,
    const vector<string>& serverIPs,
    const vector<int>& serverPorts,
    const string& selfIP,
    int selfPort)
{
    Twistpoint nextGenerator = PrsonaServer::EL_GAMAL_GENERATOR;

    struct synchronization_tool updateSync;

    unique_lock<mutex> lck(*updateMtx, defer_lock);
    obtain_update_locks(
        lck,
        serverIPs,
        serverPorts,
        selfIP,
        selfPort,
        &updateSync);

    vector<Proof> generatorProof =
        epoch_build_up(
            prsonaServer,
            rng,
            serverIPs,
            serverPorts,
            selfIP,
            selfPort,
            nextGenerator);

    std::vector<EGCiphertext> currentUserEncryptedTallies;
    std::vector<CurveBipoint> currentServerEncryptedTallies;

    tally_scores(
        prsonaServer,
        serverIPs,
        serverPorts,
        selfIP,
        selfPort,
        nextGenerator,
        currentUserEncryptedTallies,
        currentServerEncryptedTallies);

    distribute_tallied_scores(
        prsonaServer,
        serverIPs,
        serverPorts,
        selfIP,
        selfPort,
        nextGenerator,
        currentUserEncryptedTallies,
        currentServerEncryptedTallies);

    epoch_break_down(
        prsonaServer,
        rng,
        serverIPs,
        serverPorts,
        selfIP,
        selfPort,
        generatorProof,
        nextGenerator);

    epochNum->fetch_add(1);

    release_update_locks(
        lck,
        serverIPs,
        serverPorts,
        selfIP,
        selfPort,
        &updateSync);
}

class EpochReadyHandler : public CivetHandler
{
    public:
        EpochReadyHandler(struct synchronization_tool *exitSync, struct synchronization_tool *readySync, mutex *updateMtx, size_t numServers)
            : exitSync(exitSync), readySync(readySync), updateMtx(updateMtx), numServers(numServers)
        { /* */ }

        bool handleGet(CivetServer *server, struct mg_connection *conn)
        {
            unique_lock<mutex> exitLock(exitSync->mtx, defer_lock);
            unique_lock<mutex> updateLock(*updateMtx, defer_lock);
            unique_lock<mutex> readyLock(readySync->mtx);

            if (readySync->val < numServers)
            {
                mg_printf(conn,
                          "HTTP/1.1 503 Service Unavailable\r\nContent-Type: "
                          "text/plain\r\nConnection: close\r\n\r\n");
                mg_printf(conn, "Server is waiting for other servers to begin.\n");
            }
            else if (!exitLock.try_lock())
            {
                mg_printf(conn,
                          "HTTP/1.1 503 Service Unavailable\r\nContent-Type: "
                          "text/plain\r\nConnection: close\r\n\r\n");
                mg_printf(conn, "Server is still in a previous epoch.\n");   
            }
            else if (!updateLock.try_lock())
            {
                mg_printf(conn,
                          "HTTP/1.1 503 Service Unavailable\r\nContent-Type: "
                          "text/plain\r\nConnection: close\r\n\r\n");
                mg_printf(conn, "Server is handling other updates.\n");   
            }
            else
            {
                mg_printf(conn,
                          "HTTP/1.1 200 OK\r\nContent-Type: "
                          "text/plain\r\nConnection: close\r\n\r\n");
                mg_printf(conn, "Server is ready for epoch.\n");
            }

            return true;
        }

    private:
        struct synchronization_tool *exitSync, *readySync;
        mutex *updateMtx;
        const size_t numServers;

};

class EpochNumHandler : public CivetHandler
{
    public:
        EpochNumHandler(atomic<size_t> *epochNum)
            : epochNum(epochNum)
        { /* */ }

        bool handleGet(CivetServer *server, struct mg_connection *conn)
        {
            mg_printf(conn,
                      "HTTP/1.1 200 OK\r\nContent-Type: "
                      "text/plain\r\nConnection: close\r\n\r\n");
            mg_printf(conn, "Epoch num: %lu\n", epochNum->load());

            return true;
        }

    private:
        atomic<size_t> *epochNum;
};

class UpdateLockWebSocketHandler : public CivetWebSocketHandler
{
    public:
        UpdateLockWebSocketHandler(mutex *updateMtx, unique_lock<mutex> **lockHolder, bool isLocking)
            : updateMtx(updateMtx), lockHolder(lockHolder), isLocking(isLocking)
        { /* */ }

        ~UpdateLockWebSocketHandler()
        { delete *lockHolder; }

        bool handleConnection(CivetServer *server, const struct mg_connection *conn)
        { return true; }

        void handleReadyState(CivetServer *server, struct mg_connection *conn)
        { /* */ }

        bool handleData(CivetServer *server, struct mg_connection *conn, int bits, char *data, size_t data_len)
        {
            switch (bits & 0xf)
            {
                case MG_WEBSOCKET_OPCODE_DATACOMPLETE:
                    if (isLocking)
                    {
                        unique_lock<mutex> *tempHolder = new unique_lock<mutex>(*updateMtx);
                        
                        // Once you get to this line, we now hold the lock,
                        // and lockHolder is guaranteed to be NULL
                        *lockHolder = tempHolder;

                        // Respond to notify that the requesting process holds the lock
                        mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_DATACOMPLETE, "", 0);
                    }
                    else
                    {
                        // You must do things in this order so that *lockHolder will be
                        // guaranteed to be NULL at the time the lock unlocks
                        // (deletion of the lock object)
                        unique_lock<mutex> *tempHolder = *lockHolder;
                        *lockHolder = NULL;

                        delete tempHolder;
                        
                        // Respond to notify that the requesting process has released the lock
                        mg_websocket_write(conn, MG_WEBSOCKET_OPCODE_DATACOMPLETE, "", 0);
                    }
                    break;

                case MG_WEBSOCKET_OPCODE_CONNECTION_CLOSE:
                    break;

                default:
                    cerr << "Unknown opcode: failing." << endl;
                    break;
            }

            return false;
        }

        void handleClose(CivetServer *server, const struct mg_connection *conn)
        { /* */ }


    private:
        mutex *updateMtx;
        unique_lock<mutex> **lockHolder;
        const bool isLocking;
};

int main(int argc, char *argv[])
{
    initialize_prsona_classes();

#if USE_SSL
    mg_init_library(0);
#else
    mg_init_library(MG_FEATURES_SSL);
#endif

    string id = "";
    if (argc > 1)
        id = argv[1];

    string seedStr;
    if (id.empty())
        seedStr = "default-server";
    else
    {
        seedStr = id;
        seedStr += "-server";
    }

    vector<string> serverIPs;
    vector<int> serverPorts;
    string selfIP, selfPortStr, dealerIP;
    int selfPort = 0, dealerPort = 0;

    char buffer[46], *helper;
    ifstream serverConfig("cfg/serverIPs.cfg");
    while (!serverConfig.eof())
    {
        serverConfig.getline(buffer, 46);
        if (strlen(buffer) > 0)
        {
            helper = buffer;
            if (strchr(helper, ':'))
            {
                helper = strtok(helper, ":");
                serverIPs.push_back(string(helper));
                helper = strtok(NULL, ":");
                serverPorts.push_back(atoi(helper));
            }
            else
            {
                serverIPs.push_back(string(helper));
                serverPorts.push_back(atoi(PRSONA_PORT_STR));
            }
        }
    }

    string selfConfigFilename = "cfg/selfIP";
    if (!id.empty())
    {
        selfConfigFilename += "-";
        selfConfigFilename += id;
    }
    selfConfigFilename += ".cfg";

    ifstream selfConfig(selfConfigFilename);
    while (!selfConfig.eof())
    {
        selfConfig.getline(buffer, 46);
        if (strlen(buffer) > 0)
        {
            helper = buffer;
            if (strchr(helper, ':'))
            {
                helper = strtok(helper, ":");
                selfIP = helper;
                helper = strtok(NULL, ":");
                selfPortStr = helper;
                selfPort = atoi(helper);
            }
            else
            {
                selfIP = helper;
                selfPortStr = PRSONA_PORT_STR;
                selfPort = atoi(PRSONA_PORT_STR);
            }
        }
    }

    ifstream dealerConfig("cfg/dealerIP.cfg");
    while (!dealerConfig.eof())
    {
        dealerConfig.getline(buffer, 46);
        if (strlen(buffer) > 0)
        {
            helper = buffer;
            if (strchr(helper, ':'))
            {
                helper = strtok(helper, ":");
                dealerIP = helper;
                helper = strtok(NULL, ":");
                dealerPort = atoi(helper);
            }
            else
            {
                dealerIP = helper;
                dealerPort = atoi(PRSONA_PORT_STR);
            }
        }
    }

    // Defaults
    size_t numServers = serverIPs.size();
    bool bgnDealer = selfIP == dealerIP && selfPort == dealerPort;
    bool maliciousServers = true;

    const char *options[] = {"listening_ports", selfPortStr.c_str(), 0};

    if (argc > 2)
    {
        bool setting = argv[2][0] == 't' || argv[2][0] == 'T';
        maliciousServers = setting;
    }

    seed_seq seed(seedStr.begin(), seedStr.end());
    default_random_engine rng(seed);

    cout << "[" << seedStr << "] Establishing PRSONA server with the following parameters: " << endl;
    cout << "[" << seedStr << "] " << numServers << " PRSONA servers" << endl;
    cout << "[" << seedStr << "] This server " << (bgnDealer ? "IS" : "is NOT") << " the trusted BGN dealer" << endl;
    cout << "[" << seedStr << "] Servers are set to " << (maliciousServers ? "MALICIOUS" : "HBC") << " security" << endl;
    cout << "[" << seedStr << "] This server is at IP address: " << selfIP << ":" << selfPort << endl;
    cout << "[" << seedStr << "] The BGN dealer is at IP address: " << dealerIP << ":" << dealerPort << endl;
    cout << endl;

    // Set malicious flags where necessary
    if (maliciousServers)
        PrsonaBase::set_server_malicious();

    struct synchronization_tool exitSync, readySync;
    mutex updateMtx;
    unique_lock<mutex> *updateLockHolder;
    atomic<size_t> epochNum(0);

    cout << "[" << seedStr << "] Creating PrsonaServer entity." << endl;

    // Entities we operate with
    PrsonaServer *prsonaServer;
    if (bgnDealer)
        prsonaServer = new PrsonaServer(numServers);
    else
    {
        cout << "[" << seedStr << "] Retrieving BGN details." << endl;

        struct synchronization_tool bgnSync;
        char *bgnFilename;

        bool flag = false;
        while (!flag)
        {
            struct mg_connection *conn =
                mg_connect_websocket_client(
                    dealerIP.c_str(),
                    dealerPort,
                    USE_SSL,
                    NULL,
                    0,
                    PRIVATE_BGN_URI,
                    "null",
                    file_websocket_data_handler,
                    file_websocket_close_handler,
                    &bgnSync);

            if (!conn)
            {
                cerr << "[" << seedStr << "] Couldn't obtain BGN details." << endl;
                continue;
            }

            unique_lock<mutex> lck(bgnSync.mtx);
            bgnFilename = set_temp_filename(rng, conn);
            bgnSync.val = 0;
            mg_websocket_client_write(
                conn,
                MG_WEBSOCKET_OPCODE_DATACOMPLETE,
                "",
                0);

            while (!bgnSync.val)
                bgnSync.cv.wait(lck);

            mg_close_connection(conn);

            flag = true;
        }

        prsonaServer = create_server_from_bgn_file(numServers, &bgnSync, bgnFilename);

        remove(bgnFilename);
        delete [] bgnFilename;
    }

    CivetServer server(options);

    PrsonaServerWebSocketHandler wsHandler(prsonaServer, &rng, &updateMtx, &epochNum, serverIPs, serverPorts, selfIP, selfPort);
    server.addWebSocketHandler("/ws", wsHandler);

    if (bgnDealer)
    {
        cout << "[" << seedStr << "] Waiting for other servers to check in and retrieve BGN details." << endl;

        unique_lock<mutex> lck(readySync.mtx);
        RemoteControlHandler serverReadyHandler(&readySync, "ACK");
        server.addHandler(SERVER_READY_URI, serverReadyHandler);

        readySync.val++;
        while (readySync.val < numServers)
            readySync.cv.wait(lck);

        vector<Proof> pi;
        Twistpoint freshGenerator =
            get_generator(rng, pi, prsonaServer, serverIPs, serverPorts, selfIP, selfPort, true);

        handout_generator(pi, freshGenerator, prsonaServer, serverIPs, serverPorts, selfIP, selfPort, true);

        Twistpoint blindGenerator =
            get_generator(rng, pi, prsonaServer, serverIPs, serverPorts, selfIP, selfPort, false);

        handout_generator(pi, blindGenerator, prsonaServer, serverIPs, serverPorts, selfIP, selfPort, false);
    }
    else
    {
        cout << "[" << seedStr << "] Notifying BGN dealer that this server is ready." << endl;

        stringstream sysString;
        string data;

        struct mg_connection *conn =
            mg_connect_client(
                dealerIP.c_str(), 
                dealerPort,
                USE_SSL, 
                NULL,
                0);

        sysString << "GET " << SERVER_READY_URI << " HTTP/1.1\r\n";
        sysString << "Host: " << dealerIP << ":" << dealerPort << "\r\n\r\n";
        data = sysString.str();

        mg_write(conn, data.c_str(), data.length());
        mg_close_connection(conn);
    }

    unique_lock<mutex> exitLock(exitSync.mtx);
    exitSync.val = 0;
    exitSync.val2 = 0;
    RemoteControlHandler exitHandler(&exitSync, "Server coming down!");
    server.addHandler(EXIT_URI, exitHandler);

    UpdateLockWebSocketHandler lockHandler(&updateMtx, &updateLockHolder, true);
    UpdateLockWebSocketHandler unlockHandler(&updateMtx, &updateLockHolder, false);

    server.addWebSocketHandler(UPDATE_LOCK_URI, lockHandler);
    server.addWebSocketHandler(UPDATE_UNLOCK_URI, unlockHandler);

    cout << "[" << seedStr << "] Entering main ready loop." << endl;

    if (bgnDealer)
    {
        AltRemoteControlHandler triggerEpochHandler(1, &exitSync, "Server will initiate epoch!");
        server.addHandler(TRIGGER_EPOCH_URI, triggerEpochHandler);

        EpochReadyHandler epochReadyHandler(&exitSync, &readySync, &updateMtx, numServers);
        server.addHandler(EPOCH_READY_URI, epochReadyHandler);

        EpochNumHandler epochNumHandler(&epochNum);
        server.addHandler(WHICH_EPOCH_URI, epochNumHandler);

        while (!exitSync.val)
        {
            while (!exitSync.val && !exitSync.val2)
                exitSync.cv.wait(exitLock);

            if (exitSync.val2)
            {
                size_t currEpoch = epochNum.load();
                cout << "[" << seedStr << "] Executing epoch calculations (going from t = " << currEpoch << " to " << currEpoch + 1 << ")." << endl;

                epoch(&updateMtx, &epochNum, prsonaServer, rng, serverIPs, serverPorts, selfIP, selfPort);

                currEpoch = epochNum.load();
                cout << "[" << seedStr << "] Epoch calculations complete (now in t = " << currEpoch << ")." << endl;

                exitSync.val2 = 0;
            }
        }
    }
    else
    {
        while (!exitSync.val)
            exitSync.cv.wait(exitLock);
    }

    cout << "[" << seedStr << "] Shutting down." << endl;

    mg_exit_library();
    delete prsonaServer;

    return 0;
}