From 2e879e0043cd244a6adafbc756aab729a31e7d34 Mon Sep 17 00:00:00 2001
From: Nic Klaassen <nic@goteleport.com>
Date: Thu, 5 Sep 2024 21:38:26 -0700
Subject: [PATCH] lazily PrecomputeKeys everywhere

---
 lib/auth/auth.go           |  3 ---
 lib/auth/sessions.go       | 10 ++++++++++
 lib/reversetunnel/cache.go |  9 ++++-----
 lib/service/service.go     |  8 --------
 4 files changed, 14 insertions(+), 16 deletions(-)

diff --git a/lib/auth/auth.go b/lib/auth/auth.go
index 5a10eb2a9e92e..d6685bbd1a010 100644
--- a/lib/auth/auth.go
+++ b/lib/auth/auth.go
@@ -82,7 +82,6 @@ import (
 	"github.com/gravitational/teleport/entitlements"
 	"github.com/gravitational/teleport/lib/auth/authclient"
 	"github.com/gravitational/teleport/lib/auth/keystore"
-	"github.com/gravitational/teleport/lib/auth/native"
 	"github.com/gravitational/teleport/lib/auth/userloginstate"
 	wanlib "github.com/gravitational/teleport/lib/auth/webauthn"
 	wantypes "github.com/gravitational/teleport/lib/auth/webauthntypes"
@@ -375,8 +374,6 @@ func NewServer(cfg *InitConfig, opts ...ServerOption) (*Server, error) {
 		if !modules.GetModules().Features().GetEntitlement(entitlements.HSM).Enabled {
 			return nil, fmt.Errorf("AWS KMS support requires a license with the HSM feature enabled: %w", ErrRequiresEnterprise)
 		}
-	} else {
-		native.PrecomputeKeys()
 	}
 	keyStore, err := keystore.NewManager(context.Background(), &cfg.KeyStoreConfig, keystoreOpts)
 	if err != nil {
diff --git a/lib/auth/sessions.go b/lib/auth/sessions.go
index c84160c19cf40..94b9e2487e887 100644
--- a/lib/auth/sessions.go
+++ b/lib/auth/sessions.go
@@ -21,6 +21,7 @@ package auth
 import (
 	"context"
 	"crypto"
+	"crypto/rsa"
 	"time"
 
 	"github.com/gravitational/trace"
@@ -36,6 +37,7 @@ import (
 	apievents "github.com/gravitational/teleport/api/types/events"
 	"github.com/gravitational/teleport/api/utils/keys"
 	"github.com/gravitational/teleport/entitlements"
+	"github.com/gravitational/teleport/lib/auth/native"
 	"github.com/gravitational/teleport/lib/cryptosuites"
 	"github.com/gravitational/teleport/lib/defaults"
 	dtconfig "github.com/gravitational/teleport/lib/devicetrust/config"
@@ -253,6 +255,14 @@ func (a *Server) newWebSession(
 		if err != nil {
 			return nil, nil, trace.Wrap(err)
 		}
+		if _, isRSA := sshKey.Public().(*rsa.PublicKey); isRSA {
+			// Ensure the native package is precomputing RSA keys if we ever
+			// generate one. [native.PrecomputeKeys] is idempotent.
+			// Doing this lazily easily handles changing signature algorithm
+			// suites and won't start precomputing keys if they are never needed
+			// (a major benefit in tests).
+			native.PrecomputeKeys()
+		}
 	}
 
 	sessionTTL := req.SessionTTL
diff --git a/lib/reversetunnel/cache.go b/lib/reversetunnel/cache.go
index d0a882a71ffd9..a30a70b3c2843 100644
--- a/lib/reversetunnel/cache.go
+++ b/lib/reversetunnel/cache.go
@@ -142,11 +142,10 @@ func (c *certificateCache) generateHostCert(ctx context.Context, principals []st
 
 	if _, isRSA := hostKey.Public().(*rsa.PublicKey); isRSA {
 		// Ensure the native package is precomputing RSA keys if we ever
-		// generate one. [native.PrecomputeKeys] is idempotent. Do this here
-		// instead of in newHostCertificateCache for 2 reasons:
-		// 1. This will handle changes to the configured algorithm suite.
-		// 2. This won't start precomputing keys if a host key is never actually
-		//    needed, which can be a major benefit in tests.
+		// generate one. [native.PrecomputeKeys] is idempotent.
+		// Doing this lazily easily handles changing signature algorithm suites
+		// and won't start precomputing keys if they are never needed (a major
+		// benefit in tests).
 		native.PrecomputeKeys()
 	}
 
diff --git a/lib/service/service.go b/lib/service/service.go
index 86c728de539d0..72eb7c7ae4d58 100644
--- a/lib/service/service.go
+++ b/lib/service/service.go
@@ -92,7 +92,6 @@ import (
 	"github.com/gravitational/teleport/lib/auth/authclient"
 	"github.com/gravitational/teleport/lib/auth/keygen"
 	"github.com/gravitational/teleport/lib/auth/machineid/machineidv1"
-	"github.com/gravitational/teleport/lib/auth/native"
 	"github.com/gravitational/teleport/lib/auth/state"
 	"github.com/gravitational/teleport/lib/auth/storage"
 	"github.com/gravitational/teleport/lib/authz"
@@ -1053,13 +1052,6 @@ func waitAndReload(ctx context.Context, sigC <-chan os.Signal, cfg servicecfg.Co
 func NewTeleport(cfg *servicecfg.Config) (*TeleportProcess, error) {
 	var err error
 
-	// auth and proxy benefit from precomputing keys since they can experience spikes in key
-	// generation due to web session creation and recorded session creation respectively.
-	// for all other agents precomputing keys consumes excess resources.
-	if cfg.Auth.Enabled || cfg.Proxy.Enabled {
-		native.PrecomputeKeys()
-	}
-
 	// Before we do anything reset the SIGINT handler back to the default.
 	system.ResetInterruptSignalHandler()