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 <bryan@probo.com>
This commit is contained in:
Bryan Frimin
2026-05-23 11:28:51 -07:00
parent a862faee39
commit 5e9ff656d3
4 changed files with 267 additions and 1 deletions

View File

@@ -0,0 +1,80 @@
// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
//
// 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
}
}

View File

@@ -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
}

View File

@@ -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,

View File

@@ -48,6 +48,7 @@ import (
type (
Resolver struct {
authorize authz.AuthorizeFunc
batchAuthorize authz.BatchAuthorizeFunc
probo *probo.Service
iam *iam.Service
esign *esign.Service