From f16eebb2f40255a49d92eac1b4286acf3ddfcdb0 Mon Sep 17 00:00:00 2001 From: rosstimothy <39066650+rosstimothy@users.noreply.github.com> Date: Wed, 23 Jul 2025 09:23:40 -0400 Subject: [PATCH] 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 | 3 + integration/helpers/testmain.go | 13 ++- lib/auth/auth_test.go | 9 +- lib/client/api_test.go | 10 ++- .../cryptosuitestest/precompute.go | 90 +++++++++++++++++++ lib/cryptosuites/internal/rsa/rsa.go | 27 +++--- lib/cryptosuites/internal/rsa/rsa_test.go | 2 +- lib/cryptosuites/precompute.go | 2 + lib/srv/app/server_test.go | 8 +- lib/srv/db/access_test.go | 10 ++- lib/tbot/tbot_test.go | 10 ++- lib/web/apiserver_test.go | 6 +- tool/tsh/common/tsh_test.go | 9 +- 13 files changed, 167 insertions(+), 32 deletions(-) create mode 100644 lib/cryptosuites/cryptosuitestest/precompute.go diff --git a/.golangci.yml b/.golangci.yml index a748b6b7e0e..4770889daba 100644 --- a/.golangci.yml +++ b/.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' diff --git a/integration/helpers/testmain.go b/integration/helpers/testmain.go index b6f6e57851f..bf0d84a3ee5 100644 --- a/integration/helpers/testmain.go +++ b/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) } diff --git a/lib/auth/auth_test.go b/lib/auth/auth_test.go index 7f108f404cb..d3da8479f07 100644 --- a/lib/auth/auth_test.go +++ b/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) { diff --git a/lib/client/api_test.go b/lib/client/api_test.go index 0cfed2171c4..a5f294b58e9 100644 --- a/lib/client/api_test.go +++ b/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 { diff --git a/lib/cryptosuites/cryptosuitestest/precompute.go b/lib/cryptosuites/cryptosuitestest/precompute.go new file mode 100644 index 00000000000..7b2f15aeacf --- /dev/null +++ b/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 . + +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) +} diff --git a/lib/cryptosuites/internal/rsa/rsa.go b/lib/cryptosuites/internal/rsa/rsa.go index 2203bd22bf0..b0dddf883f1 100644 --- a/lib/cryptosuites/internal/rsa/rsa.go +++ b/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() }) } diff --git a/lib/cryptosuites/internal/rsa/rsa_test.go b/lib/cryptosuites/internal/rsa/rsa_test.go index 15bbf508f37..fb84cd5fb17 100644 --- a/lib/cryptosuites/internal/rsa/rsa_test.go +++ b/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.") } diff --git a/lib/cryptosuites/precompute.go b/lib/cryptosuites/precompute.go index 6e25378f9ea..042dff608ef 100644 --- a/lib/cryptosuites/precompute.go +++ b/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) } diff --git a/lib/srv/app/server_test.go b/lib/srv/app/server_test.go index f07ee085809..01f303c833c 100644 --- a/lib/srv/app/server_test.go +++ b/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 { diff --git a/lib/srv/db/access_test.go b/lib/srv/db/access_test.go index 4bf95e97cdf..dc458528e13 100644 --- a/lib/srv/db/access_test.go +++ b/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 diff --git a/lib/tbot/tbot_test.go b/lib/tbot/tbot_test.go index 9d35b0c663d..a175e0aa644 100644 --- a/lib/tbot/tbot_test.go +++ b/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 { diff --git a/lib/web/apiserver_test.go b/lib/web/apiserver_test.go index 6f936a58e6b..5d63f949e26 100644 --- a/lib/web/apiserver_test.go +++ b/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) } diff --git a/tool/tsh/common/tsh_test.go b/tool/tsh/common/tsh_test.go index e8ca47c27b7..5fdb1b20808 100644 --- a/tool/tsh/common/tsh_test.go +++ b/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) {