From dfabad37d95dbd4c1a82fefbf41358380cf94c0d Mon Sep 17 00:00:00 2001
From: Roman D <me@iamroman.org>
Date: Thu, 17 Aug 2023 12:41:34 +0200
Subject: [PATCH] Configurable number of CPU cores.

---
 client/src/proxmark3.c | 17 +++++++++++++++++
 client/src/util.c      |  6 ++++++
 client/src/util.h      |  1 +
 3 files changed, 24 insertions(+)

diff --git a/client/src/proxmark3.c b/client/src/proxmark3.c
index f6c25a9eb..26f4a5698 100644
--- a/client/src/proxmark3.c
+++ b/client/src/proxmark3.c
@@ -581,6 +581,7 @@ static void show_help(bool showFullHelp, char *exec_name) {
         PrintAndLogEx(NORMAL, "      --unlock-bootloader                 Enable flashing of bootloader area *DANGEROUS* (need --flash)");
         PrintAndLogEx(NORMAL, "      --force                             Enable flashing even if firmware seems to not match client version");
         PrintAndLogEx(NORMAL, "      --image <imagefile>                 image to flash. Can be specified several times.");
+        PrintAndLogEx(NORMAL, "      --ncpu <num_cores>                  override number of CPU cores");
         PrintAndLogEx(NORMAL, "\nExamples:");
         PrintAndLogEx(NORMAL, "\n  to run Proxmark3 client:\n");
         PrintAndLogEx(NORMAL, "      %s "SERIAL_PORT_EXAMPLE_H"                       -- runs the pm3 client", exec_name);
@@ -996,6 +997,22 @@ int main(int argc, char *argv[]) {
             continue;
         }
 
+        if (strcmp(argv[i], "--ncpu") == 0) {
+            if (i + 1 == argc) {
+                PrintAndLogEx(ERR, _RED_("ERROR:") " missing CPU number specification after --ncpu\n");
+                show_help(false, exec_name);
+                return 1;
+            }
+            long int ncpus = strtol(argv[i + 1], NULL, 10);
+            if (ncpus < 0 || ncpus >= INT_MAX) {
+                PrintAndLogEx(ERR, _RED_("ERROR:") " invalid number of CPU cores: --ncpu " _YELLOW_("%s") "\n", argv[i + 1]);
+                return 1;
+            }
+            g_numCPUs = ncpus;
+            i++;
+            continue;
+        }
+
         // We got an unknown parameter
         PrintAndLogEx(ERR, _RED_("ERROR:") " invalid parameter: " _YELLOW_("%s") "\n", argv[i]);
         show_help(false, exec_name);
diff --git a/client/src/util.c b/client/src/util.c
index 51408a524..150b042a5 100644
--- a/client/src/util.c
+++ b/client/src/util.c
@@ -40,6 +40,8 @@ uint8_t g_debugMode = 0;
 uint8_t g_printAndLog = PRINTANDLOG_PRINT | PRINTANDLOG_LOG;
 // global client tell if a pending prompt is present
 bool g_pendingPrompt = false;
+// global CPU core count override
+int g_numCPUs = 0;
 
 #ifdef _WIN32
 #include <windows.h>
@@ -1079,6 +1081,10 @@ uint64_t HornerScheme(uint64_t num, uint64_t divider, uint64_t factor) {
 
 // determine number of logical CPU cores (use for multithreaded functions)
 int num_CPUs(void) {
+    if (g_numCPUs > 0) {
+        return g_numCPUs;
+    }
+
 #if defined(_WIN32)
 #include <sysinfoapi.h>
     SYSTEM_INFO sysinfo;
diff --git a/client/src/util.h b/client/src/util.h
index 4dab9ae06..0712d8d18 100644
--- a/client/src/util.h
+++ b/client/src/util.h
@@ -32,6 +32,7 @@
 extern uint8_t g_debugMode;
 extern uint8_t g_printAndLog;
 extern bool g_pendingPrompt;
+extern int g_numCPUs;
 
 #define PRINTANDLOG_PRINT 1
 #define PRINTANDLOG_LOG   2