From 1d9e2936ac4c00b6227fa3c04e694049b104c67e Mon Sep 17 00:00:00 2001
From: gbrochar <gaetanbrochard@protonmail.com>
Date: Sun, 18 Feb 2024 17:14:48 +0100
Subject: [PATCH] opti: rm malloc in mul and useless cmp in modulo

---
 rsa/bigint.c | 53 +++++++++++++++++++++++++++++++++++-----------------
 1 file changed, 36 insertions(+), 17 deletions(-)

diff --git a/rsa/bigint.c b/rsa/bigint.c
index 58e2e19..4ecc0d6 100644
--- a/rsa/bigint.c
+++ b/rsa/bigint.c
@@ -106,11 +106,10 @@ void custom_bigint_modulo(bigint_t a, bigint_t b, bigint_t result, bigint_t mod,
 	if (bigint_cmp(result, b) < 0) {
 		return ;
 	}
-	bigint_bitwise_left_shift(mod);
+	while (bigint_cmp(result, mod) > 0) {
+		bigint_bitwise_left_shift(mod);
+	}
 	while (bigint_cmp(b, mod) < 0) {
-		while (bigint_cmp(result, mod) > 0) {
-			bigint_bitwise_left_shift(mod);
-		}
 		bigint_bitwise_right_shift(mod);
 		if (bigint_cmp(result, mod) > 0) {
 			bigint_substraction(result, mod, borrow_sub, y_sub);
@@ -175,17 +174,13 @@ void bigint_set_zeros(bigint_t n) {
 	}
 }
 
-void custom_bigint_mul(bigint_t a, bigint_t b, bigint_t result) {
+void custom_bigint_mul(bigint_t a, bigint_t b, bigint_t result, bigint_t *b_tool) {
 	int width = a.len * 32;
 	bigint_set_zeros(result);
-	bigint_t *b_tool = (bigint_t *)malloc(32 * sizeof(bigint_t));
-
-	b_tool[0] = bigint_zero(a.len + b.len);
 	bigint_set_zeros(b_tool[0]);
 	memcpy(b_tool[0].data, b.data, b.len * sizeof(uint32_t));
 
 	for (int i = 1; i < 32; i++) {
-		b_tool[i] = bigint_zero(a.len + b.len);
 		bigint_set_zeros(b_tool[i]);
 		memcpy(b_tool[i].data, b_tool[i - 1].data, b.len * sizeof(uint32_t));
 		bigint_bitwise_left_shift(b_tool[i]);
@@ -198,14 +193,11 @@ void custom_bigint_mul(bigint_t a, bigint_t b, bigint_t result) {
 			custom_bigint_add(result, b_tool[offset], index);
 		}
 	}
-	for (int i = 0; i < 32; i++) {
-		bigint_destroy(b_tool[i]);
-	}
 }
 
 // a^e mod n
 // clean memory tricks !!!
-void custom_bigint_pow_mod(bigint_t a, bigint_t e, bigint_t n, bigint_t result, bigint_t custom, bigint_t custom2, bigint_t mod, bigint_t borrow_sub, bigint_t y_sub) {
+void custom_bigint_pow_mod(bigint_t a, bigint_t e, bigint_t n, bigint_t result, bigint_t custom, bigint_t custom2, bigint_t mod, bigint_t borrow_sub, bigint_t y_sub, bigint_t *b_tool) {
 	bigint_set_zeros(result);
 	bigint_set_zeros(custom);
 	bigint_set_zeros(custom2);
@@ -217,12 +209,12 @@ void custom_bigint_pow_mod(bigint_t a, bigint_t e, bigint_t n, bigint_t result,
 	}
 	cursor--;
 	while (cursor >= 0) {
-		custom_bigint_mul(result, result, custom);
+		custom_bigint_mul(result, result, custom, b_tool);
 		custom_bigint_modulo(custom, n, custom2, mod, borrow_sub, y_sub);
 		bigint_set_zeros(result);
 		memcpy(result.data, custom2.data, custom2.len * sizeof(uint32_t));
 		if (e.data[cursor / 32] & 1 << (cursor % 32)) {
-			custom_bigint_mul(result, a, custom);
+			custom_bigint_mul(result, a, custom, b_tool);
 			custom_bigint_modulo(custom, n, custom2, mod, borrow_sub, y_sub);
 			memcpy(result.data, custom2.data, custom2.len * sizeof(uint32_t));
 		}
@@ -269,6 +261,12 @@ bigint_t bigint_prime(size_t len) {
 	bigint_t y = bigint_zero(RSA_BLOCK_SIZE / 8 / sizeof(uint32_t) * 2);
 	bigint_t custom = bigint_zero(RSA_BLOCK_SIZE / 8 / sizeof(uint32_t) * 2);
 	bigint_t custom2 = bigint_zero(RSA_BLOCK_SIZE / 8 / sizeof(uint32_t) * 2);
+	bigint_t *b_tool = (bigint_t *)protected_malloc(32 * sizeof(bigint_t));
+
+	for (int i = 0; i < 32; i++) {
+		b_tool[i] = bigint_zero(RSA_BLOCK_SIZE / 8 / sizeof(uint32_t) * 4);
+	}
+
 
 	bigint_t two = bigint_zero(RSA_BLOCK_SIZE / 8 / sizeof(uint32_t) * 2);
 	bigint_t one = bigint_zero(RSA_BLOCK_SIZE / 8 / sizeof(uint32_t) * 2);
@@ -285,14 +283,21 @@ bigint_t bigint_prime(size_t len) {
 		while (bigint_cmp(a, two) < 0 || bigint_cmp(a, n_minus_two) > 0) {
 			bigint_set_random_bytes(a, len);
 		}
-		custom_bigint_pow_mod(a, d, n, x, custom, custom2, mod, borrow_sub, y_sub);
+		custom_bigint_pow_mod(a, d, n, x, custom, custom2, mod, borrow_sub, y_sub, b_tool);
 		for (uint32_t i = 0; i < s; i++) {
-			custom_bigint_pow_mod(x, two, n, y, custom, custom2, mod, borrow_sub, y_sub);
+			custom_bigint_pow_mod(x, two, n, y, custom, custom2, mod, borrow_sub, y_sub, b_tool);
 			if (!bigint_dif(y, one) && bigint_dif(x, one) && bigint_dif(x, n_minus_one)) {
 				bulk_destroy(x, y, n, d, two, one, n_minus_two, n_minus_one);
 				bigint_destroy(custom);
 				bigint_destroy(custom2);
 				bigint_destroy(a);
+				bigint_destroy(mod);
+				bigint_destroy(borrow_sub);
+				bigint_destroy(y_sub);
+				for (int i = 0; i < 32; i++) {
+					bigint_destroy(b_tool[i]);
+				}
+				free(b_tool);
 				return bigint_prime(len);
 			}
 			bigint_destroy(x);
@@ -303,12 +308,26 @@ bigint_t bigint_prime(size_t len) {
 			bigint_destroy(custom);
 			bigint_destroy(custom2);
 			bigint_destroy(a);
+			bigint_destroy(mod);
+			bigint_destroy(borrow_sub);
+			bigint_destroy(y_sub);
+			for (int i = 0; i < 32; i++) {
+				bigint_destroy(b_tool[i]);
+			}
+			free(b_tool);
 			return bigint_prime(len);
 		}
 	}
 	bulk_destroy(x, y, custom, d, two, one, n_minus_two, n_minus_one);
 	bigint_destroy(custom2);
 	bigint_destroy(a);
+	bigint_destroy(mod);
+	bigint_destroy(borrow_sub);
+	bigint_destroy(y_sub);
+	for (int i = 0; i < 32; i++) {
+		bigint_destroy(b_tool[i]);
+	}
+	free(b_tool);
 	return n;
 }