From 5e9ff656d3a4b65285df2a0a7076f2aae5b5250d Mon Sep 17 00:00:00 2001 From: Bryan Frimin Date: Sat, 23 May 2026 11:28:51 -0700 Subject: [PATCH] Add per-request Authorize dataloader to console v1 Resolving a typical Console GraphQL query triggers many parallel authorize calls (one per resource per field resolver). This commit collapses them via a dataloader: parallel calls within the same request are gathered into a single iam.Authorizer.AuthorizeMulti pass, and only fall back to per-item Authorize when AuthorizeMulti rejects the whole batch (e.g. mixed organizations). The loader key encodes resource id, action, options, and a canonical JSON-encoded attribute map so logically identical calls share a key while differing ones do not. The loader is created without caching so repeated calls within a request still produce one audit log entry per call. dataloader.NewAuthorizeFunc preserves the existing authz.AuthorizeFunc signature and error mapping. Signed-off-by: Bryan Frimin --- .../api/console/v1/dataloader/authorize.go | 80 ++++++++ .../api/console/v1/dataloader/dataloader.go | 183 ++++++++++++++++++ pkg/server/api/console/v1/graphql_handler.go | 4 +- pkg/server/api/console/v1/resolver.go | 1 + 4 files changed, 267 insertions(+), 1 deletion(-) create mode 100644 pkg/server/api/console/v1/dataloader/authorize.go diff --git a/pkg/server/api/console/v1/dataloader/authorize.go b/pkg/server/api/console/v1/dataloader/authorize.go new file mode 100644 index 000000000..d506ba725 --- /dev/null +++ b/pkg/server/api/console/v1/dataloader/authorize.go @@ -0,0 +1,80 @@ +// Copyright (c) 2026 Probo Inc . +// +// Permission to use, copy, modify, and/or distribute this software for any +// purpose with or without fee is hereby granted, provided that the above +// copyright notice and this permission notice appear in all copies. +// +// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +// PERFORMANCE OF THIS SOFTWARE. + +package dataloader + +import ( + "context" + "errors" + + "go.gearno.de/kit/log" + "go.probo.inc/probo/pkg/coredata" + "go.probo.inc/probo/pkg/gid" + "go.probo.inc/probo/pkg/iam" + "go.probo.inc/probo/pkg/server/api/authz" + "go.probo.inc/probo/pkg/server/gqlutils" +) + +// NewAuthorizeFunc returns an authz.AuthorizeFunc that batches authorize +// calls through the per-request dataloader. Parallel field resolvers within +// the same request collapse into a single iam.Authorizer.AuthorizeMulti +// call. Requires the dataloader middleware to have populated Loaders in +// context. +func NewAuthorizeFunc(logger *log.Logger) authz.AuthorizeFunc { + return func( + ctx context.Context, + objectID gid.GID, + action string, + options ...authz.AuthorizeFuncOption, + ) (*coredata.Scope, error) { + loaders := FromContext(ctx) + + applied := iam.AuthorizeParams{ + ResourceAttributes: make(map[string]string), + } + for _, opt := range options { + opt(&applied) + } + + result, err := loaders.Authorize.Load( + ctx, + AuthorizeKey{ + ResourceID: objectID, + Action: action, + ResourceAttributes: EncodeAuthorizeKeyAttributes(applied.ResourceAttributes), + DryRun: applied.DryRun, + SkipAssumptionCheck: applied.SkipAssumptionCheck, + }, + ) + if err != nil { + if _, ok := errors.AsType[*iam.ErrAssumptionRequired](err); ok { + return nil, gqlutils.AssumptionRequired(ctx, err) + } + + if _, ok := errors.AsType[*iam.ErrInsufficientPermissions](err); ok { + return nil, gqlutils.Forbidden(ctx, err) + } + + if errors.Is(err, coredata.ErrResourceNotFound) { + return nil, gqlutils.NotFoundf(ctx, "resource not found") + } + + logger.ErrorCtx(ctx, "cannot authorize", log.Error(err)) + + return nil, gqlutils.Internal(ctx) + } + + return result.Scope, nil + } +} diff --git a/pkg/server/api/console/v1/dataloader/dataloader.go b/pkg/server/api/console/v1/dataloader/dataloader.go index fa257d28d..fac69abdc 100644 --- a/pkg/server/api/console/v1/dataloader/dataloader.go +++ b/pkg/server/api/console/v1/dataloader/dataloader.go @@ -16,20 +16,42 @@ package dataloader import ( "context" + "encoding/json" "fmt" + "maps" "net/http" + "slices" + "strings" "github.com/vikstrous/dataloadgen" "go.probo.inc/probo/pkg/cookiebanner" "go.probo.inc/probo/pkg/coredata" "go.probo.inc/probo/pkg/gid" "go.probo.inc/probo/pkg/iam" + "go.probo.inc/probo/pkg/iam/policy" "go.probo.inc/probo/pkg/probo" + "go.probo.inc/probo/pkg/server/api/authn" ) type ( ctxKey struct{ name string } + // AuthorizeKey identifies an authorize call. ResourceAttributes is the + // canonical (sorted-key) JSON encoding of the attributes passed via + // authz.WithAttr, so calls with the same logical inputs share a key and + // batch together. + AuthorizeKey struct { + ResourceID gid.GID + Action string + ResourceAttributes string + DryRun bool + SkipAssumptionCheck bool + } + + AuthorizeResult struct { + Scope *coredata.Scope + } + Loaders struct { Organization *dataloadgen.Loader[gid.GID, *coredata.Organization] Framework *dataloadgen.Loader[gid.GID, *coredata.Framework] @@ -44,6 +66,7 @@ type ( Report *dataloadgen.Loader[gid.GID, *coredata.Report] CookieBanner *dataloadgen.Loader[gid.GID, *coredata.CookieBanner] CookieCategory *dataloadgen.Loader[gid.GID, *coredata.CookieCategory] + Authorize *dataloadgen.Loader[AuthorizeKey, AuthorizeResult] } batchFetcher struct { @@ -87,6 +110,10 @@ func (f *batchFetcher) newLoaders() *Loaders { Report: dataloadgen.NewMappedLoader(f.fetchReports), CookieBanner: dataloadgen.NewMappedLoader(f.fetchCookieBanners), CookieCategory: dataloadgen.NewMappedLoader(f.fetchCookieCategories), + Authorize: dataloadgen.NewMappedLoader( + f.fetchAuthorizes, + dataloadgen.WithoutCache(), + ), } } @@ -297,3 +324,159 @@ func (f *batchFetcher) fetchCookieCategories(ctx context.Context, keys []gid.GID return result, nil } + +// fetchAuthorizes evaluates the batch with a single AuthorizeMulti call and +// surfaces per-key denials via dataloadgen.MappedFetchError. When +// AuthorizeMulti cannot evaluate the batch as a whole (e.g. mixed +// organizations), we fall back to per-item Authorize so every key still +// gets a scope or iam error. +// +// The Authorize loader is created with WithoutCache() so repeated calls +// with the same (resource, action) within a single request still produce +// one audit log entry per call. +func (f *batchFetcher) fetchAuthorizes( + ctx context.Context, + keys []AuthorizeKey, +) (map[AuthorizeKey]AuthorizeResult, error) { + identity := authn.IdentityFromContext(ctx) + if identity == nil { + return nil, fmt.Errorf("cannot authorize without an identity in context") + } + + session := authn.SessionFromContext(ctx) + + items := make([]iam.MultiAuthorizeItem, 0, len(keys)) + for _, key := range keys { + attrs, err := decodeAuthorizeKeyAttributes(key.ResourceAttributes) + if err != nil { + return nil, fmt.Errorf("cannot decode authorize key attributes: %w", err) + } + + items = append(items, iam.MultiAuthorizeItem{ + Resource: key.ResourceID, + Action: key.Action, + ResourceAttributes: attrs, + DryRun: key.DryRun, + SkipAssumptionCheck: key.SkipAssumptionCheck, + }) + } + + multiParams := iam.AuthorizeMultiParams{ + Principal: identity.ID, + Items: items, + } + if session != nil { + multiParams.Session = &session.ID + } + + scope, decisions, err := f.iam.Authorizer.AuthorizeMulti(ctx, multiParams) + if err != nil { + return f.fetchAuthorizesIndividually(ctx, keys, identity.ID, session) + } + + result := make(map[AuthorizeKey]AuthorizeResult, len(keys)) + perKeyErrs := make(dataloadgen.MappedFetchError[AuthorizeKey]) + + for i, key := range keys { + if decisions[i] != nil { + perKeyErrs[key] = decisions[i] + continue + } + + result[key] = AuthorizeResult{Scope: scope} + } + + if len(perKeyErrs) > 0 { + return result, perKeyErrs + } + + return result, nil +} + +// fetchAuthorizesIndividually is the per-item fallback used when +// AuthorizeMulti cannot evaluate the batch as a whole. +func (f *batchFetcher) fetchAuthorizesIndividually( + ctx context.Context, + keys []AuthorizeKey, + principalID gid.GID, + session *coredata.Session, +) (map[AuthorizeKey]AuthorizeResult, error) { + result := make(map[AuthorizeKey]AuthorizeResult, len(keys)) + perKeyErrs := make(dataloadgen.MappedFetchError[AuthorizeKey]) + + for _, key := range keys { + attrs, err := decodeAuthorizeKeyAttributes(key.ResourceAttributes) + if err != nil { + return nil, fmt.Errorf("cannot decode authorize key attributes: %w", err) + } + + params := iam.AuthorizeParams{ + Principal: principalID, + Resource: key.ResourceID, + Action: key.Action, + ResourceAttributes: make(map[string]string, len(attrs)), + DryRun: key.DryRun, + SkipAssumptionCheck: key.SkipAssumptionCheck, + } + maps.Copy(params.ResourceAttributes, attrs) + + if session != nil { + params.Session = &session.ID + } + + scope, err := f.iam.Authorizer.Authorize(ctx, params) + if err != nil { + perKeyErrs[key] = err + continue + } + + result[key] = AuthorizeResult{Scope: scope} + } + + if len(perKeyErrs) > 0 { + return result, perKeyErrs + } + + return result, nil +} + +// EncodeAuthorizeKeyAttributes returns a canonical (sorted-key) JSON +// encoding of attrs for use as AuthorizeKey.ResourceAttributes. Maps that +// compare equal produce identical strings; nil or empty attrs encode to "". +func EncodeAuthorizeKeyAttributes(attrs policy.Attributes) string { + if len(attrs) == 0 { + return "" + } + + keys := slices.Sorted(maps.Keys(attrs)) + + var sb strings.Builder + sb.WriteByte('{') + for i, k := range keys { + if i > 0 { + sb.WriteByte(',') + } + + kb, _ := json.Marshal(k) + sb.Write(kb) + sb.WriteByte(':') + vb, _ := json.Marshal(attrs[k]) + sb.Write(vb) + } + sb.WriteByte('}') + + return sb.String() +} + +func decodeAuthorizeKeyAttributes(s string) (policy.Attributes, error) { + if s == "" { + return nil, nil + } + + attrs := policy.Attributes{} + if err := json.Unmarshal([]byte(s), &attrs); err != nil { + return nil, err + } + + return attrs, nil +} diff --git a/pkg/server/api/console/v1/graphql_handler.go b/pkg/server/api/console/v1/graphql_handler.go index a291031f7..535d86686 100644 --- a/pkg/server/api/console/v1/graphql_handler.go +++ b/pkg/server/api/console/v1/graphql_handler.go @@ -27,6 +27,7 @@ import ( "go.probo.inc/probo/pkg/probo" "go.probo.inc/probo/pkg/riskmanagement" "go.probo.inc/probo/pkg/server/api/authz" + "go.probo.inc/probo/pkg/server/api/console/v1/dataloader" "go.probo.inc/probo/pkg/server/api/console/v1/schema" "go.probo.inc/probo/pkg/server/gqlutils" "go.probo.inc/probo/pkg/thirdparty" @@ -47,7 +48,8 @@ func NewGraphQLHandler( ) http.Handler { config := schema.Config{ Resolvers: &Resolver{ - authorize: authz.NewAuthorizeFunc(iamSvc, logger), + authorize: dataloader.NewAuthorizeFunc(logger), + batchAuthorize: authz.NewBatchAuthorizeFunc(iamSvc, logger), probo: proboSvc, iam: iamSvc, esign: esignSvc, diff --git a/pkg/server/api/console/v1/resolver.go b/pkg/server/api/console/v1/resolver.go index dbe81e560..92c6e9337 100644 --- a/pkg/server/api/console/v1/resolver.go +++ b/pkg/server/api/console/v1/resolver.go @@ -48,6 +48,7 @@ import ( type ( Resolver struct { authorize authz.AuthorizeFunc + batchAuthorize authz.BatchAuthorizeFunc probo *probo.Service iam *iam.Service esign *esign.Service