Configurable number of CPU cores.

RfidResearchGroup · Aug 17, 2023 · dfabad3 · dfabad3
1 parent 4811c84
commit dfabad3
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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