Don't fetch the role list twice for Kubernetes connection locks (#68959)

* Add MaxKubernetesConnections to the AccessChecker interface

* Don't fetch the role list again in the Kube forwarder

* Move tracing span to acquireConnectionLock

* Add TODO for scoped identities

lib/kube/proxy/forwarder.go

@@ -659,23 +659,19 @@ func (f *Forwarder) withAuthStd(handler handlerWithAuthFuncStd) http.HandlerFunc
 
 // acquireConnectionLockWithIdentity acquires a connection lock under a given identity.
 func (f *Forwarder) acquireConnectionLockWithIdentity(ctx context.Context, identity *authContext) error {
-	ctx, span := f.cfg.tracer.Start(
-		ctx,
-		"kube.Forwarder/acquireConnectionLockWithIdentity",
-		oteltrace.WithSpanKind(oteltrace.SpanKindServer),
-		oteltrace.WithAttributes(
-			semconv.RPCServiceKey.String(f.cfg.KubeServiceType),
-			semconv.RPCSystemKey.String("kube"),
-		),
-	)
-	defer span.End()
-	user := identity.Identity.GetIdentity().Username
-	roles, err := getRolesByName(ctx, f, identity.Identity.GetIdentity().Groups)
-	if err != nil {
-		return trace.Wrap(err)
+	unscopedContext, isUnscoped := identity.UnscopedContext()
+	if !isUnscoped {
+		// TODO(espadolini) TODO(eriktate): scoped identities don't currently
+		// support max_kubernetes_connections, this should be updated when they
+		// do
+		return nil
 	}
-
-	if err := f.acquireConnectionLock(ctx, user, roles); err != nil {
+	maxConnections := unscopedContext.Checker.MaxKubernetesConnections()
+	if maxConnections == 0 {
+		return nil
+	}
+	user := unscopedContext.Identity.GetIdentity().Username
+	if err := f.acquireConnectionLock(ctx, user, maxConnections); err != nil {
 		return trace.Wrap(err)
 	}
 
@@ -1681,8 +1677,18 @@ func wsProxy(ctx context.Context, log *slog.Logger, wsSource *gwebsocket.Conn, w
 // acquireConnectionLock acquires a semaphore used to limit connections to the Kubernetes agent.
 // The semaphore is releasted when the request is returned/connection is closed.
 // Returns an error if a semaphore could not be acquired.
-func (f *Forwarder) acquireConnectionLock(ctx context.Context, user string, roles services.RoleSet) error {
-	maxConnections := roles.MaxKubernetesConnections()
+func (f *Forwarder) acquireConnectionLock(ctx context.Context, user string, maxConnections int64) error {
+	ctx, span := f.cfg.tracer.Start(
+		ctx,
+		"kube.Forwarder/acquireConnectionLock",
+		oteltrace.WithSpanKind(oteltrace.SpanKindServer),
+		oteltrace.WithAttributes(
+			semconv.RPCServiceKey.String(f.cfg.KubeServiceType),
+			semconv.RPCSystemKey.String("kube"),
+		),
+	)
+	defer span.End()
+
 	if maxConnections == 0 {
 		return nil
 	}

lib/kube/proxy/forwarder_test.go

@@ -2004,6 +2004,13 @@ func TestKubernetesConnectionLimit(t *testing.T) {
 						Username: user.GetName(),
 						Groups:   []string{testCase.role.GetName()},
 					}),
+					Checker: services.NewAccessCheckerWithRoleSet(
+						&services.AccessInfo{
+							Roles: []string{testCase.role.GetName()},
+						},
+						"clustername",
+						services.NewRoleSet(testCase.role),
+					),
 				}),
 			}
 

lib/services/access_checker.go

@@ -309,6 +309,12 @@ type AccessChecker interface {
 
 	// DelegationSessionID returns the ID of the current Delegation Session.
 	DelegationSessionID() string
+
+	// MaxKubernetesConnections returns the maximum number of concurrent
+	// Kubernetes connections allowed. If MaxKubernetesConnections is zero then
+	// no maximum was defined and the number of concurrent connections is
+	// unconstrained.
+	MaxKubernetesConnections() int64
 }
 
 // AccessInfo hold information about an identity necessary to check whether that

lib/services/role.go

@@ -1599,9 +1599,7 @@ func (set RoleSet) MaxSessions() int64 {
 	return ms
 }
 
-// MaxConnections returns the maximum number of concurrent Kubernetes connections
-// allowed.  If MaxConnections is zero then no maximum was defined
-// and the number of concurrent connections is unconstrained.
+// MaxKubernetesConnections implements [AccessChecker].
 func (set RoleSet) MaxKubernetesConnections() int64 {
 	var mcs int64
 	for _, role := range set {