From 7887d086f2e7a4f1e3f81aec30cf2237ad3ef54b Mon Sep 17 00:00:00 2001
From: cathugger <cathugger@cock.li>
Date: Mon, 25 Sep 2017 20:49:47 +0300
Subject: [PATCH] improvements all over the place

---
 Makefile      |   7 +--
 base16.h      |  14 +++--
 base16_from.c |  77 ++++++++++++++++++++++------
 base32.h      |  13 ++++-
 base32_from.c |  68 +++++++++++++++++++-----
 base32_to.c   |   2 +-
 main.c        | 139 +++++++++++++++++++++++++++++++++++++++-----------
 7 files changed, 250 insertions(+), 70 deletions(-)

diff --git a/Makefile b/Makefile
index 868dfb9..eeecd4c 100644
--- a/Makefile
+++ b/Makefile
@@ -1,7 +1,8 @@
 
 CC= gcc
-CFLAGS= -O3 -march=native -Wall
-#CFLAGS= -O0 -g3 -fsanitize=address
+CSTD= -std=c99 -Wall -D_POSIX_C_SOURCE=200112L
+CFLAGS= $(CSTD) -O3 -march=native
+#CFLAGS= $(CSTD) -O0 -g3 -fsanitize=address
 MV= mv
 
 ED25519OBJ= $(patsubst %.c,%.o,$(wildcard ed25519/ref10/*.c))
@@ -61,7 +62,7 @@ clean:
 	$(RM) $(EXE)
 
 depend:
-	makedepend -Y -- $(CFLAGS) -- $(MAINOBJ:.o=.c) $(TEST_BASE16OBJ:.o=.c) $(TEST_BASE32OBJ:.o=.c) $(TEST_ED25519OBJ:.o=.c)
+	makedepend -Y -- $(CSTD) -- $(MAINOBJ:.o=.c) $(TEST_BASE16OBJ:.o=.c) $(TEST_BASE32OBJ:.o=.c) $(TEST_ED25519OBJ:.o=.c)
 
 # DO NOT DELETE THIS LINE
 
diff --git a/base16.h b/base16.h
index 74583c3..6ab6baf 100644
--- a/base16.h
+++ b/base16.h
@@ -1,3 +1,11 @@
-
-char *base16_to(char *dst, const u8 *src, size_t slen);
-size_t base16_from(u8 *dst, u8 *dmask, const char *src);
+// converts src[0:slen] to base16 string
+char *base16_to(char *dst,const u8 *src,size_t slen);
+// calculates length needed to store data converted to base16
+#define BASE16_TO_LEN(l) (((l) * 8 + 3) / 4)
+// converts src string from base16
+size_t base16_from(u8 *dst,u8 *dmask,const char *src);
+// calculates length needed to store data converted from base16
+#define BASE16_FROM_LEN(l) (((l) * 4 + 7) / 8)
+// validates base16 string and optionally stores length of valid data
+// returns 1 if whole string is good, 0 if string contains invalid data
+int base16_valid(const char *src,size_t *count);
diff --git a/base16_from.c b/base16_from.c
index 0bce356..9260621 100644
--- a/base16_from.c
+++ b/base16_from.c
@@ -4,32 +4,77 @@
 #include "base16.h"
 
 static const u8 base16f[256] = {
-	['0'] =  0, ['1'] =  1, ['2'] =  2, ['3'] =  3,
-	['4'] =  4, ['5'] =  5, ['6'] =  6, ['7'] =  7,
-	['8'] =  8, ['9'] =  9, ['A'] = 10, ['B'] = 11,
-	['C'] = 12, ['D'] = 13, ['E'] = 14, ['F'] = 15,
-	                        ['a'] = 10, ['b'] = 11,
-	['c'] = 12, ['d'] = 13, ['e'] = 14, ['f'] = 15,
+	//00    01    02    03    04    05    06    07
+	//08    09    0A    0B    0C    0D    0E    0F
+	// 0x00..0x3F
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, // 0x00
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, // 0x08
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, // 0x10
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, // 0x18
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, // 0x20
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, // 0x28
+	0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, // 0x30
+	0x08, 0x09, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, // 0x38
+	// 0x40..0x7F
+	0xFF, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F, 0xFF, // 0x40
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, // 0x48
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, // 0x50
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, // 0x58
+	0xFF, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F, 0xFF, // 0x60
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, // 0x68
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, // 0x70
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, // 0x78
+	// 0x80..0xBF
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	// 0xC0..0xFF
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
 };
 
-size_t base16_from(u8 *dst, u8 *dmask, const char *src)
+size_t base16_from(u8 *dst,u8 *dmask,const char *src)
 {
-	int i, j, k = -1, l;
-	u8 mask = 0;
+	size_t i, j, k = (size_t)-1, l, sk = 0;
+	u8 mask = 0, cmask = 0;
 	for (i = 0;;i += 4) {
 		j = i/4;
-		l = i%8;
 		if (!src[j]) {
-			if (!l)
-				mask = 0xFF;
-			if (k >= 0)
-				dst[k] &= mask;
-			*dmask = mask;
-			return (size_t)(k+1);
+			if (k != (size_t)-1)
+				dst[k] &= cmask;
+			*dmask = cmask;
+			return k + 1;
 		}
+		l = i%8;
 		k = i/8;
+		if (k != sk)
+			cmask = 0;
+		sk = k;
 		mask = (0x0F << 4) >> l;
+		cmask |= mask;
 		dst[k] &= ~mask;
 		dst[k] |= (base16f[(u8)src[j]] << 4) >> l;
 	}
 }
+
+int base16_valid(const char *src,size_t *count)
+{
+	const char *p;
+
+	for (p = src;base16f[(u8)*p] != 0xFF;++p) ;
+
+	if (count)
+		*count = p - src;
+	return !*p;
+}
diff --git a/base32.h b/base32.h
index be8f476..2e18c16 100644
--- a/base32.h
+++ b/base32.h
@@ -1,2 +1,11 @@
-char *base32_to(char *dst, const u8 *src, size_t slen);
-size_t base32_from(u8 *dst, u8 *dmask, const char *src);
+// converts src[0:slen] to base32 string
+char *base32_to(char *dst,const u8 *src,size_t slen);
+// calculates length needed to store data converted to base32
+#define BASE32_TO_LEN(l) (((l) * 8 + 4) / 5)
+// converts src string from base32
+size_t base32_from(u8 *dst,u8 *dmask,const char *src);
+// calculates length needed to store data converted from base
+#define BASE32_FROM_LEN(l) (((l) * 5 + 7) / 8)
+// validates base32 string and optionally stores length of valid data
+// returns 1 if whole string is good, 0 if string contains invalid data
+int base32_valid(const char *src,size_t *count);
diff --git a/base32_from.c b/base32_from.c
index 4a6d577..f66c1e8 100644
--- a/base32_from.c
+++ b/base32_from.c
@@ -4,14 +4,44 @@
 #include "base32.h"
 
 static const u8 base32f[256] = {
-	['a'] =  0, ['b'] =  1, ['c'] =  2, ['d'] =  3,
-	['e'] =  4, ['f'] =  5, ['g'] =  6, ['h'] =  7,
-	['i'] =  8, ['j'] =  9, ['k'] = 10, ['l'] = 11,
-	['m'] = 12, ['n'] = 13, ['o'] = 14, ['p'] = 15,
-	['q'] = 16, ['r'] = 17, ['s'] = 18, ['t'] = 19,
-	['u'] = 20, ['v'] = 21, ['w'] = 22, ['x'] = 23,
-	['y'] = 24, ['z'] = 25, ['2'] = 26, ['3'] = 27,
-	['4'] = 28, ['5'] = 29, ['6'] = 30, ['7'] = 31,
+	//00    01    02    03    04    05    06    07
+	//08    09    0A    0B    0C    0D    0E    0F
+	// 0x00..0x3F
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, // 0x00
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, // 0x08
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, // 0x10
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, // 0x18
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, // 0x20
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, // 0x28
+	0xFF, 0xFF, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E, 0x1F, // 0x30
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, // 0x38
+	// 0x40..0x7F
+	0xFF, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, // 0x40
+	0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, // 0x48
+	0x0F, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, // 0x50
+	0x17, 0x18, 0x19, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, // 0x58
+	0xFF, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, // 0x60
+	0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, // 0x68
+	0x0F, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, // 0x70
+	0x17, 0x18, 0x19, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, // 0x78
+	// 0x80..0xBF
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	// 0xC0..0xFF
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+	0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
 };
 /*
 +--first octet--+-second octet--+--third octet--+--forth octet--+--fifth octet--+
@@ -20,20 +50,19 @@ static const u8 base32f[256] = {
 |4 3 2 1 0|4 3 2 1 0|4 3 2 1 0|4 3 2 1 0|4 3 2 1 0|4 3 2 1 0|4 3 2 1 0|4 3 2 1 0|
 +-1.index-+-2.index-+-3.index-+-4.index-+-5.index-+-6.index-+-7.index-+-8.index-+
 */
-size_t base32_from(u8 *dst, u8 *dmask, const char *src)
+size_t base32_from(u8 *dst,u8 *dmask,const char *src)
 {
-	int i, j, k = -1, l, sk = 0;
+	size_t i, j, k = (size_t)-1, l, sk = 0;
 	u8 mask = 0, cmask = 0;
 	for (i = 0;;i += 5) {
 		j = i/5;
-		l = i%8;
 		if (!src[j]) {
-			if (k >= 0)
+			if (k != (size_t)-1)
 				dst[k] &= cmask;
-			//printf("dst[k]:%02X mask:%02X l:%d\n", dst[k], (cmask & 0xFF), l);
 			*dmask = cmask;
-			return (size_t)(k + 1);
+			return k + 1;
 		}
+		l = i%8;
 		k = i/8;
 		if (k != sk)
 			cmask = 0;
@@ -57,3 +86,14 @@ size_t base32_from(u8 *dst, u8 *dmask, const char *src)
 		//printf("mask1: %02x\n", ((0x1F << 8) >> (i%8+5-8)) & 0xFF);
 	}
 }
+
+int base32_valid(const char *src,size_t *count)
+{
+	const char *p;
+
+	for (p = src;base32f[(u8)*p] != 0xFF;++p) ;
+
+	if (count)
+		*count = p - src;
+	return !*p;
+}
diff --git a/base32_to.c b/base32_to.c
index c0d307c..781c7e4 100644
--- a/base32_to.c
+++ b/base32_to.c
@@ -23,7 +23,7 @@ static const char base32t[32] = {
 // masks:
 // 0xFF 0x7F 0x3F 0x1F 0x0F 0x07 0x03 0x01
 //  255  127  63    31   15   7     3    1
-char *base32_to(char *dst, const u8 *src, size_t slen)
+char *base32_to(char *dst,const u8 *src,size_t slen)
 {
 	//printf("slen = %d\n", slen);
 	//printhex(base32t, 32);
diff --git a/main.c b/main.c
index 9aa111a..0dba49d 100644
--- a/main.c
+++ b/main.c
@@ -2,7 +2,10 @@
 #include <stdlib.h>
 #include <stdint.h>
 #include <string.h>
+#include <assert.h>
+#include <time.h>
 #include <pthread.h>
+#include <signal.h>
 #include <sodium/randombytes.h>
 #include "ed25519/ref10/ed25519_ref10.h"
 #include "ed25519/ref10/ge.h"
@@ -41,11 +44,22 @@ static size_t direndpos;     // end of dir before .onion within string
 static size_t printstartpos; // where to start printing from
 static size_t printlen;      // precalculated, related to printstartpos
 
-static FILE *fout;
 static pthread_mutex_t fout_mutex;
-
-static volatile int endwork = 0;
+static FILE *fout;
+static size_t numneedgenerate = 0;
+static pthread_mutex_t keysgenerated_mutex;
 static volatile size_t keysgenerated = 0;
+static volatile int endwork = 0;
+
+static void termhandler(int sig)
+{
+	switch (sig) {
+	case SIGTERM:
+	case SIGINT:
+		endwork = 1;
+		break;
+	}
+}
 
 struct binfilter {
 	u8 *f;
@@ -62,17 +76,25 @@ static void filters_init()
 
 static void filters_add(const char *filter)
 {
-	u8 buf[256];
 	struct binfilter bf;
-	size_t ret;
+	size_t ret, ret2;
 
-	ret = base32_from(buf,&bf.mask,filter);
-	//printf("--m:%02X\n", bf.mask);
+	if (!base32_valid(filter,&ret)) {
+		fprintf(stderr, "filter \"%s\" is invalid\n", filter);
+		return;
+	}
+	ret = BASE32_FROM_LEN(ret);
 	if (!ret)
 		return;
+	if (ret > PUBLIC_LEN) {
+		fprintf(stderr, "filter \"%s\" is too long\n", filter);
+		return;
+	}
 	bf.f = malloc(ret);
+	ret2 = base32_from(bf.f,&bf.mask,filter);
+	assert(ret == ret2);
+	//printf("--m:%02X\n", bf.mask);
 	bf.len = ret - 1;
-	memcpy(bf.f,buf,ret);
 	VEC_ADD(bfilters,bf);
 }
 
@@ -100,8 +122,13 @@ static void loadfilterfile(const char *fname)
 
 static void printfilters()
 {
-	fprintf(stderr, "current filters:\n");
-	for (size_t i = 0; i < VEC_LENGTH(bfilters); ++i) {
+	size_t i,l = VEC_LENGTH(bfilters);
+	if (l)
+		fprintf(stderr, "filters:\n");
+	else
+		fprintf(stderr, "no filters defined\n");
+
+	for (i = 0;i < l;++i) {
 		char buf0[256],buf1[256];
 		u8 bufx[128];
 		size_t len = VEC_BUF(bfilters,i).len + 1;
@@ -122,9 +149,37 @@ static void onionready(char *sname, const u8 *secret, const u8 *pubonion)
 {
 	FILE *fh;
 
-	if (mkdir(sname, 0700) != 0)
+	if (endwork)
 		return;
-	
+
+	if (numneedgenerate) {
+		pthread_mutex_lock(&keysgenerated_mutex);
+		if (keysgenerated >= numneedgenerate) {
+			pthread_mutex_unlock(&keysgenerated_mutex);
+			return;
+		}
+	}
+
+	if (mkdir(sname, 0700) != 0) {
+		if (numneedgenerate)
+			pthread_mutex_unlock(&keysgenerated_mutex);
+		return;
+	}
+
+	if (numneedgenerate) {
+		++keysgenerated;
+		if (keysgenerated >= numneedgenerate)
+			endwork = 1;
+		pthread_mutex_unlock(&keysgenerated_mutex);
+	}
+
+	strcpy(&sname[onionendpos], "/hs_ed25519_secret_key");
+	fh = fopen(sname, "wb");
+	if (fh) {
+		fwrite(secret, skprefixlen + SECRET_LEN, 1, fh);
+		fclose(fh);
+	}
+
 	strcpy(&sname[onionendpos], "/hostname");
 	fh = fopen(sname, "w");
 	if (fh) {
@@ -140,21 +195,13 @@ static void onionready(char *sname, const u8 *secret, const u8 *pubonion)
 		fclose(fh);
 	}
 
-	strcpy(&sname[onionendpos], "/hs_ed25519_secret_key");
-	fh = fopen(sname, "wb");
-	if (fh) {
-		fwrite(secret, skprefixlen + SECRET_LEN, 1, fh);
-		fclose(fh);
-	}
-
 	sname[onionendpos] = '\n';
-	pthread_mutex_lock(&fout_mutex);
 	if (fout) {
+		pthread_mutex_lock(&fout_mutex);
 		fwrite(&sname[printstartpos], printlen, 1, fout);
 		fflush(fout);
+		pthread_mutex_unlock(&fout_mutex);
 	}
-	++keysgenerated;
-	pthread_mutex_unlock(&fout_mutex);
 }
 
 // little endian inc
@@ -185,7 +232,7 @@ static void *dowork(void *task)
 	memcpy(hashsrc,checksumstr,checksumstrlen);
 	hashsrc[checksumstrlen + PUBLIC_LEN] = 0x03; // version
 
-	sname = alloca(workdirlen + ONIONLEN + 63 + 1);
+	sname = malloc(workdirlen + ONIONLEN + 63 + 1);
 	if (workdir)
 		memcpy(sname,workdir,workdirlen);
 
@@ -215,6 +262,7 @@ again:
 	goto again;
 
 end:
+	free(sname);
 	return 0;
 }
 
@@ -274,7 +322,7 @@ static void *dofastwork(void *task)
 	memcpy(hashsrc, checksumstr, checksumstrlen);
 	hashsrc[checksumstrlen + PUBLIC_LEN] = 0x03; // version
 
-	sname = alloca(workdirlen + ONIONLEN + 63 + 1);
+	sname = malloc(workdirlen + ONIONLEN + 63 + 1);
 	if (workdir)
 		memcpy(sname, workdir, workdirlen);
 
@@ -327,6 +375,7 @@ initseed:
 	goto initseed;
 
 end:
+	free(sname);
 	return 0;
 }
 
@@ -392,6 +441,7 @@ int main(int argc, char **argv)
 	filters_init();
 	
 	fout = stdout;
+	pthread_mutex_init(&keysgenerated_mutex, 0);
 	pthread_mutex_init(&fout_mutex, 0);
 
 	const char *progname = argv[0];
@@ -467,6 +517,14 @@ int main(int argc, char **argv)
 					exit(1);
 				}
 			}
+			else if (*arg == 'n') {
+				if (argc--)
+					numneedgenerate = (size_t)atoll(*argv++);
+				else {
+					fprintf(stderr, "additional argument required\n");
+					exit(1);
+				}
+			}
 			else if (*arg == 'z')
 				fastkeygen = 1;
 			else {
@@ -478,7 +536,7 @@ int main(int argc, char **argv)
 		}
 		else filters_add(arg);
 	}
-	
+
 	if (outfile)
 		fout = fopen(outfile, "w");
 
@@ -487,10 +545,10 @@ int main(int argc, char **argv)
 
 	if (workdir)
 		mkdir(workdir, 0700);
-	
+
 	direndpos = workdirlen;
 	onionendpos = workdirlen + ONIONLEN;
-	
+
 	if (!dirnameflag) {
 		printstartpos = direndpos;
 		printlen = ONIONLEN + 1;
@@ -498,15 +556,18 @@ int main(int argc, char **argv)
 		printstartpos = 0;
 		printlen = onionendpos + 1;
 	}
-	
+
 	if (numthreads <= 0) {
 		// TODO: autodetect
 		numthreads = 1;
 	}
-	
+
+	signal(SIGTERM, termhandler);
+	signal(SIGINT, termhandler);
+
 	VEC_INIT(threads);
 	VEC_ADDN(threads, pthread_t, numthreads);
-	
+
 	for (size_t i = 0; i < VEC_LENGTH(threads); ++i) {
 		tret = pthread_create(&VEC_BUF(threads, i), 0, fastkeygen ? dofastwork : dowork, 0);
 		if (tret) {
@@ -514,14 +575,30 @@ int main(int argc, char **argv)
 			exit(1);
 		}
 	}
+
+	struct timespec ts;
+	memset(&ts,0,sizeof(ts));
+	ts.tv_nsec = 100000000;
+	while (!endwork) {
+		if (numneedgenerate && keysgenerated >= numneedgenerate) {
+			endwork = 1;
+			break;
+		}
+		nanosleep(&ts,0);
+	}
+
 	fprintf(stderr, "waiting for threads to finish...\n");
 	for (size_t i = 0; i < VEC_LENGTH(threads); ++i) {
 		pthread_join(VEC_BUF(threads, i), 0);
 	}
 	fprintf(stderr, "done, quitting\n");
 
+	pthread_mutex_destroy(&keysgenerated_mutex);
 	pthread_mutex_destroy(&fout_mutex);
 	filters_clean();
-	
+
+	if (outfile)
+		fclose(fout);
+
 	return 0;
 }