Skip to content

Commit 04e9ffb

Browse files
smallinskysmallinsky
authored
Improve userstate.Generate performance (#58387)
1 parent 7a477f1 commit 04e9ffb

File tree

3 files changed

+191
-103
lines changed

3 files changed

+191
-103
lines changed

‎lib/auth/userloginstate/generator.go‎

Lines changed: 62 additions & 93 deletions
Original file line numberDiff line numberDiff line change
@@ -27,14 +27,14 @@ import (
2727
"github.com/jonboulle/clockwork"
2828

2929
"github.com/gravitational/teleport/api/client/proto"
30-
accesslistv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/accesslist/v1"
3130
usageeventsv1 "github.com/gravitational/teleport/api/gen/proto/go/usageevents/v1"
3231
"github.com/gravitational/teleport/api/types"
3332
"github.com/gravitational/teleport/api/types/accesslist"
3433
apievents "github.com/gravitational/teleport/api/types/events"
3534
"github.com/gravitational/teleport/api/types/header"
3635
"github.com/gravitational/teleport/api/types/userloginstate"
3736
"github.com/gravitational/teleport/api/utils"
37+
"github.com/gravitational/teleport/api/utils/clientutils"
3838
"github.com/gravitational/teleport/lib/accesslists"
3939
"github.com/gravitational/teleport/lib/events"
4040
"github.com/gravitational/teleport/lib/modules"
@@ -212,127 +212,96 @@ func (g *Generator) generate(ctx context.Context, user types.User, ulsService se
212212

213213
// addAccessListsToState will add the user's applicable access lists to the user login state after validating them, returning any inherited roles and traits.
214214
func (g *Generator) addAccessListsToState(ctx context.Context, user types.User, state *userloginstate.UserLoginState) (inheritedRoles []string, inheritedTraits map[string][]string, err error) {
215-
accessLists, err := g.accessLists.GetAccessLists(ctx)
215+
locks, err := g.accessLists.GetLocks(ctx, true, types.LockTarget{
216+
User: user.GetName(),
217+
})
216218
if err != nil {
217219
return nil, nil, trace.Wrap(err)
218220
}
221+
if len(locks) > 0 {
222+
return inheritedRoles, inheritedTraits, nil
223+
}
219224

220225
var allInheritedRoles []string
221226
allInheritedTraits := make(map[string][]string)
222-
223-
for _, accessList := range accessLists {
224-
// Grants are inherited if the user is a member of the access list, explicitly or via inheritance.
225-
inheritedRoles, inheritedTraits, err := g.handleAccessListMembership(ctx, user, accessList, state)
227+
applyHierarchy := func(hierarchy []*accesslist.AccessList, accessListName string, ownerHierarchy bool) error {
228+
if len(hierarchy) == 0 {
229+
return nil
230+
}
231+
inheritedRoles, inheritedTraits, err := g.applyGrantsAcrossACLs(ctx, hierarchy, accessListName, user, state, ownerHierarchy)
226232
if err != nil {
227-
return nil, nil, trace.Wrap(err)
233+
return trace.Wrap(err)
228234
}
229235
allInheritedRoles = append(allInheritedRoles, inheritedRoles...)
230236
for k, values := range inheritedTraits {
231237
allInheritedTraits[k] = append(allInheritedTraits[k], values...)
232238
}
239+
return nil
240+
}
233241

234-
// OwnerGrants are inherited if the user is an owner of the access list, explicitly or via inheritance.
235-
inheritedRoles, inheritedTraits, err = g.handleAccessListOwnership(ctx, user, accessList, state)
242+
h, err := accesslists.NewHierarchy(accesslists.HierarchyConfig{
243+
AccessListsService: g.accessLists,
244+
Clock: g.clock,
245+
})
246+
if err != nil {
247+
return nil, nil, trace.Wrap(err)
248+
}
249+
250+
for acl, err := range clientutils.Resources(ctx, g.accessLists.ListAccessLists) {
236251
if err != nil {
237252
return nil, nil, trace.Wrap(err)
238253
}
239-
allInheritedRoles = append(allInheritedRoles, inheritedRoles...)
240-
for k, values := range inheritedTraits {
241-
allInheritedTraits[k] = append(allInheritedTraits[k], values...)
254+
memberOf, ownerOf, err := h.GetHierarchyForUser(ctx, acl, user)
255+
if err != nil {
256+
return nil, nil, trace.Wrap(err)
257+
}
258+
// Grants are inherited if the user is a member of the access list, explicitly or via inheritance.
259+
if err := applyHierarchy(memberOf, acl.GetName(), false); err != nil {
260+
return nil, nil, trace.Wrap(err)
261+
}
262+
// OwnerGrants are inherited if the user is an owner of the access list, explicitly or via inheritance.
263+
if err := applyHierarchy(ownerOf, acl.GetName(), true); err != nil {
264+
return nil, nil, trace.Wrap(err)
242265
}
243266
}
244-
245267
return allInheritedRoles, allInheritedTraits, nil
246268
}
247269

248-
// handleAccessListMembership validates the access list and applies the grants and traits from the access list to the user if they are a member of the access list.
249-
// If the access list is invalid (because it references a non-existent role, for example,
250-
// then it will not be applied.
251-
func (g *Generator) handleAccessListMembership(ctx context.Context, user types.User, accessList *accesslist.AccessList, state *userloginstate.UserLoginState) ([]string, map[string][]string, error) {
252-
var inheritedRoles []string
253-
inheritedTraits := make(map[string][]string)
254-
255-
membershipKind, err := accesslists.IsAccessListMember(ctx, user, accessList, g.accessLists, g.accessLists, g.clock)
256-
// Return early if there was an error or the user isn't a member of the access list.
257-
if err != nil || membershipKind == accesslistv1.AccessListUserAssignmentType_ACCESS_LIST_USER_ASSIGNMENT_TYPE_UNSPECIFIED {
258-
// Log any error besides user being locked.
259-
if err != nil && !accesslists.IsUserLocked(err) {
260-
g.log.WarnContext(ctx, "checking access list membership", "error", err)
270+
func (g *Generator) applyGrantsAcrossACLs(ctx context.Context, acls []*accesslist.AccessList, baseName string, user types.User, state *userloginstate.UserLoginState, owner bool) (inheritedRoles []string, inheritedTraits map[string][]string, err error) {
271+
inheritedTraits = make(map[string][]string)
272+
for _, acl := range acls {
273+
grants := selectGrants(acl, owner)
274+
missing, err := g.identifyMissingRoles(ctx, grants.Roles)
275+
if err != nil {
276+
return nil, nil, trace.Wrap(err)
261277
}
262-
return inheritedRoles, inheritedTraits, nil
263-
}
264-
265-
// Validate that all the roles in the access list exist.
266-
missingRoles, err := g.identifyMissingRoles(ctx, accessList.Spec.Grants.Roles)
267-
if err != nil {
268-
return nil, nil, trace.Wrap(err)
269-
}
270-
271-
// If there are any missing roles, then we cannot apply the access list.
272-
// Emit an audit event and return early.
273-
// This flow is designed to skip the entire access list rather than processing individual roles within it.
274-
// This approach ensures that access lists are treated as cohesive units of access control. Partial
275-
// application of an access list could result in unintended permission configurations, potentially leading
276-
// to security vulnerabilities or unpredictable behavior.
277-
if missingRoles != nil {
278-
g.emitSkippedAccessListEvent(ctx, accessList.Spec.Title, missingRoles, user.GetName())
279-
return nil, nil, nil
280-
}
281-
282-
g.grantRolesAndTraits(accessList.Spec.Grants, state)
283-
if membershipKind == accesslistv1.AccessListUserAssignmentType_ACCESS_LIST_USER_ASSIGNMENT_TYPE_INHERITED {
284-
inheritedRoles = append(inheritedRoles, accessList.Spec.Grants.Roles...)
285-
for k, values := range accessList.Spec.Grants.Traits {
278+
if len(missing) > 0 {
279+
// If there are any missing roles, then we cannot apply the access list.
280+
// Emit an audit event and skip the access list.
281+
// This flow is designed to skip the entire access list rather than processing individual roles within it.
282+
// This approach ensures that access lists are treated as cohesive units of access control. Partial
283+
// application of an access list could result in unintended permission configurations, potentially leading
284+
// to security vulnerabilities or unpredictable behavior.
285+
g.emitSkippedAccessListEvent(ctx, acl.Spec.Title, missing, user.GetName())
286+
continue
287+
}
288+
g.grantRolesAndTraits(grants, state)
289+
if acl.GetName() == baseName {
290+
continue
291+
}
292+
inheritedRoles = append(inheritedRoles, grants.Roles...)
293+
for k, values := range grants.Traits {
286294
inheritedTraits[k] = append(inheritedTraits[k], values...)
287295
}
288296
}
289-
290297
return inheritedRoles, inheritedTraits, nil
291298
}
292299

293-
// handleAccessListOwnership validates the access list and applies the grants and traits from the access list to the user if they are an owner of the access list.
294-
// If the access list is invalid (because it references a non-existent role, for example,
295-
// then it will not be applied.
296-
func (g *Generator) handleAccessListOwnership(ctx context.Context, user types.User, accessList *accesslist.AccessList, state *userloginstate.UserLoginState) ([]string, map[string][]string, error) {
297-
var inheritedRoles []string
298-
inheritedTraits := make(map[string][]string)
299-
300-
ownershipType, err := accesslists.IsAccessListOwner(ctx, user, accessList, g.accessLists, g.accessLists, g.clock)
301-
// Return early if there was an error or the user isn't an owner of the access list.
302-
if err != nil || ownershipType == accesslistv1.AccessListUserAssignmentType_ACCESS_LIST_USER_ASSIGNMENT_TYPE_UNSPECIFIED {
303-
// Log any error besides user being locked.
304-
if err != nil && !accesslists.IsUserLocked(err) {
305-
g.log.WarnContext(ctx, "checking access list ownership", "error", err)
306-
}
307-
return inheritedRoles, inheritedTraits, nil
308-
}
309-
310-
// Validate that all the roles in the access list exist.
311-
missingRoles, err := g.identifyMissingRoles(ctx, accessList.Spec.OwnerGrants.Roles)
312-
if err != nil {
313-
return nil, nil, trace.Wrap(err)
314-
}
315-
316-
// If there are any missing roles, then we cannot apply the access list.
317-
// Emit an audit event and return early.
318-
// This flow is designed to skip the entire access list rather than processing individual roles within it.
319-
// This approach ensures that access lists are treated as cohesive units of access control. Partial
320-
// application of an access list could result in unintended permission configurations, potentially leading
321-
// to security vulnerabilities or unpredictable behavior.
322-
if missingRoles != nil {
323-
g.emitSkippedAccessListEvent(ctx, accessList.Spec.Title, missingRoles, user.GetName())
324-
return nil, nil, nil
300+
func selectGrants(acl *accesslist.AccessList, owner bool) accesslist.Grants {
301+
if owner {
302+
return acl.Spec.OwnerGrants
325303
}
326-
327-
g.grantRolesAndTraits(accessList.Spec.OwnerGrants, state)
328-
if ownershipType == accesslistv1.AccessListUserAssignmentType_ACCESS_LIST_USER_ASSIGNMENT_TYPE_INHERITED {
329-
inheritedRoles = append(inheritedRoles, accessList.Spec.OwnerGrants.Roles...)
330-
for k, values := range accessList.Spec.OwnerGrants.Traits {
331-
inheritedTraits[k] = append(inheritedTraits[k], values...)
332-
}
333-
}
334-
335-
return inheritedRoles, inheritedTraits, nil
304+
return acl.Spec.Grants
336305
}
337306

338307
// grantRolesAndTraits will append the roles and traits from the provided Grants to the UserLoginState,

‎lib/auth/userloginstate/generator_bench_test.go‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
/*
2+
* Teleport
3+
* Copyright (C) 2025 Gravitational, Inc.
4+
*
5+
* This program is free software: you can redistribute it and/or modify
6+
* it under the terms of the GNU Affero General Public License as published by
7+
* the Free Software Foundation, either version 3 of the License, or
8+
* (at your option) any later version.
9+
*
10+
* This program is distributed in the hope that it will be useful,
11+
* but WITHOUT ANY WARRANTY; without even the implied warranty of
12+
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
13+
* GNU Affero General Public License for more details.
14+
*
15+
* You should have received a copy of the GNU Affero General Public License
16+
* along with this program. If not, see <http://www.gnu.org/licenses/>.
17+
*/
18+
19+
package userloginstate
20+
21+
import (
22+
"fmt"
23+
"testing"
24+
"time"
25+
26+
"github.com/gravitational/trace"
27+
"github.com/jonboulle/clockwork"
28+
29+
"github.com/gravitational/teleport/api/types"
30+
"github.com/gravitational/teleport/api/types/accesslist"
31+
"github.com/gravitational/teleport/api/types/header"
32+
)
33+
34+
func BenchmarkGenerate(b *testing.B) {
35+
user, err := types.NewUser("alice")
36+
if err != nil {
37+
b.Fatalf("NewUser: %v", err)
38+
}
39+
40+
accessList, err := makeAccessList("acl")
41+
if err != nil {
42+
b.Fatalf("makeAccessList: %v", err)
43+
}
44+
45+
sizes := []int{100, 1_000, 10_000}
46+
for _, n := range sizes {
47+
b.Run(fmt.Sprintf("members=%d", n), func(b *testing.B) {
48+
svc, backendSvc, err := initGeneratorSvc()
49+
if err != nil {
50+
b.Fatalf("initGeneratorSvc: %v", err)
51+
}
52+
_, err = backendSvc.UpsertAccessList(b.Context(), accessList)
53+
if err != nil {
54+
b.Fatalf("UpsertAccessList: %v", err)
55+
}
56+
57+
for i := 0; i < n; i++ {
58+
member := makeUserMember(accessList.GetName(), fmt.Sprintf("u%d", i))
59+
_, err = backendSvc.UpsertAccessListMember(b.Context(), member)
60+
if err != nil {
61+
b.Fatalf("UpsertAccessListMember: %v", err)
62+
}
63+
}
64+
65+
b.ReportAllocs()
66+
b.ResetTimer()
67+
68+
for b.Loop() {
69+
_, err := svc.Generate(b.Context(), user, backendSvc)
70+
if err != nil {
71+
b.Fatalf("Generate: %v", err)
72+
}
73+
}
74+
})
75+
}
76+
}
77+
78+
func makeAccessList(name string) (*accesslist.AccessList, error) {
79+
accessList, err := accesslist.NewAccessList(header.Metadata{
80+
Name: name,
81+
}, accesslist.Spec{
82+
Title: "title",
83+
Audit: accesslist.Audit{
84+
NextAuditDate: clockwork.NewRealClock().Now().Add(time.Hour * 48),
85+
},
86+
Owners: []accesslist.Owner{{Name: ownerUser}},
87+
})
88+
if err != nil {
89+
return nil, trace.Wrap(err)
90+
}
91+
return accessList, nil
92+
}
93+
94+
func makeUserMember(acl, username string) *accesslist.AccessListMember {
95+
m := &accesslist.AccessListMember{
96+
ResourceHeader: header.ResourceHeader{
97+
Metadata: header.Metadata{
98+
Name: username,
99+
},
100+
},
101+
Spec: accesslist.AccessListMemberSpec{
102+
Name: username,
103+
AccessList: acl,
104+
MembershipKind: accesslist.MembershipKindUser,
105+
Joined: clockwork.NewRealClock().Now(),
106+
AddedBy: "system",
107+
},
108+
}
109+
return m
110+
}

‎lib/auth/userloginstate/generator_test.go‎

Lines changed: 19 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import (
2525

2626
"github.com/google/go-cmp/cmp"
2727
"github.com/google/go-cmp/cmp/cmpopts"
28+
"github.com/gravitational/trace"
2829
"github.com/jonboulle/clockwork"
2930
"github.com/stretchr/testify/require"
3031

@@ -579,7 +580,8 @@ func TestAccessLists(t *testing.T) {
579580
})
580581

581582
ctx := context.Background()
582-
svc, backendSvc := initGeneratorSvc(t)
583+
svc, backendSvc, err := initGeneratorSvc()
584+
require.NoError(t, err)
583585

584586
for _, accessList := range test.accessLists {
585587
_, err = backendSvc.UpsertAccessList(ctx, accessList)
@@ -632,7 +634,8 @@ func TestAccessLists(t *testing.T) {
632634

633635
func TestGitHubIdentity(t *testing.T) {
634636
ctx := context.Background()
635-
svc, backendSvc := initGeneratorSvc(t)
637+
svc, backendSvc, err := initGeneratorSvc()
638+
require.NoError(t, err)
636639

637640
noGitHubIdentity, err := types.NewUser("alice")
638641
require.NoError(t, err)
@@ -715,20 +718,24 @@ func (s *svc) SubmitUsageEvent(ctx context.Context, req *proto.SubmitUsageEventR
715718
return nil
716719
}
717720

718-
func initGeneratorSvc(t *testing.T) (*Generator, *svc) {
719-
t.Helper()
720-
721+
func initGeneratorSvc() (*Generator, *svc, error) {
721722
clock := clockwork.NewFakeClock()
722723
mem, err := memory.New(memory.Config{
723724
Clock: clock,
724725
})
725-
require.NoError(t, err)
726+
if err != nil {
727+
return nil, nil, trace.Wrap(err)
728+
}
726729

727730
accessListsSvc, err := local.NewAccessListService(mem, clock)
728-
require.NoError(t, err)
731+
if err != nil {
732+
return nil, nil, trace.Wrap(err)
733+
}
729734
accessSvc := local.NewAccessService(mem)
730735
ulsService, err := local.NewUserLoginStateService(mem)
731-
require.NoError(t, err)
736+
if err != nil {
737+
return nil, nil, trace.Wrap(err)
738+
}
732739

733740
svc := &svc{
734741
AccessLists: accessListsSvc,
@@ -746,8 +753,10 @@ func initGeneratorSvc(t *testing.T) (*Generator, *svc) {
746753
Clock: clock,
747754
Emitter: emitter,
748755
})
749-
require.NoError(t, err)
750-
return generator, svc
756+
if err != nil {
757+
return nil, nil, trace.Wrap(err)
758+
}
759+
return generator, svc, nil
751760
}
752761

753762
func grants(roles []string, traits trait.Traits) accesslist.Grants {

0 commit comments

Comments
 (0)