mkproof/src/mkproof.c

#include <assert.h>
#include <errno.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>
#include "argon2.h"
#include "random.h"
#include "util.h"
#include "proof.h"


#define MAX_PIDS 512

static void
die(int check, char *why)
{
	if (check) {
		fprintf(stderr, "Error: %s\n", why);
		exit(1);
	}
}

static void
draw_spinner(void)
{
	static size_t i = 0;
	const char *spinners[] = {"|", "/", "-", "\\"};
	size_t nspinners = sizeof(spinners) / sizeof(*spinners);

	fprintf(stderr, "\rRunning... %s", spinners[i % nspinners]);

	i++;
}

static void
run_mkproof(argon2_context context, int digits)
{
	unsigned char password[16];
	unsigned char hash[32];

	context.out = hash;
	context.outlen = sizeof(hash);
	context.pwd = password;
	context.pwdlen = sizeof(password);

	int valid = 0;
	while (!valid) {
		ssize_t b = get_random_bytes(password, sizeof(password));
		assert(b == sizeof(password));

		int r = argon2id_ctx(&context);
		die(r != 0, "argon2id failed\n");

		valid = hash_msb(hash) >= digits;
	}

	if (isatty(STDERR_FILENO)) {
		fprintf(stderr, "\rHere is your proof:\n");
	}

	char proof[33];
	enchex(password, sizeof(password), proof, sizeof(proof));
	printf("%s\n", proof);
	_exit(0);
}

int
main(int argc, char *argv[])
{
#ifdef _SC_NPROCESSORS_ONLN
	long nprocs = sysconf(_SC_NPROCESSORS_ONLN) - 1;
#else
	long nprocs = 2;
#endif

	char *endptr;
	int c;
	while ((c = getopt(argc, argv, "j:")) != -1) {
		switch (c) {
		case 'j':
			nprocs = strtol(optarg, &endptr, 10);
			die(*endptr, "error parsing nprocs");
			break;
		}
	}

	if ((argc - optind) != 1) {
		fprintf(stderr, "Usage: %s [-j nprocs] <challenge>\n", argv[0]);
		return 1;
	}

	if (nprocs <= 0 || nprocs > MAX_PIDS) {
		fprintf(stderr, "please specify nprocs value between 1 and "
				"%d\n", MAX_PIDS);
		return 1;
	}

	int iters, memory, digits;
	unsigned char salt[16];

	char *challenge = argv[optind];
	char *algo = strtok(challenge, ":");
	if (strcmp(algo, "argon2id") != 0) {
		fprintf(stderr, "Error: unknown challenge type %s\n", algo);
		return 1;
	}

	char *iterstr = strtok(NULL, ":");
	iters = strtoul(iterstr, &endptr, 10);
	die(*endptr, "Invalid challenge");

	char *memorystr = strtok(NULL, ":");
	memory = strtoul(memorystr, &endptr, 10);
	die(*endptr, "Invalid challenge");

	char *digitsstr = strtok(NULL, ":");
	digits = strtoul(digitsstr, &endptr, 10);
	die(*endptr, "Invalid challenge");

	char *saltstr = strtok(NULL, ":");
	die(strlen(saltstr) != 32, "Invalid challenge");
	int r = dechex(saltstr, strlen(saltstr), salt, sizeof(salt));
	die(r == -1, "Invalid challenge");

	if (isatty(STDERR_FILENO)) {
		// TODO: Provide a better estimate based on the difficulty
		fprintf(stderr, "Generating a proof of work.\n");
		fprintf(stderr, "This may take anywhere from several minutes to a few hours on slow computers.\n");
	}

	pid_t pids[MAX_PIDS];
	assert(nprocs <= MAX_PIDS);

	for (size_t i = 0; i < (size_t)nprocs; i++) {
		pid_t pid = fork();
		die(pid < 0, "fork()");

		if (pid == 0) {
			argon2_context context = {
				.salt = salt,
				.saltlen = sizeof(salt),
				.t_cost = iters,
				.m_cost = memory,
				.lanes = 1,
				.threads = 1,
				.flags = ARGON2_DEFAULT_FLAGS,
				.version = ARGON2_VERSION_NUMBER,
			};
			run_mkproof(context, digits);
			// unreachable
			assert(0);
		}

		pids[i] = pid;
	}

	pid_t pid;
	while (1) {
		int wstatus;
		pid = waitpid(-1, &wstatus, WNOHANG);
		if (pid > 0) {
			break;
		}

		draw_spinner();
		sleep(1);
	}

	for (size_t i = 0; i < (size_t)nprocs; i++) {
		if (pids[i] == pid) {
			continue;
		}

		r = kill(pids[i], SIGTERM);
		die(r == -1, strerror(errno));
	}

	return 0;
}