Skip to content

Commit 032338f

Browse files
ammarioammario
committed
chore: use global cache for RBAC
1 parent 2624ee8 commit 032338f

File tree

8 files changed

+58
-84
lines changed

8 files changed

+58
-84
lines changed

‎coderd.test

85.4 MB
Binary file not shown.

‎coderd/coderdtest/coderdtest.go

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,7 @@ type Options struct {
137137
}
138138

139139
// New constructs a codersdk client connected to an in-memory API instance.
140-
func New(t *testing.T, options *Options) *codersdk.Client {
140+
func New(t testing.TB, options *Options) *codersdk.Client {
141141
client, _ := newWithCloser(t, options)
142142
return client
143143
}
@@ -162,12 +162,12 @@ func NewWithProvisionerCloser(t *testing.T, options *Options) (*codersdk.Client,
162162
// upon thee. Even the io.Closer that is exposed here shouldn't be exposed
163163
// and is a temporary measure while the API to register provisioners is ironed
164164
// out.
165-
func newWithCloser(t *testing.T, options *Options) (*codersdk.Client, io.Closer) {
165+
func newWithCloser(t testing.TB, options *Options) (*codersdk.Client, io.Closer) {
166166
client, closer, _ := NewWithAPI(t, options)
167167
return client, closer
168168
}
169169

170-
func NewOptions(t *testing.T, options *Options) (func(http.Handler), context.CancelFunc, *url.URL, *coderd.Options) {
170+
func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.CancelFunc, *url.URL, *coderd.Options) {
171171
if options == nil {
172172
options = &Options{}
173173
}
@@ -359,7 +359,7 @@ func NewOptions(t *testing.T, options *Options) (func(http.Handler), context.Can
359359
// NewWithAPI constructs an in-memory API instance and returns a client to talk to it.
360360
// Most tests never need a reference to the API, but AuthorizationTest in this module uses it.
361361
// Do not expose the API or wrath shall descend upon thee.
362-
func NewWithAPI(t *testing.T, options *Options) (*codersdk.Client, io.Closer, *coderd.API) {
362+
func NewWithAPI(t testing.TB, options *Options) (*codersdk.Client, io.Closer, *coderd.API) {
363363
if options == nil {
364364
options = &Options{}
365365
}
@@ -384,7 +384,7 @@ func NewWithAPI(t *testing.T, options *Options) (*codersdk.Client, io.Closer, *c
384384
// NewProvisionerDaemon launches a provisionerd instance configured to work
385385
// well with coderd testing. It registers the "echo" provisioner for
386386
// quick testing.
387-
func NewProvisionerDaemon(t *testing.T, coderAPI *coderd.API) io.Closer {
387+
func NewProvisionerDaemon(t testing.TB, coderAPI *coderd.API) io.Closer {
388388
echoClient, echoServer := provisionersdk.MemTransportPipe()
389389
ctx, cancelFunc := context.WithCancel(context.Background())
390390
t.Cleanup(func() {
@@ -465,7 +465,7 @@ var FirstUserParams = codersdk.CreateFirstUserRequest{
465465

466466
// CreateFirstUser creates a user with preset credentials and authenticates
467467
// with the passed in codersdk client.
468-
func CreateFirstUser(t *testing.T, client *codersdk.Client) codersdk.CreateFirstUserResponse {
468+
func CreateFirstUser(t testing.TB, client *codersdk.Client) codersdk.CreateFirstUserResponse {
469469
resp, err := client.CreateFirstUser(context.Background(), FirstUserParams)
470470
require.NoError(t, err)
471471

@@ -1111,7 +1111,7 @@ sz9Di8sGIaUbLZI2rd0CQQCzlVwEtRtoNCyMJTTrkgUuNufLP19RZ5FpyXxBO5/u
11111111
QastnN77KfUwdj3SJt44U/uh1jAIv4oSLBr8HYUkbnI8
11121112
-----END RSA PRIVATE KEY-----`
11131113

1114-
func DeploymentValues(t *testing.T) *codersdk.DeploymentValues {
1114+
func DeploymentValues(t testing.TB) *codersdk.DeploymentValues {
11151115
var cfg codersdk.DeploymentValues
11161116
opts := cfg.Options()
11171117
err := opts.SetDefaults()

‎coderd/database/dbtestutil/db.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ import (
1313
"github.com/coder/coder/coderd/database/postgres"
1414
)
1515

16-
func NewDB(t *testing.T) (database.Store, database.Pubsub) {
16+
func NewDB(t testing.TB) (database.Store, database.Pubsub) {
1717
t.Helper()
1818

1919
db := dbfake.New()

‎coderd/rbac/authz.go

Lines changed: 27 additions & 75 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,12 @@ package rbac
33
import (
44
"context"
55
_ "embed"
6+
"encoding/json"
67
"strings"
78
"sync"
89
"time"
910

11+
"github.com/ammario/tlru"
1012
"github.com/open-policy-agent/opa/ast"
1113
"github.com/open-policy-agent/opa/rego"
1214
"github.com/prometheus/client_golang/prometheus"
@@ -310,6 +312,7 @@ func (a RegoAuthorizer) Authorize(ctx context.Context, subject Subject, action A
310312
defer span.End()
311313

312314
err := a.authorize(ctx, subject, action, object)
315+
313316
span.SetAttributes(attribute.Bool("authorized", err == nil))
314317

315318
dur := time.Since(start)
@@ -605,7 +608,13 @@ func (a *authorizedSQLFilter) SQLString() string {
605608
return a.sqlString
606609
}
607610

608-
type cachedCalls struct {
611+
type authCache struct {
612+
// cache is a cache of JSON marshaled inputs to "authorize" to
613+
// corresponding errors. When the Authorizer is immutable (e.g. in the case
614+
// of Rego)"
615+
// determistic function.
616+
cache *tlru.Cache[string, error]
617+
609618
authz Authorizer
610619
}
611620

@@ -617,94 +626,37 @@ type cachedCalls struct {
617626
//
618627
// Cacher is safe for multiple actors.
619628
func Cacher(authz Authorizer) Authorizer {
620-
return &cachedCalls{authz: authz}
629+
return &authCache{
630+
authz: authz,
631+
cache: tlru.New[string](tlru.ConstantCost[error], 4096),
632+
}
621633
}
622634

623-
func (c *cachedCalls) Authorize(ctx context.Context, subject Subject, action Action, object Object) error {
624-
cache := cacheFromContext(ctx)
635+
func (c *authCache) Authorize(ctx context.Context, subject Subject, action Action, object Object) error {
636+
var authorizeCacheKey strings.Builder
637+
enc := json.NewEncoder(&authorizeCacheKey)
638+
_ = enc.Encode(subject)
639+
_ = enc.Encode(action)
640+
_ = enc.Encode(object)
625641

626-
resp, ok := cache.Load(subject, action, object)
627-
if ok {
628-
return resp
642+
var err error
643+
err, _, ok := c.cache.Get(authorizeCacheKey.String())
644+
if !ok {
645+
err := c.authz.Authorize(ctx, subject, action, object)
646+
// In case there is a caching bug, bound the TTL to 1 minute.
647+
c.cache.Set(authorizeCacheKey.String(), err, time.Minute)
629648
}
630649

631-
err := c.authz.Authorize(ctx, subject, action, object)
632-
cache.Save(subject, action, object, err)
633650
return err
634651
}
635652

636653
// Prepare returns the underlying PreparedAuthorized. The cache does not apply
637654
// to prepared authorizations. These should be using a SQL filter, and
638655
// therefore the cache is not needed.
639-
func (c *cachedCalls) Prepare(ctx context.Context, subject Subject, action Action, objectType string) (PreparedAuthorized, error) {
656+
func (c *authCache) Prepare(ctx context.Context, subject Subject, action Action, objectType string) (PreparedAuthorized, error) {
640657
return c.authz.Prepare(ctx, subject, action, objectType)
641658
}
642659

643-
// authorizeCache enabled caching of Authorizer calls for a given request. This
644-
// prevents the cost of running the same rbac checks multiple times.
645-
// A cache hit must match on all 3 values: subject, action, and object.
646-
type authorizeCache struct {
647-
sync.Mutex
648-
// calls is a list of all calls made to the Authorizer.
649-
// This list is cached per request context. The size of this list is expected
650-
// to be incredibly small. Often 1 or 2 calls.
651-
calls []cachedAuthCall
652-
}
653-
654-
type cachedAuthCall struct {
655-
AuthCall
656-
Err error
657-
}
658-
659-
// cacheContextKey is a context key used to store the cache in the context.
660-
type cacheContextKey struct{}
661-
662-
// cacheFromContext returns the cache from the context.
663-
// If there is no cache, a nil value is returned.
664-
// The nil cache can still be called as a normal cache, but will not cache or
665-
// return any values.
666-
func cacheFromContext(ctx context.Context) *authorizeCache {
667-
cache, _ := ctx.Value(cacheContextKey{}).(*authorizeCache)
668-
return cache
669-
}
670-
671-
func WithCacheCtx(ctx context.Context) context.Context {
672-
return context.WithValue(ctx, cacheContextKey{}, &authorizeCache{})
673-
}
674-
675-
//nolint:revive
676-
func (c *authorizeCache) Load(subject Subject, action Action, object Object) (error, bool) {
677-
if c == nil {
678-
return nil, false
679-
}
680-
c.Lock()
681-
defer c.Unlock()
682-
683-
for _, call := range c.calls {
684-
if call.Action == action && call.Object.Equal(object) && call.Actor.Equal(subject) {
685-
return call.Err, true
686-
}
687-
}
688-
return nil, false
689-
}
690-
691-
func (c *authorizeCache) Save(subject Subject, action Action, object Object, err error) {
692-
if c == nil {
693-
return
694-
}
695-
c.Lock()
696-
defer c.Unlock()
697-
698-
c.calls = append(c.calls, cachedAuthCall{
699-
AuthCall: AuthCall{
700-
Actor: subject,
701-
Action: action,
702-
Object: object,
703-
},
704-
Err: err,
705-
})
706-
}
707-
708660
// rbacTraceAttributes are the attributes that are added to all spans created by
709661
// the rbac package. These attributes should help to debug slow spans.
710662
func rbacTraceAttributes(actor Subject, action Action, objectType string, extra ...attribute.KeyValue) trace.SpanStartOption {

‎coderd/rbac/authz_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -291,7 +291,7 @@ func BenchmarkCacher(b *testing.B) {
291291
sizes := []int{1, 10, 100, 1000}
292292
for _, size := range sizes {
293293
b.Run(fmt.Sprintf("Size%d", size), func(b *testing.B) {
294-
ctx := rbac.WithCacheCtx(context.Background())
294+
ctx := context.Background()
295295
authz := rbac.Cacher(&coderdtest.FakeAuthorizer{AlwaysReturn: nil})
296296
for i := 0; i < size; i++ {
297297
// Preload the cache of a given size

‎coderd/users_test.go

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1671,3 +1671,18 @@ func sortUsers(users []codersdk.User) {
16711671
return users[i].CreatedAt.Before(users[j].CreatedAt)
16721672
})
16731673
}
1674+
1675+
func BenchmarkUsersMe(b *testing.B) {
1676+
client := coderdtest.New(b, nil)
1677+
_ = coderdtest.CreateFirstUser(b, client)
1678+
1679+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
1680+
defer cancel()
1681+
1682+
b.ReportAllocs()
1683+
b.ResetTimer()
1684+
for i := 0; i < b.N; i++ {
1685+
_, err := client.User(ctx, codersdk.Me)
1686+
require.NoError(b, err)
1687+
}
1688+
}

‎go.mod

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -174,6 +174,8 @@ require (
174174
tailscale.com v1.32.2
175175
)
176176

177+
require github.com/armon/go-radix v1.0.0 // indirect
178+
177179
require (
178180
cloud.google.com/go/compute v1.18.0 // indirect
179181
cloud.google.com/go/logging v1.6.1 // indirect
@@ -189,6 +191,7 @@ require (
189191
github.com/akutz/memconn v0.1.0 // indirect
190192
github.com/alecthomas/chroma v0.10.0 // indirect
191193
github.com/alexbrainman/sspi v0.0.0-20210105120005-909beea2cc74 // indirect
194+
github.com/ammario/tlru v0.3.0
192195
github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be // indirect
193196
github.com/apparentlymart/go-textseg/v13 v13.0.0 // indirect
194197
github.com/aymanbagabas/go-osc52 v1.2.1 // indirect

‎go.sum

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,8 @@ github.com/alexbrainman/sspi v0.0.0-20210105120005-909beea2cc74 h1:Kk6a4nehpJ3Uu
170170
github.com/alexbrainman/sspi v0.0.0-20210105120005-909beea2cc74/go.mod h1:cEWa1LVoE5KvSD9ONXsZrj0z6KqySlCCNKHlLzbqAt4=
171171
github.com/alexflint/go-filemutex v0.0.0-20171022225611-72bdc8eae2ae/go.mod h1:CgnQgUtFrFz9mxFNtED3jI5tLDjKlOM+oUF/sTk6ps0=
172172
github.com/alexflint/go-filemutex v1.1.0/go.mod h1:7P4iRhttt/nUvUOrYIhcpMzv2G6CY9UnI16Z+UJqRyk=
173+
github.com/ammario/tlru v0.3.0 h1:yK8ESoFlEyz/BVVL8yZQKAUzJwFJR/j9EfxjnKxtR/Q=
174+
github.com/ammario/tlru v0.3.0/go.mod h1:aYzRFu0XLo4KavE9W8Lx7tzjkX+pAApz+NgcKYIFUBQ=
173175
github.com/andybalholm/brotli v1.0.4 h1:V7DdXeJtZscaqfNuAdSRuRFzuiKlHSC/Zh3zl9qY3JY=
174176
github.com/andybalholm/brotli v1.0.4/go.mod h1:fO7iG3H7G2nSZ7m0zPUDn85XEX2GTukHGRSepvi9Eig=
175177
github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be h1:9AeTilPcZAjCFIImctFaOjnTIavg87rW78vTPkQqLI8=
@@ -191,6 +193,8 @@ github.com/armon/circbuf v0.0.0-20190214190532-5111143e8da2/go.mod h1:3U/XgcO3hC
191193
github.com/armon/consul-api v0.0.0-20180202201655-eb2c6b5be1b6/go.mod h1:grANhF5doyWs3UAsr3K4I6qtAmlQcZDesFNEHPZAzj8=
192194
github.com/armon/go-metrics v0.0.0-20180917152333-f0300d1749da/go.mod h1:Q73ZrmVTwzkszR9V5SSuryQ31EELlFMUz1kKyl939pY=
193195
github.com/armon/go-radix v0.0.0-20180808171621-7fddfc383310/go.mod h1:ufUuZ+zHj4x4TnLV4JWEpy2hxWSpsRywHrMgIH9cCH8=
196+
github.com/armon/go-radix v1.0.0 h1:F4z6KzEeeQIMeLFa97iZU6vupzoecKdU5TX24SNppXI=
197+
github.com/armon/go-radix v1.0.0/go.mod h1:ufUuZ+zHj4x4TnLV4JWEpy2hxWSpsRywHrMgIH9cCH8=
194198
github.com/asaskevich/govalidator v0.0.0-20190424111038-f61b66f89f4a/go.mod h1:lB+ZfQJz7igIIfQNfa7Ml4HSf2uFQQRzpGGRXenZAgY=
195199
github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI=
196200
github.com/awalterschulze/gographviz v2.0.3+incompatible h1:9sVEXJBJLwGX7EQVhLm2elIKCm7P2YHFC8v6096G09E=

0 commit comments

Comments
 (0)