Implement publickey whitelist (using libsodium)

Renamed "mutex" to "conMapMutex" since it is mainly used to lock access
to the connection map.

Removed UDPC_client_initiate_connection_pk() as publickey whitelisting
replaces its functionality.

src/UDPC_Defines.hpp

@@ -82,6 +82,16 @@ struct IPV6_Hasher {
     std::size_t operator()(const UDPC_IPV6_ADDR_TYPE& addr) const;
 };
 
+struct PKContainer {
+    PKContainer();
+    PKContainer(unsigned char *pk);
+
+    unsigned char pk[crypto_sign_PUBLICKEYBYTES];
+
+    std::size_t operator()(const PKContainer& container) const;
+    bool operator==(const PKContainer& other) const;
+};
+
 struct ConnectionData {
     ConnectionData(bool isUsingLibsodium);
     ConnectionData(
@@ -112,7 +122,6 @@ struct ConnectionData {
      * 4 - is id set
      * 5 - error initializing keys for public key encryption
      * 6 - using libsodium for header verification
-     * 7 - peer_pk pre-set
      */
     std::bitset<8> flags;
     uint32_t id;
@@ -225,6 +234,7 @@ public:
     // id to ipv6 address and port (as UDPC_ConnectionId)
     std::unordered_map<uint32_t, UDPC_ConnectionId> idMap;
     std::unordered_set<UDPC_ConnectionId, ConnectionIdHasher> deletionMap;
+    std::unordered_set<PKContainer, PKContainer> peerPKWhitelist;
     TSLQueue<UDPC_PacketInfo> receivedPkts;
     TSLQueue<UDPC_PacketInfo> cSendPkts;
     // handled internally
@@ -236,7 +246,8 @@ public:
 
     std::thread thread;
     std::atomic_bool threadRunning;
-    std::mutex mutex;
+    std::mutex conMapMutex;
+    std::mutex peerPKWhitelistMutex;
 
     std::chrono::milliseconds threadedSleepTime;
     unsigned char sk[crypto_sign_SECRETKEYBYTES];

src/UDPConnection.cpp

@@ -65,6 +65,24 @@ std::size_t UDPC::IPV6_Hasher::operator()(const UDPC_IPV6_ADDR_TYPE& addr) const
     return std::hash<std::string>()(std::string((const char*)UDPC_IPV6_ADDR_SUB(addr), 16));
 }
 
+UDPC::PKContainer::PKContainer() {
+    std::memset(pk, 0, crypto_sign_PUBLICKEYBYTES);
+}
+
+UDPC::PKContainer::PKContainer(unsigned char *pk) {
+    std::memcpy(this->pk, pk, crypto_sign_PUBLICKEYBYTES);
+}
+
+std::size_t UDPC::PKContainer::operator()(const PKContainer& container) const {
+    return std::hash<std::string>()(std::string(
+        (const char*)container.pk,
+        crypto_sign_PUBLICKEYBYTES));
+}
+
+bool UDPC::PKContainer::operator==(const PKContainer& other) const {
+    return std::memcmp(pk, other.pk, crypto_sign_PUBLICKEYBYTES) == 0;
+}
+
 bool operator ==(const UDPC_ConnectionId& a, const UDPC_ConnectionId& b) {
     return a.addr == b.addr && a.scope_id == b.scope_id && a.port == b.port;
 }
@@ -208,7 +226,7 @@ atostrBufIndex(0),
 receivedPkts(),
 cSendPkts(),
 rng_engine(),
-mutex()
+conMapMutex()
 {
     for(unsigned int i = 0; i < UDPC_ATOSTR_SIZE; ++i) {
         atostrBuf[i] = 0;
@@ -317,103 +335,6 @@ void UDPC::Context::update_impl() {
                     std::memcpy(newCon.verifyMessage.get(), &timeInt, 8);
                 }
 
-                if(conMap.find(optE->conId) == conMap.end()) {
-                    conMap.insert(std::make_pair(
-                        optE->conId,
-                        std::move(newCon)));
-                    auto addrConIter = addrConMap.find(optE->conId.addr);
-                    if(addrConIter == addrConMap.end()) {
-                        auto insertResult = addrConMap.insert(std::make_pair(
-                            optE->conId.addr,
-                            std::unordered_set<UDPC_ConnectionId, UDPC::ConnectionIdHasher>{}));
-                        assert(insertResult.second &&
-                            "new connection insert into addrConMap must not fail");
-                        addrConIter = insertResult.first;
-                    }
-                    addrConIter->second.insert(optE->conId);
-                    UDPC_CHECK_LOG(this,
-                        UDPC_LoggingType::UDPC_INFO,
-                        "Client initiating connection to ",
-                        UDPC_atostr((UDPC_HContext)this, optE->conId.addr),
-                        " port ",
-                        optE->conId.port,
-                        " ...");
-                } else {
-                    UDPC_CHECK_LOG(this,
-                        UDPC_LoggingType::UDPC_WARNING,
-                        "Client initiate connection, already connected to peer ",
-                        UDPC_atostr((UDPC_HContext)this, optE->conId.addr),
-                        " port ",
-                        optE->conId.port);
-                }
-            }
-                break;
-            case UDPC_ET_REQUEST_CONNECT_PK:
-            {
-                assert(flags.test(2) &&
-                    "libsodium should be explictly enabled");
-                unsigned char *sk = nullptr;
-                unsigned char *pk = nullptr;
-                if(keysSet.load()) {
-                    sk = this->sk;
-                    pk = this->pk;
-                }
-                UDPC::ConnectionData newCon(
-                    false,
-                    this,
-                    optE->conId.addr,
-                    optE->conId.scope_id,
-                    optE->conId.port,
-#ifdef UDPC_LIBSODIUM_ENABLED
-                    true,
-                    sk, pk);
-#else
-                    false,
-                    sk, pk);
-                assert(!"compiled without libsodium support");
-                delete[] optE->v.pk;
-                break;
-#endif
-                if(newCon.flags.test(5)) {
-                    delete[] optE->v.pk;
-                    UDPC_CHECK_LOG(this,
-                        UDPC_LoggingType::UDPC_ERROR,
-                        "Failed to init ConnectionData instance (libsodium "
-                        "init fail) while client establishing connection with ",
-                        UDPC_atostr((UDPC_HContext)this, optE->conId.addr),
-                        " port ",
-                        optE->conId.port);
-                    continue;
-                }
-                newCon.sent = std::chrono::steady_clock::now() - UDPC::INIT_PKT_INTERVAL_DT;
-                if(flags.test(2) && newCon.flags.test(6)) {
-                    // set up verification string to send to server
-                    std::time_t time = std::time(nullptr);
-                    if(time <= 0) {
-                        UDPC_CHECK_LOG(this, UDPC_LoggingType::UDPC_ERROR,
-                            "Failed to get current epoch time");
-                        continue;
-                    }
-                    uint64_t timeInt = time;
-#ifndef NDEBUG
-                    UDPC_CHECK_LOG(this, UDPC_LoggingType::UDPC_DEBUG,
-                        "Client set up verification epoch time \"",
-                        timeInt, "\"");
-#endif
-                    UDPC::be64((char*)&timeInt);
-                    newCon.verifyMessage = std::unique_ptr<char[]>(new char[8]);
-                    std::memcpy(newCon.verifyMessage.get(), &timeInt, 8);
-
-                    // set peer public key
-                    std::memcpy(
-                        newCon.peer_pk,
-                        optE->v.pk,
-                        crypto_sign_PUBLICKEYBYTES);
-                    newCon.flags.set(7);
-                }
-
-                delete[] optE->v.pk;
-
                 if(conMap.find(optE->conId) == conMap.end()) {
                     conMap.insert(std::make_pair(
                         optE->conId,
@@ -1303,6 +1224,15 @@ void UDPC::Context::update_impl() {
                     newConnection.peer_pk,
                     recvBuf + UDPC_MIN_HEADER_SIZE + 4,
                     crypto_sign_PUBLICKEYBYTES);
+                {
+                    std::lock_guard<std::mutex> lock(peerPKWhitelistMutex);
+                    if(!peerPKWhitelist.empty() && peerPKWhitelist.find(UDPC::PKContainer(newConnection.peer_pk)) == peerPKWhitelist.end()) {
+                        UDPC_CHECK_LOG(this, UDPC_LoggingType::UDPC_WARNING,
+                            "peer_pk is not in whitelist, not establishing "
+                            "connection with client");
+                        return;
+                    }
+                }
                 newConnection.verifyMessage = std::unique_ptr<char[]>(new char[crypto_sign_BYTES]);
                 std::time_t currentTime = std::time(nullptr);
                 uint64_t receivedTime;
@@ -1403,19 +1333,17 @@ void UDPC::Context::update_impl() {
 
             if(pktType == 2 && flags.test(2) && iter->second.flags.test(6)) {
 #ifdef UDPC_LIBSODIUM_ENABLED
-                if(iter->second.flags.test(7)) {
+                std::memcpy(iter->second.peer_pk,
-                    if(std::memcmp(iter->second.peer_pk,
+                    recvBuf + UDPC_MIN_HEADER_SIZE + 4,
-                            recvBuf + UDPC_MIN_HEADER_SIZE + 4,
+                    crypto_sign_PUBLICKEYBYTES);
-                            crypto_sign_PUBLICKEYBYTES) != 0) {
+                {
+                    std::lock_guard<std::mutex> lock(peerPKWhitelistMutex);
+                    if(!peerPKWhitelist.empty() && peerPKWhitelist.find(UDPC::PKContainer(iter->second.peer_pk)) == peerPKWhitelist.end()) {
                         UDPC_CHECK_LOG(this, UDPC_LoggingType::UDPC_WARNING,
-                            "peer_pk did not match pre-set peer_pk, not "
+                            "peer_pk is not in whitelist, not establishing "
-                            "establishing connection");
+                            "connection with server");
                         return;
                     }
-                } else {
-                    std::memcpy(iter->second.peer_pk,
-                        recvBuf + UDPC_MIN_HEADER_SIZE + 4,
-                        crypto_sign_PUBLICKEYBYTES);
                 }
                 if(crypto_sign_verify_detached(
                     (unsigned char*)(recvBuf + UDPC_MIN_HEADER_SIZE + 4 + crypto_sign_PUBLICKEYBYTES),
@@ -1849,7 +1777,7 @@ void UDPC::threadedUpdate(Context *ctx) {
     while(ctx->threadRunning.load()) {
         now = std::chrono::steady_clock::now();
         {
-            std::lock_guard<std::mutex> lock(ctx->mutex);
+            std::lock_guard<std::mutex> lock(ctx->conMapMutex);
             ctx->update_impl();
         }
         nextNow = std::chrono::steady_clock::now();
@@ -2122,12 +2050,6 @@ void UDPC_destroy(UDPC_HContext ctx) {
 #if UDPC_PLATFORM == UDPC_PLATFORM_WINDOWS
         WSACleanup();
 #endif
-        while(!UDPC_ctx->internalEvents.empty()) {
-            auto optE = UDPC_ctx->internalEvents.top_and_pop();
-            if(optE && optE->type == UDPC_ET_REQUEST_CONNECT_PK) {
-                delete[] optE->v.pk;
-            }
-        }
         UDPC_ctx->_contextIdentifier = 0;
         delete UDPC_ctx;
     }
@@ -2140,7 +2062,7 @@ void UDPC_update(UDPC_HContext ctx) {
         return;
     }
 
-    std::lock_guard<std::mutex> lock(c->mutex);
+    std::lock_guard<std::mutex> lock(c->conMapMutex);
     c->update_impl();
 }
 
@@ -2163,34 +2085,6 @@ void UDPC_client_initiate_connection(
     c->internalEvents.push(UDPC_Event{UDPC_ET_REQUEST_CONNECT, connectionId, enableLibSodium});
 }
 
-void UDPC_client_initiate_connection_pk(
-        UDPC_HContext ctx,
-        UDPC_ConnectionId connectionId,
-        unsigned char *serverPK) {
-    UDPC::Context *c = UDPC::verifyContext(ctx);
-    if(!c || !c->flags.test(1)) {
-        return;
-    }
-#ifndef UDPC_LIBSODIUM_ENABLED
-    UDPC_CHECK_LOG(c, UDPC_LoggingType::UDPC_ERROR,
-        "Cannot initiate connection with public key, UDPC was compiled "
-        "without libsodium");
-    return;
-#else
-    else if(!c->flags.test(2)) {
-        UDPC_CHECK_LOG(c, UDPC_LoggingType::UDPC_ERROR,
-            "Cannot initiate connection with public key, libsodium is not "
-            "enabled");
-        return;
-    }
-#endif
-
-    UDPC_Event event{UDPC_ET_REQUEST_CONNECT_PK, connectionId, 0};
-    event.v.pk = new unsigned char[crypto_sign_PUBLICKEYBYTES];
-    std::memcpy(event.v.pk, serverPK, crypto_sign_PUBLICKEYBYTES);
-    c->internalEvents.push(event);
-}
-
 void UDPC_queue_send(UDPC_HContext ctx, UDPC_ConnectionId destinationId,
                      int isChecked, void *data, uint32_t size) {
     if(size == 0 || !data) {
@@ -2229,7 +2123,7 @@ unsigned long UDPC_get_queued_size(UDPC_HContext ctx, UDPC_ConnectionId id, int
         return 0;
     }
 
-    std::lock_guard<std::mutex> lock(c->mutex);
+    std::lock_guard<std::mutex> lock(c->conMapMutex);
     auto iter = c->conMap.find(id);
     if(iter != c->conMap.end()) {
         if(exists) {
@@ -2272,7 +2166,7 @@ int UDPC_has_connection(UDPC_HContext ctx, UDPC_ConnectionId connectionId) {
         return 0;
     }
 
-    std::lock_guard<std::mutex> lock(c->mutex);
+    std::lock_guard<std::mutex> lock(c->conMapMutex);
 
     return c->conMap.find(connectionId) == c->conMap.end() ? 0 : 1;
 }
@@ -2283,7 +2177,7 @@ UDPC_ConnectionId* UDPC_get_list_connected(UDPC_HContext ctx, unsigned int *size
         return nullptr;
     }
 
-    std::lock_guard<std::mutex> lock(c->mutex);
+    std::lock_guard<std::mutex> lock(c->conMapMutex);
 
     if(c->conMap.empty()) {
         if(size) {
@@ -2396,7 +2290,7 @@ int UDPC_set_libsodium_keys(UDPC_HContext ctx, unsigned char *sk, unsigned char
         return 0;
     }
 
-    std::lock_guard<std::mutex> lock(c->mutex);
+    std::lock_guard<std::mutex> lock(c->conMapMutex);
     std::memcpy(c->sk, sk, crypto_sign_SECRETKEYBYTES);
     std::memcpy(c->pk, pk, crypto_sign_PUBLICKEYBYTES);
     c->keysSet.store(true);
@@ -2417,13 +2311,64 @@ int UDPC_unset_libsodium_keys(UDPC_HContext ctx) {
         return 0;
     }
 
-    std::lock_guard<std::mutex> lock(c->mutex);
+    std::lock_guard<std::mutex> lock(c->conMapMutex);
     c->keysSet.store(false);
     std::memset(c->pk, 0, crypto_sign_PUBLICKEYBYTES);
     std::memset(c->sk, 0, crypto_sign_SECRETKEYBYTES);
     return 1;
 }
 
+int UDPC_add_whitelist_pk(UDPC_HContext ctx, unsigned char *pk) {
+    UDPC::Context *c = UDPC::verifyContext(ctx);
+    if(!c || !c->flags.test(2)) {
+        return 0;
+    }
+
+    std::lock_guard<std::mutex> lock(c->peerPKWhitelistMutex);
+    auto result = c->peerPKWhitelist.insert(UDPC::PKContainer(pk));
+    if(result.second) {
+        return c->peerPKWhitelist.size();
+    }
+    return 0;
+}
+
+int UDPC_has_whitelist_pk(UDPC_HContext ctx, unsigned char *pk) {
+    UDPC::Context *c = UDPC::verifyContext(ctx);
+    if(!c || !c->flags.test(2)) {
+        return 0;
+    }
+
+    std::lock_guard<std::mutex> lock(c->peerPKWhitelistMutex);
+    if(c->peerPKWhitelist.find(UDPC::PKContainer(pk)) != c->peerPKWhitelist.end()) {
+        return 1;
+    }
+    return 0;
+}
+
+int UDPC_remove_whitelist_pk(UDPC_HContext ctx, unsigned char *pk) {
+    UDPC::Context *c = UDPC::verifyContext(ctx);
+    if(!c || !c->flags.test(2)) {
+        return 0;
+    }
+
+    std::lock_guard<std::mutex> lock(c->peerPKWhitelistMutex);
+    if(c->peerPKWhitelist.erase(UDPC::PKContainer(pk)) != 0) {
+        return 1;
+    }
+    return 0;
+}
+
+int UDPC_clear_whitelist(UDPC_HContext ctx) {
+    UDPC::Context *c = UDPC::verifyContext(ctx);
+    if(!c || !c->flags.test(2)) {
+        return 0;
+    }
+
+    std::lock_guard<std::mutex> lock(c->peerPKWhitelistMutex);
+    c->peerPKWhitelist.clear();
+    return 1;
+}
+
 int UDPC_get_auth_policy(UDPC_HContext ctx) {
     UDPC::Context *c = UDPC::verifyContext(ctx);
     if(!c) {

src/UDPConnection.h

@@ -196,8 +196,7 @@ typedef enum {
     UDPC_ET_CONNECTED,
     UDPC_ET_DISCONNECTED,
     UDPC_ET_GOOD_MODE,
-    UDPC_ET_BAD_MODE,
+    UDPC_ET_BAD_MODE
-    UDPC_ET_REQUEST_CONNECT_PK
 } UDPC_EventType;
 
 /*!
@@ -216,7 +215,6 @@ typedef struct {
     union Value {
         int dropAllWithAddr;
         int enableLibSodium;
-        unsigned char *pk;
     } v;
 } UDPC_Event;
 
@@ -403,25 +401,6 @@ void UDPC_client_initiate_connection(
     UDPC_ConnectionId connectionId,
     int enableLibSodium);
 
-/*!
- * \brief Initiate a connection to a server peer with an expected public key
- *
- * Note that this function does nothing on a server context.
- *
- * \param ctx The context to initiate a connection from
- * \param connectionId The server peer to initiate a connection to
- * \param serverPK A pointer to the public key that the server is expected to
- * use (if the server does not use this public key, then the connection will
- * fail; it must point to a buffer of size \p crypto_sign_PUBLICKEYBYTES)
- *
- * This function assumes that support for libsodium was enabled when UDPC was
- * compiled. If it has not, then this function will fail.
- */
-void UDPC_client_initiate_connection_pk(
-    UDPC_HContext ctx,
-    UDPC_ConnectionId connectionId,
-    unsigned char *serverPK);
-
 /*!
  * \brief Queues a packet to be sent to the specified peer
  *
@@ -510,6 +489,11 @@ int UDPC_set_libsodium_key_easy(UDPC_HContext ctx, unsigned char *sk);
 
 int UDPC_unset_libsodium_keys(UDPC_HContext ctx);
 
+int UDPC_add_whitelist_pk(UDPC_HContext ctx, unsigned char *pk);
+int UDPC_has_whitelist_pk(UDPC_HContext ctx, unsigned char *pk);
+int UDPC_remove_whitelist_pk(UDPC_HContext ctx, unsigned char *pk);
+int UDPC_clear_whitelist(UDPC_HContext ctx);
+
 int UDPC_get_auth_policy(UDPC_HContext ctx);
 int UDPC_set_auth_policy(UDPC_HContext ctx, int value);
 

src/test/UDPC_NetworkTest.c

@@ -12,6 +12,7 @@
 
 #define QUEUED_MAX_SIZE 32
 #define SEND_IDS_SIZE 64
+#define WHITELIST_FILES_SIZE 64
 
 void usage() {
     puts("[-c | -s] - client or server (default server)");
@@ -24,7 +25,7 @@ void usage() {
     puts("-l (silent|error|warning|info|verbose|debug) - log level, default debug");
     puts("-e - enable receiving events");
     puts("-ls - enable libsodium");
-    puts("-ck <pubkey_file> - connect to server expecting this public key");
+    puts("-ck <pubkey_file> - add pubkey to whitelist");
     puts("-sk <pubkey> <seckey> - start with pub/sec key pair");
     puts("-p <\"fallback\" or \"strict\"> - set auth policy");
 }
@@ -57,6 +58,9 @@ int main(int argc, char **argv) {
     const char *seckey_file = NULL;
     unsigned char pubkey[crypto_sign_PUBLICKEYBYTES];
     unsigned char seckey[crypto_sign_SECRETKEYBYTES];
+    const char *whitelist_pk_files[WHITELIST_FILES_SIZE];
+    unsigned int whitelist_pk_files_index = 0;
+    unsigned char whitelist_pks[WHITELIST_FILES_SIZE][crypto_sign_PUBLICKEYBYTES];
     int authPolicy = UDPC_AUTH_POLICY_FALLBACK;
 
     while(argc > 0) {
@@ -111,7 +115,11 @@ int main(int argc, char **argv) {
             puts("Enabled libsodium");
         } else if(strcmp(argv[0], "-ck") == 0 && argc > 1) {
             --argc; ++argv;
-            pubkey_file = argv[0];
+            if(whitelist_pk_files_index >= WHITELIST_FILES_SIZE) {
+                puts("ERROR: limit reached for whitelisted pks");
+                return 1;
+            }
+            whitelist_pk_files[whitelist_pk_files_index++] = argv[0];
         } else if(strcmp(argv[0], "-sk") == 0 && argc > 2) {
             --argc; ++argv;
             pubkey_file = argv[0];
@@ -141,38 +149,53 @@ int main(int argc, char **argv) {
     if(isLibSodiumEnabled == 0) {
         puts("Disabled libsodium");
     } else {
-        if(pubkey_file) {
+        if(pubkey_file && seckey_file) {
-            FILE *pubkey_f = fopen(pubkey_file, "r");
+            FILE *pubkey_f = fopen(pubkey_file, "rb");
             if(!pubkey_f) {
                 printf("ERROR: Failed to open pubkey_file \"%s\"\n", pubkey_file);
                 return 1;
             }
-            size_t count = fread(pubkey, 1, crypto_sign_PUBLICKEYBYTES, pubkey_f);
+            size_t count = fread(pubkey, crypto_sign_PUBLICKEYBYTES, 1, pubkey_f);
-            if(count != crypto_sign_PUBLICKEYBYTES) {
+            if(count != 1) {
                 fclose(pubkey_f);
                 printf("ERROR: Failed to read pubkey_file \"%s\"\n", pubkey_file);
                 return 1;
             }
             fclose(pubkey_f);
 
-            if(seckey_file) {
+            FILE *seckey_f = fopen(seckey_file, "rb");
-                FILE *seckey_f = fopen(seckey_file, "r");
+            if(!seckey_f) {
-                if(!seckey_f) {
+                printf("ERROR: Failed to open seckey_file \"%s\"\n", seckey_file);
-                    printf("ERROR: Failed to open seckey_file \"%s\"\n", seckey_file);
+                return 1;
-                    return 1;
-                }
-                count = fread(seckey, 1, crypto_sign_SECRETKEYBYTES, seckey_f);
-                if(count != crypto_sign_SECRETKEYBYTES) {
-                    fclose(seckey_f);
-                    printf("ERROR: Failed to read seckey_file \"%s\"\n", seckey_file);
-                    return 1;
-                }
-                fclose(seckey_f);
             }
-        } else if(seckey_file) {
+            count = fread(seckey, crypto_sign_SECRETKEYBYTES, 1, seckey_f);
-            printf("ERROR: Invalid state (seckey_file defined but not pubkey_file)\n");
+            if(count != 1) {
+                fclose(seckey_f);
+                printf("ERROR: Failed to read seckey_file \"%s\"\n", seckey_file);
+                return 1;
+            }
+            fclose(seckey_f);
+        } else if(pubkey_file || seckey_file) {
+            printf("ERROR: Invalid state (pubkey_file and seckey_file not "
+                "defined)\n");
             return 1;
         }
+        for(unsigned int i = 0; i < whitelist_pk_files_index; ++i) {
+            FILE *pkf = fopen(whitelist_pk_files[i], "rb");
+            if(!pkf) {
+                printf("ERROR: Failed to open whitelist pubkey file \"%s\"\n",
+                    whitelist_pk_files[i]);
+                return 1;
+            }
+            size_t count = fread(whitelist_pks[i], crypto_sign_PUBLICKEYBYTES, 1, pkf);
+            if(count != 1) {
+                fclose(pkf);
+                printf("ERROR: Failed to read whitelist pubkey file \"%s\"\n",
+                    whitelist_pk_files[i]);
+                return 1;
+            }
+            fclose(pkf);
+        }
     }
 
     if(!listenAddr) {
@@ -207,9 +230,9 @@ int main(int argc, char **argv) {
 
     UDPC_set_logging_type(context, logLevel);
     UDPC_set_receiving_events(context, isReceivingEvents);
-    if(!isClient && pubkey_file && seckey_file) {
+    if(pubkey_file && seckey_file) {
         UDPC_set_libsodium_keys(context, seckey, pubkey);
-        puts("Set pubkey/seckey for server");
+        puts("Set pubkey/seckey");
     }
 
     UDPC_set_auth_policy(context, authPolicy);
@@ -220,6 +243,16 @@ int main(int argc, char **argv) {
         puts("Auth policy set to \"strict\"");
     }
 
+    if(isLibSodiumEnabled && whitelist_pk_files_index > 0) {
+        puts("Enabling pubkey whitelist...");
+        for(unsigned int i = 0; i < whitelist_pk_files_index; ++i) {
+            if(UDPC_add_whitelist_pk(context, whitelist_pks[i]) != i + 1) {
+                puts("Failed to add pubkey to whitelist");
+                return 1;
+            }
+        }
+    }
+
     UDPC_enable_threaded_update(context);
 
     unsigned int tick = 0;
@@ -235,11 +268,8 @@ int main(int argc, char **argv) {
     while(1) {
         sleep_seconds(1);
         if(isClient && UDPC_has_connection(context, connectionId) == 0) {
-            if(isLibSodiumEnabled && pubkey_file) {
+
-                UDPC_client_initiate_connection_pk(context, connectionId, pubkey);
+            UDPC_client_initiate_connection(context, connectionId, isLibSodiumEnabled);
-            } else {
-                UDPC_client_initiate_connection(context, connectionId, isLibSodiumEnabled);
-            }
         }
         if(!noPayload) {
             list = UDPC_get_list_connected(context, &temp);