main.cpp

/*
 * This file is part of the BSGS distribution (https://github.com/JeanLucPons/Kangaroo).
 * Copyright (c) 2020 Jean Luc PONS.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, version 3.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program. If not, see <http://www.gnu.org/licenses/>.
*/

#include "Kangaroo.h"
#include "Timer.h"
#include "SECPK1/SECP256k1.h"
#include "GPU/GPUEngine.h"
#include <errno.h>
#include <fstream>
#include <sys/stat.h>
#include <string>
#include <string.h>
#include <stdexcept>

using namespace std;

class MainError : public std::runtime_error {
public:
  MainError(RunResult result,const std::string& message)
    : std::runtime_error(message), result(result) {
  }

  RunResult result;
};

#define CHECKARG(opt,n) if(a>=argc-1) {::printf(opt " missing argument #%d\n",n); throw MainError(RESULT_INVALID_INPUT, opt " missing argument");} else {a++;}

// ------------------------------------------------------------------------------------------

void printUsage() {

  printf("Kangaroo [-v] [-t nbThread] [-d dpBit] [gpu] [-check]\n");
  printf("         [-gpuId gpuId1[,gpuId2,...]] [-g g1x,g1y[,g2x,g2y,...]]\n");
  printf("         inFile\n");
  printf(" -v: Print version\n");
  printf(" -gpu: Enable gpu calculation\n");
  printf(" -gpuId gpuId1,gpuId2,...: List of GPU(s) to use, default is 0\n");
  printf(" -g g1x,g1y,g2x,g2y,...: Specify GPU(s) kernel gridsize, default is 2*(MP),2*(Core/MP)\n");
  printf(" -d: Specify number of leading zeros for the DP method (default is auto)\n");
  printf(" -t nbThread: Secify number of thread\n");
  printf(" -w workfile: Specify file to save work into (current processed key only)\n");
  printf(" -i workfile: Specify file to load work from (current processed key only)\n");
  printf(" --resume-if-present <file> : Resume from workfile if it exists, otherwise start fresh\n");
  printf(" -wi workInterval: Periodic interval (in seconds) for saving work\n");
  printf(" -ws: Save kangaroos in the work file\n");
  printf(" -wss: Save kangaroos via the server\n");
  printf(" -wsplit: Split work file of server and reset hashtable\n");
  printf(" -wm file1 file2 destfile: Merge work file\n");
  printf(" -wmdir dir destfile: Merge directory of work files\n");
  printf(" -wt timeout: Save work timeout in millisec (default is 3000ms)\n");
  printf(" -winfo file1: Work file info file\n");
  printf(" -wpartcreate name: Create empty partitioned work file (name is a directory)\n");
  printf(" -wcheck worfile: Check workfile integrity\n");
  printf(" -m maxStep: number of operations before give up the search (maxStep*expected operation)\n");
  printf(" -s: Start in server mode\n");
  printf(" -c server_ip: Start in client mode and connect to server server_ip\n");
  printf(" -sp port: Server port, default is 17403\n");
  printf(" -nt timeout: Network timeout in millisec (default is 3000ms)\n");
  printf(" -o fileName: output result to fileName\n");
  printf(" --pubkey fileName: encrypt recovered result to fileName using an RSA public key\n");
  printf(" --cleanup-on-found: best-effort scrub/delete workfile and config after encrypted success\n");
  printf(" -l: List cuda enabled devices\n");
  printf(" -check: Check GPU kernel vs CPU\n");
  printf(" inFile: intput configuration file\n");
  exit(0);

}

// ------------------------------------------------------------------------------------------

int getInt(string name,char *v) {

  int r;

  try {

    r = std::stoi(string(v));

  } catch(std::invalid_argument&) {

    printf("Invalid %s argument, number expected\n",name.c_str());
    throw MainError(RESULT_INVALID_INPUT,"Invalid integer argument");

  }

  return r;

}

double getDouble(string name,char *v) {

  double r;

  try {

    r = std::stod(string(v));

  } catch(std::invalid_argument&) {

    printf("Invalid %s argument, number expected\n",name.c_str());
    throw MainError(RESULT_INVALID_INPUT,"Invalid floating-point argument");

  }

  return r;

}

// ------------------------------------------------------------------------------------------

void getInts(string name,vector<int> &tokens,const string &text,char sep) {

  size_t start = 0,end = 0;
  tokens.clear();
  int item;

  try {

    while((end = text.find(sep,start)) != string::npos) {
      item = std::stoi(text.substr(start,end - start));
      tokens.push_back(item);
      start = end + 1;
    }

    item = std::stoi(text.substr(start));
    tokens.push_back(item);

  }
  catch(std::invalid_argument &) {

    printf("Invalid %s argument, number expected\n",name.c_str());
    throw MainError(RESULT_INVALID_INPUT,"Invalid integer list argument");

  }

}
// ------------------------------------------------------------------------------------------

// Default params
static int dp = -1;
static int nbCPUThread;
static string configFile = "";
static bool checkFlag = false;
static bool gpuEnable = false;
static vector<int> gpuId = { 0 };
static vector<int> gridSize;
static string workFile = "";
static string checkWorkFile = "";
static string iWorkFile = "";
static string resumeIfPresentFile = "";
static uint32_t savePeriod = 60;
static bool saveKangaroo = false;
static bool saveKangarooByServer = false;
static string merge1 = "";
static string merge2 = "";
static string mergeDest = "";
static string mergeDir = "";
static string infoFile = "";
static double maxStep = 0.0;
static int wtimeout = 3000;
static int ntimeout = 3000;
static int port = 17403;
static bool serverMode = false;
static string serverIP = "";
static string outputFile = "";
static string publicKeyFile = "";
static bool cleanupOnFound = false;
static bool splitWorkFile = false;

struct ResumeFileStatus {
  bool exists;
  bool nonEmpty;
  long long size;
};

static ResumeFileStatus GetResumeFileStatus(const string& path) {

#ifdef _WIN32
  struct _stat64 st;
  if(_stat64(path.c_str(),&st) != 0) {
#else
  struct stat st;
  if(stat(path.c_str(),&st) != 0) {
#endif
    if(errno == ENOENT)
      return { false,false,0 };
    ::printf("[resume] Cannot access %s\n",path.c_str());
    ::printf("%s\n",::strerror(errno));
    throw MainError(RESULT_BAD_WORKFILE,"Cannot access resume workfile");
  }

  if((st.st_mode & S_IFMT) == S_IFDIR) {
    ::printf("[resume] %s is a directory\n",path.c_str());
    throw MainError(RESULT_BAD_WORKFILE,"Resume workfile path is a directory");
  }

  return { true,st.st_size > 0,(long long)st.st_size };

}

static int RunResultToExitCode(RunResult result) {

  switch(result) {
    case RESULT_NO_KEY:
      return 0;
    case RESULT_KEY_FOUND:
      return 10;
    case RESULT_INVALID_INPUT:
      return 20;
    case RESULT_BAD_WORKFILE:
      return 21;
    case RESULT_RUNTIME_ERROR:
      return 30;
    case RESULT_OUTPUT_ERROR:
      return 31;
    case RESULT_SIGNAL_SAVE_OK:
      return 40;
    case RESULT_SIGNAL_SAVE_FAILED:
      return 41;
    case RESULT_FATAL_ERROR:
    default:
      return 50;
  }

}

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

  try {

#ifdef USE_SYMMETRY
  printf("Kangaroo v" RELEASE " (with symmetry)\n");
#else
  printf("Kangaroo v" RELEASE "\n");
#endif

  // Global Init
  Timer::Init();
  rseed(Timer::getSeed32());

  // Init SecpK1
  Secp256K1 *secp = new Secp256K1();
  secp->Init();

  int a = 1;
  nbCPUThread = Timer::getCoreNumber();

  while (a < argc) {

    if(strcmp(argv[a], "-t") == 0) {
      CHECKARG("-t",1);
      nbCPUThread = getInt("nbCPUThread",argv[a]);
      a++;
    } else if(strcmp(argv[a],"-d") == 0) {
      CHECKARG("-d",1);
      dp = getInt("dpSize",argv[a]);
      a++;
    } else if (strcmp(argv[a], "-h") == 0) {
      printUsage();
    } else if(strcmp(argv[a],"-l") == 0) {

#ifdef WITHGPU
      GPUEngine::PrintCudaInfo();
#else
      printf("GPU code not compiled, use -DWITHGPU when compiling.\n");
#endif
      exit(0);

    } else if(strcmp(argv[a],"-w") == 0) {
      CHECKARG("-w",1);
      workFile = string(argv[a]);
      a++;
    } else if(strcmp(argv[a],"-i") == 0) {
      CHECKARG("-i",1);
      iWorkFile = string(argv[a]);
      a++;
    } else if(strcmp(argv[a],"--resume-if-present") == 0) {
      CHECKARG("--resume-if-present",1);
      resumeIfPresentFile = string(argv[a]);
      a++;
    } else if(strcmp(argv[a],"-wm") == 0) {
      CHECKARG("-wm",1);
      merge1 = string(argv[a]);
      CHECKARG("-wm",2);
      merge2 = string(argv[a]);
      a++;
      if(a<argc) {
        // classic merge
        mergeDest = string(argv[a]);
        a++;
      }
    } else if(strcmp(argv[a],"-wmdir") == 0) {
      CHECKARG("-wmdir",1);
      mergeDir = string(argv[a]);
      CHECKARG("-wmdir",2);
      mergeDest = string(argv[a]);
      a++;
    }  else if(strcmp(argv[a],"-wcheck") == 0) {
      CHECKARG("-wcheck",1);
      checkWorkFile = string(argv[a]);
      a++;
    }  else if(strcmp(argv[a],"-winfo") == 0) {
      CHECKARG("-winfo",1);
      infoFile = string(argv[a]);
      a++;
    } else if(strcmp(argv[a],"-o") == 0) {
      CHECKARG("-o",1);
      outputFile = string(argv[a]);
      a++;
    } else if(strcmp(argv[a],"--pubkey") == 0) {
      CHECKARG("--pubkey",1);
      publicKeyFile = string(argv[a]);
      a++;
    } else if(strcmp(argv[a],"--cleanup-on-found") == 0) {
      cleanupOnFound = true;
      a++;
    } else if(strcmp(argv[a],"-wi") == 0) {
      CHECKARG("-wi",1);
      savePeriod = getInt("savePeriod",argv[a]);
      a++;
    } else if(strcmp(argv[a],"-wt") == 0) {
      CHECKARG("-wt",1);
      wtimeout = getInt("timeout",argv[a]);
      a++;
    } else if(strcmp(argv[a],"-nt") == 0) {
      CHECKARG("-nt",1);
      ntimeout = getInt("timeout",argv[a]);
      a++;
    } else if(strcmp(argv[a],"-m") == 0) {
      CHECKARG("-m",1);
      maxStep = getDouble("maxStep",argv[a]);
      a++;
    } else if(strcmp(argv[a],"-ws") == 0) {
      a++;
      saveKangaroo = true;
    } else if(strcmp(argv[a],"-wss") == 0) {
      a++;
      saveKangarooByServer = true;
    } else if(strcmp(argv[a],"-wsplit") == 0) {
      a++;
      splitWorkFile = true;
    } else if(strcmp(argv[a],"-wpartcreate") == 0) {
      CHECKARG("-wpartcreate",1);
      workFile = string(argv[a]);
      Kangaroo::CreateEmptyPartWork(workFile);
      exit(0);
    } else if(strcmp(argv[a],"-s") == 0) {
      a++;
      serverMode = true;
    } else if(strcmp(argv[a],"-c") == 0) {
      CHECKARG("-c",1);
      serverIP = string(argv[a]);
      a++;
    } else if(strcmp(argv[a],"-sp") == 0) {
      CHECKARG("-sp",1);
      port = getInt("serverPort",argv[a]);
      a++;
    } else if(strcmp(argv[a],"-gpu") == 0) {
      gpuEnable = true;
      a++;
    } else if(strcmp(argv[a],"-gpuId") == 0) {
      CHECKARG("-gpuId",1);
      getInts("gpuId",gpuId,string(argv[a]),',');
      a++;
    } else if(strcmp(argv[a],"-g") == 0) {
      CHECKARG("-g",1);
      getInts("gridSize",gridSize,string(argv[a]),',');
      a++;
    } else if(strcmp(argv[a],"-v") == 0) {
      ::exit(0);
    } else if(strcmp(argv[a],"-check") == 0) {
      checkFlag = true;
      a++;
    } else if(a == argc - 1) {
      configFile = string(argv[a]);
      a++;
    } else {
      printf("Unexpected %s argument\n",argv[a]);
      throw MainError(RESULT_INVALID_INPUT,"Unexpected argument");
    }

  }

  if(iWorkFile.length() > 0 && resumeIfPresentFile.length() > 0) {
    ::printf("Cannot use -i and --resume-if-present together\n");
    throw MainError(RESULT_INVALID_INPUT,"Conflicting workfile resume options");
  }

  if(resumeIfPresentFile.length() > 0) {
    ResumeFileStatus resumeStatus = GetResumeFileStatus(resumeIfPresentFile);
    if(resumeStatus.exists && resumeStatus.nonEmpty) {
      ::printf("[resume] Resuming from %s (size=%lld)\n",resumeIfPresentFile.c_str(),resumeStatus.size);
      iWorkFile = resumeIfPresentFile;
    } else {
      ::printf("[resume] No valid workfile found, starting fresh\n");
    }
  }

  if(gridSize.size() == 0) {
    for(int i = 0; i < gpuId.size(); i++) {
      gridSize.push_back(0);
      gridSize.push_back(0);
    }
  } else if(gridSize.size() != gpuId.size() * 2) {
    printf("Invalid gridSize or gpuId argument, must have coherent size\n");
    throw MainError(RESULT_INVALID_INPUT,"Invalid gridSize or gpuId argument");
  }

  Kangaroo *v = new Kangaroo(secp,dp,gpuEnable,workFile,iWorkFile,savePeriod,saveKangaroo,saveKangarooByServer,
                             maxStep,wtimeout,port,ntimeout,serverIP,outputFile,publicKeyFile,cleanupOnFound,splitWorkFile);
  if(checkFlag) {
    v->Check(gpuId,gridSize);  
    exit(0);
  } else {
    if(checkWorkFile.length() > 0) {
      v->CheckWorkFile(nbCPUThread,checkWorkFile);
      exit(0);
    } if(infoFile.length()>0) {
      v->WorkInfo(infoFile);
      exit(0);
    } else if(mergeDir.length() > 0) {
      v->MergeDir(mergeDir,mergeDest);
      exit(0);
    } else if(merge1.length()>0) {
      v->MergeWork(merge1,merge2,mergeDest);
      exit(0);
    } if(iWorkFile.length()>0) {
      if( !v->LoadWork(iWorkFile) )
        return RunResultToExitCode(RESULT_BAD_WORKFILE);
    } else if(configFile.length()>0) {
      if( !v->ParseConfigFile(configFile) )
        return RunResultToExitCode(RESULT_INVALID_INPUT);
    } else {
      if(serverIP.length()==0) {
        ::printf("No input file to process\n");
        throw MainError(RESULT_INVALID_INPUT,"No input file to process");
      }
    }
    if(!v->ValidateOutputConfiguration())
      return RunResultToExitCode(v->GetRunResult());
    RunResult result;
    if(serverMode)
      result = v->RunServer();
    else
      result = v->Run(nbCPUThread,gpuId,gridSize);
    return RunResultToExitCode(result);
  }

  return 0;

  } catch(const MainError& e) {
    return RunResultToExitCode(e.result);
  } catch(const std::exception& e) {
    fprintf(stderr,"Fatal error: %s\n",e.what());
    return RunResultToExitCode(RESULT_FATAL_ERROR);
  } catch(...) {
    fprintf(stderr,"Fatal error: unexpected exception\n");
    return RunResultToExitCode(RESULT_FATAL_ERROR);
  }

}