Move PrecomputeRSATestKeys out of production code (#57049)

Introduces a new cryptosuitestest package to be used for generating
RSA keys in tests. This new package is a test only package and
exists for the sole purpose of removing the testing package from
the cryptosuites package. Once teleport.e has been updated the
legacy way to generate test keys will be removed.

Updates #51023.

.golangci.yml

@@ -179,6 +179,7 @@ linters:
             - '!**/e/lib/idp/saml/testenv/**'
             - '!**/e/lib/operatortest/**'
             - '!**/e/tests/**'
+            - '!**/integration/helpers/**'
           deny:
             - pkg: github.com/gravitational/teleport/e/lib/aws/identitycenter/test
               desc: testing packages should not be imported outside of _test.go files
@@ -208,6 +209,8 @@ linters:
               desc: testing packages should not be imported outside of _test.go files
             - pkg: github.com/gravitational/teleport/lib/auth/authtest
               desc: testing packages should not be imported outside of _test.go files
+            - pkg: github.com/gravitational/teleport/lib/cryptosuites/cryptosuitestest
+              desc: testing packages should not be imported outside of _test.go files
         testify:
           files:
             - '!$test'

integration/helpers/testmain.go

@@ -19,11 +19,12 @@
 package helpers
 
 import (
+	"context"
 	"os"
 	"testing"
 	"time"
 
-	"github.com/gravitational/teleport/lib/cryptosuites"
+	"github.com/gravitational/teleport/lib/cryptosuites/cryptosuitestest"
 	"github.com/gravitational/teleport/lib/modules"
 	"github.com/gravitational/teleport/lib/srv"
 	"github.com/gravitational/teleport/lib/utils/log/logtest"
@@ -34,17 +35,21 @@ import (
 // it as an argument. Otherwise, it will run tests as normal.
 func TestMainImplementation(m *testing.M) {
 	logtest.InitLogger(testing.Verbose)
-	cryptosuites.PrecomputeRSATestKeys(m)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cryptosuitestest.PrecomputeRSAKeys(ctx)
 	SetTestTimeouts(3 * time.Second)
 	modules.SetInsecureTestMode(true)
 	// If the test is re-executing itself, execute the command that comes over
 	// the pipe.
 	if srv.IsReexec() {
+		defer cancel()
 		common.Run(common.Options{Args: os.Args[1:]})
 		return
 	}
 
 	// Otherwise run tests as normal.
-	code := m.Run()
-	os.Exit(code)
+	exitCode := m.Run()
+	cancel()
+	os.Exit(exitCode)
 }

lib/auth/auth_test.go

@@ -76,6 +76,7 @@ import (
 	"github.com/gravitational/teleport/lib/backend"
 	"github.com/gravitational/teleport/lib/backend/memory"
 	"github.com/gravitational/teleport/lib/cryptosuites"
+	"github.com/gravitational/teleport/lib/cryptosuites/cryptosuitestest"
 	"github.com/gravitational/teleport/lib/defaults"
 	"github.com/gravitational/teleport/lib/events"
 	"github.com/gravitational/teleport/lib/events/eventstest"
@@ -232,9 +233,13 @@ func newAuthSuite(t *testing.T) *testPack {
 
 func TestMain(m *testing.M) {
 	logtest.InitLogger(testing.Verbose)
-	cryptosuites.PrecomputeRSATestKeys(m)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cryptosuitestest.PrecomputeRSAKeys(ctx)
 	modules.SetInsecureTestMode(true)
-	os.Exit(m.Run())
+	exitCode := m.Run()
+	cancel()
+	os.Exit(exitCode)
 }
 
 func TestSessions(t *testing.T) {

lib/client/api_test.go

@@ -44,7 +44,7 @@ import (
 	"github.com/gravitational/teleport/api/utils/grpc/interceptors"
 	"github.com/gravitational/teleport/api/utils/keys"
 	"github.com/gravitational/teleport/lib/auth/authclient"
-	"github.com/gravitational/teleport/lib/cryptosuites"
+	"github.com/gravitational/teleport/lib/cryptosuites/cryptosuitestest"
 	"github.com/gravitational/teleport/lib/defaults"
 	"github.com/gravitational/teleport/lib/modules"
 	"github.com/gravitational/teleport/lib/observability/tracing"
@@ -55,8 +55,12 @@ import (
 func TestMain(m *testing.M) {
 	logtest.InitLogger(testing.Verbose)
 	modules.SetInsecureTestMode(true)
-	cryptosuites.PrecomputeRSATestKeys(m)
-	os.Exit(m.Run())
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cryptosuitestest.PrecomputeRSAKeys(ctx)
+	exitCode := m.Run()
+	cancel()
+	os.Exit(exitCode)
 }
 
 var parseProxyHostTestCases = []struct {

lib/cryptosuites/cryptosuitestest/precompute.go

@@ -0,0 +1,90 @@
+// Teleport
+// Copyright (C) 2025 Gravitational, Inc.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+package cryptosuitestest
+
+import (
+	"context"
+	"crypto/rand"
+	"crypto/rsa"
+
+	"github.com/gravitational/teleport/api/constants"
+	internalrsa "github.com/gravitational/teleport/lib/cryptosuites/internal/rsa"
+)
+
+// PrecomputeRSATestKeys may be called from TestMain to set this package into a
+// mode where it will precompute a fixed number of RSA keys and reuse them to
+// save on CPU usage.
+func PrecomputeRSAKeys(ctx context.Context) {
+	internalrsa.StartPrecomputeOnce.Do(func() {
+		go precomputeTestKeys(ctx)
+	})
+}
+
+func precomputeTestKeys(ctx context.Context) {
+	generatedTestKeys := generateTestKeys(ctx)
+	keysToReuse := make([]*rsa.PrivateKey, 0, testKeysNumber)
+	for range testKeysNumber {
+		select {
+		case k := <-generatedTestKeys:
+			internalrsa.PrecomputedKeys <- k
+			keysToReuse = append(keysToReuse, k)
+		case <-ctx.Done():
+			return
+		}
+	}
+
+	for {
+		for _, k := range keysToReuse {
+			select {
+			case internalrsa.PrecomputedKeys <- k:
+			case <-ctx.Done():
+				return
+			}
+		}
+	}
+}
+
+// testKeysNumber is the number of RSA keys generated in tests.
+const testKeysNumber = 25
+
+func generateTestKeys(ctx context.Context) <-chan *rsa.PrivateKey {
+	generatedTestKeys := make(chan *rsa.PrivateKey, testKeysNumber)
+	for range testKeysNumber {
+		// Generate each key in a separate goroutine to take advantage of
+		// multiple cores if possible.
+		go func() {
+			private, err := generateRSAPrivateKey()
+			if err != nil {
+				// Use only in tests. Safe to panic.
+				panic(err)
+			}
+
+			select {
+			case generatedTestKeys <- private:
+			case <-ctx.Done():
+				return
+			}
+
+		}()
+	}
+	return generatedTestKeys
+}
+
+func generateRSAPrivateKey() (*rsa.PrivateKey, error) {
+	//nolint:forbidigo // This is the one function allowed to generate RSA keys.
+	return rsa.GenerateKey(rand.Reader, constants.RSAKeySize)
+}

lib/cryptosuites/internal/rsa/rsa.go

@@ -27,15 +27,18 @@ import (
 
 	"github.com/gravitational/teleport"
 	"github.com/gravitational/teleport/api/constants"
+	logutils "github.com/gravitational/teleport/lib/utils/log"
 )
 
-var log = slog.With(teleport.ComponentKey, teleport.ComponentKeyGen)
+var (
+	log = logutils.NewPackageLogger(teleport.ComponentKey, teleport.ComponentKeyGen)
 
-// precomputedKeys is a queue of cached keys ready for usage.
-var precomputedKeys = make(chan *rsa.PrivateKey, 25)
+	// PrecomputedKeys is a queue of cached keys ready for usage.
+	PrecomputedKeys = make(chan *rsa.PrivateKey, 25)
 
-// startPrecomputeOnce is used to start the background task that precomputes key pairs.
-var startPrecomputeOnce sync.Once
+	// StartPrecomputeOnce is used to start the background task that precomputes key pairs.
+	StartPrecomputeOnce sync.Once
+)
 
 // GenerateKey returns a newly generated RSA private key.
 func GenerateKey() (*rsa.PrivateKey, error) {
@@ -44,7 +47,7 @@ func GenerateKey() (*rsa.PrivateKey, error) {
 
 func getOrGenerateRSAPrivateKey() (*rsa.PrivateKey, error) {
 	select {
-	case k := <-precomputedKeys:
+	case k := <-PrecomputedKeys:
 		return k, nil
 	default:
 		rsaKeyPair, err := generateRSAPrivateKey()
@@ -70,7 +73,7 @@ func precomputeKeys() {
 			time.Sleep(backoff)
 		}
 
-		precomputedKeys <- rsaPrivateKey
+		PrecomputedKeys <- rsaPrivateKey
 	}
 }
 
@@ -79,12 +82,12 @@ func precomputeTestKeys() {
 	keysToReuse := make([]*rsa.PrivateKey, 0, testKeysNumber)
 	for range testKeysNumber {
 		k := <-generatedTestKeys
-		precomputedKeys <- k
+		PrecomputedKeys <- k
 		keysToReuse = append(keysToReuse, k)
 	}
 	for {
 		for _, k := range keysToReuse {
-			precomputedKeys <- k
+			PrecomputedKeys <- k
 		}
 	}
 }
@@ -113,7 +116,7 @@ func generateTestKeys() <-chan *rsa.PrivateKey {
 // computed in advance. This should only be enabled if large spikes in key computation
 // are expected (e.g. in auth/proxy services). Safe to double-call.
 func PrecomputeKeys() {
-	startPrecomputeOnce.Do(func() {
+	StartPrecomputeOnce.Do(func() {
 		go precomputeKeys()
 	})
 }
@@ -122,8 +125,10 @@ func PrecomputeKeys() {
 // reduce CPU usage. This method should only be in tests. Safe to call multiple
 // times. This function takes *testing.M so it can only be used from TestMain in
 // tests.
+// Deprecated: prefer using cyptosuitestest.PrecomputeRSAKeys instead
+// TODO(tross): Delete once references in teleport.e are gone.
 func PrecomputeTestKeys(_ *testing.M) {
-	startPrecomputeOnce.Do(func() {
+	StartPrecomputeOnce.Do(func() {
 		go precomputeTestKeys()
 	})
 }

lib/cryptosuites/internal/rsa/rsa_test.go

@@ -29,7 +29,7 @@ func TestPrecomputeMode(t *testing.T) {
 	PrecomputeKeys()
 
 	select {
-	case <-precomputedKeys:
+	case <-PrecomputedKeys:
 	case <-time.After(time.Second * 10):
 		t.Fatal("Key precompute routine failed to start.")
 	}

lib/cryptosuites/precompute.go

@@ -33,6 +33,8 @@ func PrecomputeRSAKeys() {
 // PrecomputeRSATestKeys may be called from TestMain to set this package into a
 // mode where it will precompute a fixed number of RSA keys and reuse them to
 // save on CPU usage.
+// Deprecated: prefer using cryptosuitestest.PrecomputeRSAKeys
+// TODO(tross): Delete once teleport.e no longer references it.
 func PrecomputeRSATestKeys(m *testing.M) {
 	rsa.PrecomputeTestKeys(m)
 }

lib/srv/app/server_test.go

@@ -60,6 +60,7 @@ import (
 	"github.com/gravitational/teleport/lib/cloud/awsconfig"
 	"github.com/gravitational/teleport/lib/cloud/mocks"
 	"github.com/gravitational/teleport/lib/cryptosuites"
+	"github.com/gravitational/teleport/lib/cryptosuites/cryptosuitestest"
 	"github.com/gravitational/teleport/lib/events"
 	"github.com/gravitational/teleport/lib/httplib/reverseproxy"
 	"github.com/gravitational/teleport/lib/inventory"
@@ -76,9 +77,12 @@ import (
 
 func TestMain(m *testing.M) {
 	logtest.InitLogger(testing.Verbose)
-	cryptosuites.PrecomputeRSATestKeys(m)
+	ctx, cancel := context.WithCancel(context.Background())
+	cryptosuitestest.PrecomputeRSAKeys(ctx)
 	modules.SetInsecureTestMode(true)
-	os.Exit(m.Run())
+	exitCode := m.Run()
+	cancel()
+	os.Exit(exitCode)
 }
 
 type Suite struct {

lib/srv/db/access_test.go

@@ -65,7 +65,7 @@ import (
 	clients "github.com/gravitational/teleport/lib/cloud"
 	"github.com/gravitational/teleport/lib/cloud/awsconfig"
 	"github.com/gravitational/teleport/lib/cloud/mocks"
-	"github.com/gravitational/teleport/lib/cryptosuites"
+	"github.com/gravitational/teleport/lib/cryptosuites/cryptosuitestest"
 	"github.com/gravitational/teleport/lib/defaults"
 	libevents "github.com/gravitational/teleport/lib/events"
 	"github.com/gravitational/teleport/lib/events/eventstest"
@@ -108,12 +108,16 @@ import (
 
 func TestMain(m *testing.M) {
 	logtest.InitLogger(testing.Verbose)
-	cryptosuites.PrecomputeRSATestKeys(m)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cryptosuitestest.PrecomputeRSAKeys(ctx)
 	modules.SetInsecureTestMode(true)
 	registerTestSnowflakeEngine()
 	registerTestElasticsearchEngine()
 	registerTestSQLServerEngine()
-	os.Exit(m.Run())
+	exitCode := m.Run()
+	cancel()
+	os.Exit(exitCode)
 }
 
 // TestAccessPostgres verifies access scenarios to a Postgres database based

lib/tbot/tbot_test.go

@@ -54,7 +54,7 @@ import (
 	"github.com/gravitational/teleport/api/utils/sshutils"
 	"github.com/gravitational/teleport/integration/helpers"
 	"github.com/gravitational/teleport/lib/auth/authclient"
-	"github.com/gravitational/teleport/lib/cryptosuites"
+	"github.com/gravitational/teleport/lib/cryptosuites/cryptosuitestest"
 	"github.com/gravitational/teleport/lib/defaults"
 	"github.com/gravitational/teleport/lib/service"
 	"github.com/gravitational/teleport/lib/service/servicecfg"
@@ -73,8 +73,12 @@ import (
 
 func TestMain(m *testing.M) {
 	logtest.InitLogger(testing.Verbose)
-	cryptosuites.PrecomputeRSATestKeys(m)
-	os.Exit(m.Run())
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cryptosuitestest.PrecomputeRSAKeys(ctx)
+	exitCode := m.Run()
+	cancel()
+	os.Exit(exitCode)
 }
 
 type defaultBotConfigOpts struct {

lib/web/apiserver_test.go

@@ -113,6 +113,7 @@ import (
 	dbrepl "github.com/gravitational/teleport/lib/client/db/repl"
 	"github.com/gravitational/teleport/lib/client/sso"
 	"github.com/gravitational/teleport/lib/cryptosuites"
+	"github.com/gravitational/teleport/lib/cryptosuites/cryptosuitestest"
 	"github.com/gravitational/teleport/lib/defaults"
 	"github.com/gravitational/teleport/lib/events"
 	"github.com/gravitational/teleport/lib/events/eventstest"
@@ -183,10 +184,13 @@ func TestMain(m *testing.M) {
 		srv.RunAndExit(os.Args[1])
 		return
 	}
-	cryptosuites.PrecomputeRSATestKeys(m)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cryptosuitestest.PrecomputeRSAKeys(ctx)
 
 	// Otherwise run tests as normal.
 	code := m.Run()
+	cancel()
 	os.Exit(code)
 }
 

tool/tsh/common/tsh_test.go

@@ -84,6 +84,7 @@ import (
 	"github.com/gravitational/teleport/lib/client"
 	"github.com/gravitational/teleport/lib/client/identityfile"
 	"github.com/gravitational/teleport/lib/cryptosuites"
+	"github.com/gravitational/teleport/lib/cryptosuites/cryptosuitestest"
 	"github.com/gravitational/teleport/lib/defaults"
 	"github.com/gravitational/teleport/lib/kube/kubeconfig"
 	"github.com/gravitational/teleport/lib/modules"
@@ -133,8 +134,12 @@ func TestMain(m *testing.M) {
 	modules.SetInsecureTestMode(true)
 
 	logtest.InitLogger(testing.Verbose)
-	cryptosuites.PrecomputeRSATestKeys(m)
-	os.Exit(m.Run())
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cryptosuitestest.PrecomputeRSAKeys(ctx)
+	exitCode := m.Run()
+	cancel()
+	os.Exit(exitCode)
 }
 
 func BenchmarkInit(b *testing.B) {