From 321446308d1fc296c221c87187da9da3c6436ef2 Mon Sep 17 00:00:00 2001
From: Edoardo Spadolini <edoardo.spadolini@goteleport.com>
Date: Tue, 28 May 2024 18:16:11 +0200
Subject: [PATCH] Improve role fetching and unmarshaling (#42007) (#42062)

* Reduce default ListRoles page size

* Avoid stdjson in UnmarshalRoleV6
---
 lib/services/local/access.go |  5 +--
 lib/services/role.go         | 61 ++++++++++++++----------------------
 2 files changed, 27 insertions(+), 39 deletions(-)

diff --git a/lib/services/local/access.go b/lib/services/local/access.go
index b7b3ccd66843d..bba8edf576d7d 100644
--- a/lib/services/local/access.go
+++ b/lib/services/local/access.go
@@ -29,7 +29,6 @@ import (
 
 	"github.com/gravitational/teleport"
 	"github.com/gravitational/teleport/api/client/proto"
-	apidefaults "github.com/gravitational/teleport/api/defaults"
 	"github.com/gravitational/teleport/api/types"
 	"github.com/gravitational/teleport/lib/backend"
 	"github.com/gravitational/teleport/lib/services"
@@ -91,7 +90,9 @@ func (s *AccessService) ListRoles(ctx context.Context, req *proto.ListRolesReque
 	limit := int(req.Limit)
 
 	if limit == 0 {
-		limit = apidefaults.DefaultChunkSize
+		// it can take a lot of effort to parse roles and until a page is done
+		// parsing, it will be held in memory - so keep this reasonably small
+		limit = 100
 	}
 
 	if limit > maxPageSize {
diff --git a/lib/services/role.go b/lib/services/role.go
index 7e2c681226bba..20abf5e11be0a 100644
--- a/lib/services/role.go
+++ b/lib/services/role.go
@@ -20,7 +20,6 @@ package services
 
 import (
 	"context"
-	"encoding/json"
 	"errors"
 	"fmt"
 	"net"
@@ -34,6 +33,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws/arn"
 	"github.com/google/uuid"
 	"github.com/gravitational/trace"
+	jsoniter "github.com/json-iterator/go"
 	log "github.com/sirupsen/logrus"
 	"github.com/vulcand/predicate"
 	"golang.org/x/crypto/ssh"
@@ -3286,50 +3286,37 @@ func UnmarshalRole(bytes []byte, opts ...MarshalOption) (types.Role, error) {
 
 // UnmarshalRoleV6 unmarshals the RoleV6 resource from JSON.
 func UnmarshalRoleV6(bytes []byte, opts ...MarshalOption) (*types.RoleV6, error) {
-	var h types.ResourceHeader
-	err := json.Unmarshal(bytes, &h)
-	if err != nil {
-		return nil, trace.Wrap(err)
-	}
-
 	cfg, err := CollectOptions(opts)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
 
-	switch h.Version {
-	case types.V7:
-		fallthrough
-	case types.V6:
-		fallthrough
-	case types.V5:
-		fallthrough
-	case types.V4:
-		// V4 roles are identical to V3 except for their defaults
-		fallthrough
-	case types.V3:
-		var role types.RoleV6
-		if err := utils.FastUnmarshal(bytes, &role); err != nil {
-			return nil, trace.BadParameter(err.Error())
-		}
-
-		if err := ValidateRole(&role); err != nil {
-			return nil, trace.Wrap(err)
-		}
+	version := jsoniter.Get(bytes, "version").ToString()
+	switch version {
+	// these are all backed by the same shape of data, they just have different semantics and defaults
+	case types.V3, types.V4, types.V5, types.V6, types.V7:
+	default:
+		return nil, trace.BadParameter("role version %q is not supported", version)
+	}
+	var role types.RoleV6
+	if err := utils.FastUnmarshal(bytes, &role); err != nil {
+		return nil, trace.BadParameter(err.Error())
+	}
 
-		if cfg.ID != 0 {
-			role.SetResourceID(cfg.ID)
-		}
-		if cfg.Revision != "" {
-			role.SetRevision(cfg.Revision)
-		}
-		if !cfg.Expires.IsZero() {
-			role.SetExpiry(cfg.Expires)
-		}
-		return &role, nil
+	if err := ValidateRole(&role); err != nil {
+		return nil, trace.Wrap(err)
 	}
 
-	return nil, trace.BadParameter("role version %q is not supported", h.Version)
+	if cfg.ID != 0 {
+		role.SetResourceID(cfg.ID)
+	}
+	if cfg.Revision != "" {
+		role.SetRevision(cfg.Revision)
+	}
+	if !cfg.Expires.IsZero() {
+		role.SetExpiry(cfg.Expires)
+	}
+	return &role, nil
 }
 
 // MarshalRole marshals the Role resource to JSON.