VAULT-6614 Enable role based quotas for lease-count quotas (OSS)

builtin/plugin/backend.go

@@ -7,6 +7,8 @@ import (
 	"reflect"
 	"sync"
 
+	log "github.com/hashicorp/go-hclog"
+
 	uuid "github.com/hashicorp/go-uuid"
 	"github.com/hashicorp/vault/sdk/framework"
 	"github.com/hashicorp/vault/sdk/helper/consts"
@@ -38,7 +40,7 @@ func Factory(ctx context.Context, conf *logical.BackendConfig) (logical.Backend,
 
 // Backend returns an instance of the backend, either as a plugin if external
 // or as a concrete implementation if builtin, casted as logical.Backend.
-func Backend(ctx context.Context, conf *logical.BackendConfig) (logical.Backend, error) {
+func Backend(ctx context.Context, conf *logical.BackendConfig) (*PluginBackend, error) {
 	var b PluginBackend
 
 	name := conf.Config["plugin_name"]
@@ -80,7 +82,7 @@ func Backend(ctx context.Context, conf *logical.BackendConfig) (logical.Backend,
 
 // PluginBackend is a thin wrapper around plugin.BackendPluginClient
 type PluginBackend struct {
-	logical.Backend
+	Backend logical.Backend
 	sync.RWMutex
 
 	config *logical.BackendConfig
@@ -118,12 +120,12 @@ func (b *PluginBackend) startBackend(ctx context.Context, storage logical.Storag
 	if !b.loaded {
 		if b.Backend.Type() != nb.Type() {
 			nb.Cleanup(ctx)
-			b.Logger().Warn("failed to start plugin process", "plugin", b.config.Config["plugin_name"], "error", ErrMismatchType)
+			b.Backend.Logger().Warn("failed to start plugin process", "plugin", b.config.Config["plugin_name"], "error", ErrMismatchType)
 			return ErrMismatchType
 		}
 		if !reflect.DeepEqual(b.Backend.SpecialPaths(), nb.SpecialPaths()) {
 			nb.Cleanup(ctx)
-			b.Logger().Warn("failed to start plugin process", "plugin", b.config.Config["plugin_name"], "error", ErrMismatchPaths)
+			b.Backend.Logger().Warn("failed to start plugin process", "plugin", b.config.Config["plugin_name"], "error", ErrMismatchPaths)
 			return ErrMismatchPaths
 		}
 	}
@@ -169,7 +171,7 @@ func (b *PluginBackend) lazyLoadBackend(ctx context.Context, storage logical.Sto
 		// Reload plugin if it's an rpc.ErrShutdown
 		b.Lock()
 		if b.canary == canary {
-			b.Logger().Debug("reloading plugin backend", "plugin", b.config.Config["plugin_name"])
+			b.Backend.Logger().Debug("reloading plugin backend", "plugin", b.config.Config["plugin_name"])
 			err := b.startBackend(ctx, storage)
 			if err != nil {
 				b.Unlock()
@@ -220,3 +222,52 @@ func (b *PluginBackend) HandleExistenceCheck(ctx context.Context, req *logical.R
 func (b *PluginBackend) Initialize(ctx context.Context, req *logical.InitializationRequest) error {
 	return nil
 }
+
+// SpecialPaths is a thin wrapper used to ensure we grab the lock for race purposes
+func (b *PluginBackend) SpecialPaths() *logical.Paths {
+	b.RLock()
+	defer b.RUnlock()
+	return b.Backend.SpecialPaths()
+}
+
+// System is a thin wrapper used to ensure we grab the lock for race purposes
+func (b *PluginBackend) System() logical.SystemView {
+	b.RLock()
+	defer b.RUnlock()
+	return b.Backend.System()
+}
+
+// Logger is a thin wrapper used to ensure we grab the lock for race purposes
+func (b *PluginBackend) Logger() log.Logger {
+	b.RLock()
+	defer b.RUnlock()
+	return b.Backend.Logger()
+}
+
+// Cleanup is a thin wrapper used to ensure we grab the lock for race purposes
+func (b *PluginBackend) Cleanup(ctx context.Context) {
+	b.RLock()
+	defer b.RUnlock()
+	b.Backend.Cleanup(ctx)
+}
+
+// InvalidateKey is a thin wrapper used to ensure we grab the lock for race purposes
+func (b *PluginBackend) InvalidateKey(ctx context.Context, key string) {
+	b.RLock()
+	defer b.RUnlock()
+	b.Backend.InvalidateKey(ctx, key)
+}
+
+// Setup is a thin wrapper used to ensure we grab the lock for race purposes
+func (b *PluginBackend) Setup(ctx context.Context, config *logical.BackendConfig) error {
+	b.RLock()
+	defer b.RUnlock()
+	return b.Backend.Setup(ctx, config)
+}
+
+// Type is a thin wrapper used to ensure we grab the lock for race purposes
+func (b *PluginBackend) Type() logical.BackendType {
+	b.RLock()
+	defer b.RUnlock()
+	return b.Backend.Type()
+}

http/util.go

@@ -64,7 +64,7 @@ func rateLimitQuotaWrapping(handler http.Handler, core *vault.Core) http.Handler
 			Type:          quotas.TypeRateLimit,
 			Path:          path,
 			MountPath:     mountPath,
-			Role:          core.DetermineRoleFromLoginRequest(mountPath, bodyBytes, r.Context()),
+			Role:          core.DetermineRoleFromLoginRequestFromBytes(mountPath, bodyBytes, r.Context()),
 			NamespacePath: ns.Path,
 			ClientAddress: parseRemoteIPAddress(r),
 		})

vault/core.go

@@ -175,7 +175,7 @@ func (e *ErrInvalidKey) Error() string {
 	return fmt.Sprintf("invalid key: %v", e.Reason)
 }
 
-type RegisterAuthFunc func(context.Context, time.Duration, string, *logical.Auth) error
+type RegisterAuthFunc func(context.Context, time.Duration, string, *logical.Auth, string) error
 
 type activeAdvertisement struct {
 	RedirectAddr     string                     `json:"redirect_addr"`
@@ -3301,22 +3301,30 @@ func (c *Core) CheckPluginPerms(pluginName string) (err error) {
 	return err
 }
 
+// DetermineRoleFromLoginRequestFromBytes will determine the role that should be applied to a quota for a given
+// login request, accepting a byte payload
+func (c *Core) DetermineRoleFromLoginRequestFromBytes(mountPoint string, payload []byte, ctx context.Context) string {
+	data := make(map[string]interface{})
+	err := jsonutil.DecodeJSON(payload, &data)
+	if err != nil {
+		// Cannot discern a role from a request we cannot parse
+		return ""
+	}
+
+	return c.DetermineRoleFromLoginRequest(mountPoint, data, ctx)
+}
+
 // DetermineRoleFromLoginRequest will determine the role that should be applied to a quota for a given
 // login request
-func (c *Core) DetermineRoleFromLoginRequest(mountPoint string, payload []byte, ctx context.Context) string {
+func (c *Core) DetermineRoleFromLoginRequest(mountPoint string, data map[string]interface{}, ctx context.Context) string {
+	c.authLock.RLock()
+	defer c.authLock.RUnlock()
 	matchingBackend := c.router.MatchingBackend(ctx, mountPoint)
 	if matchingBackend == nil || matchingBackend.Type() != logical.TypeCredential {
 		// Role based quotas do not apply to this request
 		return ""
 	}
 
-	data := make(map[string]interface{})
-	err := jsonutil.DecodeJSON(payload, &data)
-	if err != nil {
-		// Cannot discern a role from a request we cannot parse
-		return ""
-	}
-
 	resp, err := matchingBackend.HandleRequest(ctx, &logical.Request{
 		MountPoint: mountPoint,
 		Path:       "login",

vault/core_util.go

@@ -177,7 +177,7 @@ func (c *Core) quotaLeaseWalker(ctx context.Context, callback func(request *quot
 	return nil
 }
 
-func (c *Core) quotasHandleLeases(ctx context.Context, action quotas.LeaseAction, leaseIDs []string) error {
+func (c *Core) quotasHandleLeases(ctx context.Context, action quotas.LeaseAction, leases []*quotas.QuotaLeaseInformation) error {
 	return nil
 }