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) {