From bdb4f47532e7362ca66f666189ea32b20fd2a4b5 Mon Sep 17 00:00:00 2001
From: Kevaundray Wedderburn <kevtheappdev@gmail.com>
Date: Fri, 11 Jul 2025 12:39:19 +0100
Subject: [PATCH] opt: strip leading zeroes

---
 crypto/modexp/gmp/gmp.go      | 31 ++++++++++++---
 crypto/modexp/gmp/gmp_test.go | 75 +++++++++++++++++++++++++++++++++++
 2 files changed, 100 insertions(+), 6 deletions(-)

diff --git a/crypto/modexp/gmp/gmp.go b/crypto/modexp/gmp/gmp.go
index c40b83e89f..d39258e7dc 100644
--- a/crypto/modexp/gmp/gmp.go
+++ b/crypto/modexp/gmp/gmp.go
@@ -25,6 +25,18 @@ import (
 	"unsafe"
 )
 
+// stripLeadingZeros removes leading zero bytes from a slice
+// Note: This has no effect for big-endian integers
+func stripLeadingZeros(data []byte) []byte {
+	for i, b := range data {
+		if b != 0 {
+			return data[i:]
+		}
+	}
+	// All zeros, return empty slice
+	return []byte{}
+}
+
 // ModExp performs modular exponentiation using GMP
 // This is thread safe.
 func ModExp(base, exp, mod []byte) ([]byte, error) {
@@ -56,8 +68,15 @@ func ModExp(base, exp, mod []byte) ([]byte, error) {
 		return []byte{1}, nil
 	}
 
-	// Allocate result buffer (size of modulus is the max possible result)
-	result := make([]byte, len(mod))
+	// Strip leading zeros for GMP performance
+	base = stripLeadingZeros(base)
+	exp = stripLeadingZeros(exp)
+	modStripped := stripLeadingZeros(mod)
+
+	// Allocate result buffer (size of stripped modulus is the max possible result)
+	// Note: We know that the modulus stripped is non-zero because
+	// we check for all zeroes.
+	result := make([]byte, len(modStripped))
 	resultLen := C.size_t(len(result))
 
 	// Handle empty slices - pass a dummy non-nil pointer with length 0
@@ -71,22 +90,22 @@ func ModExp(base, exp, mod []byte) ([]byte, error) {
 	if len(exp) > 0 {
 		expPtr = (*C.uint8_t)(unsafe.Pointer(&exp[0]))
 	}
-	if len(mod) > 0 {
-		modPtr = (*C.uint8_t)(unsafe.Pointer(&mod[0]))
+	if len(modStripped) > 0 {
+		modPtr = (*C.uint8_t)(unsafe.Pointer(&modStripped[0]))
 	}
 
 	// Call C function
 	ret := C.modexp_bytes(
 		basePtr, C.size_t(len(base)),
 		expPtr, C.size_t(len(exp)),
-		modPtr, C.size_t(len(mod)),
+		modPtr, C.size_t(len(modStripped)),
 		(*C.uint8_t)(unsafe.Pointer(&result[0])), &resultLen,
 	)
 
 	// Keep the slices alive until after the C call completes
 	runtime.KeepAlive(base)
 	runtime.KeepAlive(exp)
-	runtime.KeepAlive(mod)
+	runtime.KeepAlive(modStripped)
 	runtime.KeepAlive(result)
 
 	// Check for errors
diff --git a/crypto/modexp/gmp/gmp_test.go b/crypto/modexp/gmp/gmp_test.go
index c6e479cf02..5a918471df 100644
--- a/crypto/modexp/gmp/gmp_test.go
+++ b/crypto/modexp/gmp/gmp_test.go
@@ -134,6 +134,81 @@ func BenchmarkModExpBigInt(b *testing.B) {
 	}
 }
 
+// TestLeadingZeros tests that leading zeros are handled correctly
+func TestLeadingZeros(t *testing.T) {
+	tests := []struct {
+		name     string
+		base     []byte
+		exp      []byte
+		mod      []byte
+		expected []byte
+	}{
+		{
+			name:     "base_with_leading_zeros",
+			base:     []byte{0, 0, 0, 0, 2},
+			exp:      []byte{3},
+			mod:      []byte{7},
+			expected: []byte{1}, // 2^3 mod 7 = 8 mod 7 = 1
+		},
+		{
+			name:     "exp_with_leading_zeros",
+			base:     []byte{2},
+			exp:      []byte{0, 0, 0, 3},
+			mod:      []byte{7},
+			expected: []byte{1}, // 2^3 mod 7 = 8 mod 7 = 1
+		},
+		{
+			name:     "mod_with_leading_zeros",
+			base:     []byte{2},
+			exp:      []byte{3},
+			mod:      []byte{0, 0, 0, 0, 7},
+			expected: []byte{1}, // 2^3 mod 7 = 8 mod 7 = 1
+		},
+		{
+			name:     "all_with_leading_zeros",
+			base:     []byte{0, 0, 2},
+			exp:      []byte{0, 0, 3},
+			mod:      []byte{0, 0, 7},
+			expected: []byte{1}, // 2^3 mod 7 = 8 mod 7 = 1
+		},
+		{
+			name:     "base_all_zeros",
+			base:     []byte{0, 0, 0},
+			exp:      []byte{5},
+			mod:      []byte{7},
+			expected: []byte{}, // 0^5 mod 7 = 0
+		},
+		{
+			name:     "exp_all_zeros",
+			base:     []byte{5},
+			exp:      []byte{0, 0, 0},
+			mod:      []byte{7},
+			expected: []byte{1}, // 5^0 mod 7 = 1
+		},
+		{
+			name:     "base_and_exp_all_zeros",
+			base:     []byte{0, 0},
+			exp:      []byte{0, 0},
+			mod:      []byte{7},
+			expected: []byte{1}, // 0^0 mod 7 = 1 (by convention)
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result, err := ModExp(tt.base, tt.exp, tt.mod)
+			if err != nil {
+				t.Fatalf("ModExp error: %v", err)
+			}
+
+			if !bytes.Equal(result, tt.expected) {
+				t.Errorf("Results differ:\nGot:      %x\nExpected: %x",
+					result, tt.expected)
+			}
+		})
+	}
+}
+
 // TestSpecialCases tests the special case optimizations
 func TestSpecialCases(t *testing.T) {
 	tests := []struct {