From 932ec2be8d01e553a768df3709182abf2b579097 Mon Sep 17 00:00:00 2001
From: Filippo Valsorda <filippo@golang.org>
Date: Thu, 9 Jan 2025 15:12:49 +0100
Subject: [PATCH] crypto/rsa: fix GenerateKey flakes for toy-sized keys

Could have fixed this some other ways, including inside the FIPS 140-3
module, but this is small and self-contained, clearly not affecting
production non-toy key sizes. This late in the freeze, a surgical fix
felt best.

Fixes #71185

Change-Id: I6a6a465641357c9d6b076c8a520b221be4210ed5
Reviewed-on: https://go-review.googlesource.com/c/go/+/641755
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Roland Shoemaker <roland@golang.org>
Auto-Submit: Filippo Valsorda <filippo@golang.org>
Reviewed-by: Russ Cox <rsc@golang.org>
---
 src/crypto/rsa/rsa.go      | 15 +++++++++++++++
 src/crypto/rsa/rsa_test.go | 17 +++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/src/crypto/rsa/rsa.go b/src/crypto/rsa/rsa.go
index fb23f003a6f217..95bb4becd2ff8c 100644
--- a/src/crypto/rsa/rsa.go
+++ b/src/crypto/rsa/rsa.go
@@ -327,6 +327,21 @@ func GenerateKey(random io.Reader, bits int) (*PrivateKey, error) {
 	}
 
 	k, err := rsa.GenerateKey(random, bits)
+	if bits < 256 && err != nil {
+		// Toy-sized keys have a non-negligible chance of hitting two hard
+		// failure cases: p == q and d <= 2^(nlen / 2).
+		//
+		// Since these are impossible to hit for real keys, we don't want to
+		// make the production code path more complex and harder to think about
+		// to handle them.
+		//
+		// Instead, just rerun the whole process a total of 8 times, which
+		// brings the chance of failure for 32-bit keys down to the same as for
+		// 256-bit keys.
+		for i := 1; i < 8 && err != nil; i++ {
+			k, err = rsa.GenerateKey(random, bits)
+		}
+	}
 	if err != nil {
 		return nil, err
 	}
diff --git a/src/crypto/rsa/rsa_test.go b/src/crypto/rsa/rsa_test.go
index 2535661040273a..73b0c3749eedb2 100644
--- a/src/crypto/rsa/rsa_test.go
+++ b/src/crypto/rsa/rsa_test.go
@@ -109,6 +109,23 @@ func TestImpossibleKeyGeneration(t *testing.T) {
 	}
 }
 
+func TestTinyKeyGeneration(t *testing.T) {
+	// Toy-sized keys can randomly hit hard failures in GenerateKey.
+	if testing.Short() {
+		t.Skip("skipping in short mode")
+	}
+	t.Setenv("GODEBUG", "rsa1024min=0")
+	for range 10000 {
+		k, err := GenerateKey(rand.Reader, 32)
+		if err != nil {
+			t.Fatalf("GenerateKey(32): %v", err)
+		}
+		if err := k.Validate(); err != nil {
+			t.Fatalf("Validate(32): %v", err)
+		}
+	}
+}
+
 func TestGnuTLSKey(t *testing.T) {
 	t.Setenv("GODEBUG", "rsa1024min=0")
 	// This is a key generated by `certtool --generate-privkey --bits 128`.