Follow up changes for recent fixes: remove len downcasts and ensure all other downcasts are validated

internal/datasets/subjectset_test.go

@@ -2598,13 +2598,13 @@ func TestIntersectConcreteWithWildcard(t *testing.T) {
 // it counts in binary and "activates" input funcs that match 1s in the binary representation
 // it doesn't check for overflow so don't go crazy
 func allSubsets[T any](objs []T, n int) [][]T {
-	maxInt := uint(math.Exp2(float64(len(objs)))) - 1
+	maxInt := uint64(math.Exp2(float64(len(objs)))) - 1
 	all := make([][]T, 0)
 
-	for i := uint(0); i < maxInt; i++ {
+	for i := uint64(0); i < maxInt; i++ {
 		set := make([]T, 0, n)
-		for digit := uint(0); digit < uint(len(objs)); digit++ {
-			mask := uint(1) << digit
+		for digit := uint64(0); digit < uint64(len(objs)); digit++ {
+			mask := uint64(1) << digit
 			if mask&i != 0 {
 				set = append(set, objs[digit])
 			}

internal/datastore/crdb/pool/balancer.go

@@ -15,6 +15,7 @@ import (
 	"golang.org/x/sync/semaphore"
 
 	log "github.com/authzed/spicedb/internal/logging"
+	"github.com/authzed/spicedb/pkg/genutil"
 )
 
 var (
@@ -108,18 +109,18 @@ func (p *nodeConnectionBalancer[P, C]) Prune(ctx context.Context) {
 		case <-p.ticker.C:
 			if p.sem.TryAcquire(1) {
 				ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
-				p.pruneConnections(ctx)
+				p.mustPruneConnections(ctx)
 				cancel()
 				p.sem.Release(1)
 			}
 		}
 	}
 }
 
-// pruneConnections prunes connections to nodes that have more than MaxConns/(# of nodes)
+// mustPruneConnections prunes connections to nodes that have more than MaxConns/(# of nodes)
 // This causes the pool to reconnect, which over time will lead to a balanced number of connections
 // across each node.
-func (p *nodeConnectionBalancer[P, C]) pruneConnections(ctx context.Context) {
+func (p *nodeConnectionBalancer[P, C]) mustPruneConnections(ctx context.Context) {
 	start := time.Now()
 	defer func() {
 		pruningTimeHistogram.WithLabelValues(p.pool.ID()).Observe(float64(time.Since(start).Milliseconds()))
@@ -224,8 +225,10 @@ func (p *nodeConnectionBalancer[P, C]) pruneConnections(ctx context.Context) {
 		if numToPrune > 1 {
 			numToPrune >>= 1
 		}
-		if uint32(len(healthyConns[node])) < numToPrune {
-			numToPrune = uint32(len(healthyConns[node]))
+
+		healthyNodeCount := genutil.MustEnsureUInt32(len(healthyConns[node]))
+		if healthyNodeCount < numToPrune {
+			numToPrune = healthyNodeCount
 		}
 		if numToPrune == 0 {
 			continue

internal/datastore/crdb/pool/balancer_test.go

@@ -162,7 +162,7 @@ func TestNodeConnectionBalancerPrune(t *testing.T) {
 			ctx, cancel := context.WithCancel(context.Background())
 			defer cancel()
 
-			p.pruneConnections(ctx)
+			p.mustPruneConnections(ctx)
 			require.Equal(t, len(tt.expectedGC), len(pool.gc))
 			gcFromNodes := make([]uint32, 0, len(tt.expectedGC))
 			for _, n := range pool.gc {

internal/services/shared/schema.go

@@ -75,7 +75,7 @@ func ValidateSchemaChanges(ctx context.Context, compiled *compiler.CompiledSchem
 type AppliedSchemaChanges struct {
 	// TotalOperationCount holds the total number of "dispatch" operations performed by the schema
 	// being applied.
-	TotalOperationCount uint32
+	TotalOperationCount int
 
 	// NewObjectDefNames contains the names of the newly added object definitions.
 	NewObjectDefNames []string
@@ -229,7 +229,7 @@ func ApplySchemaChangesOverExisting(
 		Msg("completed schema update")
 
 	return &AppliedSchemaChanges{
-		TotalOperationCount:   uint32(len(validated.compiled.ObjectDefinitions) + len(validated.compiled.CaveatDefinitions) + removedObjectDefNames.Len() + removedCaveatDefNames.Len()),
+		TotalOperationCount:   len(validated.compiled.ObjectDefinitions) + len(validated.compiled.CaveatDefinitions) + removedObjectDefNames.Len() + removedCaveatDefNames.Len(),
 		NewObjectDefNames:     validated.newObjectDefNames.Subtract(existingObjectDefNames).AsSlice(),
 		RemovedObjectDefNames: removedObjectDefNames.AsSlice(),
 		NewCaveatDefNames:     validated.newCaveatDefNames.Subtract(existingCaveatDefNames).AsSlice(),

internal/services/v1/errors.go

@@ -15,16 +15,52 @@ import (
 	"github.com/authzed/spicedb/pkg/tuple"
 )
 
+// ErrExceedsMaximumChecks occurs when too many checks are given to a call.
+type ErrExceedsMaximumChecks struct {
+	error
+	checkCount      uint64
+	maxCountAllowed uint64
+}
+
+// MarshalZerologObject implements zerolog object marshalling.
+func (err ErrExceedsMaximumChecks) MarshalZerologObject(e *zerolog.Event) {
+	e.Err(err.error).Uint64("checkCount", err.checkCount).Uint64("maxCountAllowed", err.maxCountAllowed)
+}
+
+// GRPCStatus implements retrieving the gRPC status for the error.
+func (err ErrExceedsMaximumChecks) GRPCStatus() *status.Status {
+	return spiceerrors.WithCodeAndDetails(
+		err,
+		codes.InvalidArgument,
+		spiceerrors.ForReason(
+			v1.ErrorReason_ERROR_REASON_UNSPECIFIED,
+			map[string]string{
+				"check_count":            strconv.FormatUint(err.checkCount, 10),
+				"maximum_checks_allowed": strconv.FormatUint(err.maxCountAllowed, 10),
+			},
+		),
+	)
+}
+
+// NewExceedsMaximumChecksErr creates a new error representing that too many updates were given to a BulkCheckPermissions call.
+func NewExceedsMaximumChecksErr(checkCount uint64, maxCountAllowed uint64) ErrExceedsMaximumChecks {
+	return ErrExceedsMaximumChecks{
+		error:           fmt.Errorf("check count of %d is greater than maximum allowed of %d", checkCount, maxCountAllowed),
+		checkCount:      checkCount,
+		maxCountAllowed: maxCountAllowed,
+	}
+}
+
 // ErrExceedsMaximumUpdates occurs when too many updates are given to a call.
 type ErrExceedsMaximumUpdates struct {
 	error
-	updateCount     uint16
-	maxCountAllowed uint16
+	updateCount     uint64
+	maxCountAllowed uint64
 }
 
 // MarshalZerologObject implements zerolog object marshalling.
 func (err ErrExceedsMaximumUpdates) MarshalZerologObject(e *zerolog.Event) {
-	e.Err(err.error).Uint16("updateCount", err.updateCount).Uint16("maxCountAllowed", err.maxCountAllowed)
+	e.Err(err.error).Uint64("updateCount", err.updateCount).Uint64("maxCountAllowed", err.maxCountAllowed)
 }
 
 // GRPCStatus implements retrieving the gRPC status for the error.
@@ -35,15 +71,15 @@ func (err ErrExceedsMaximumUpdates) GRPCStatus() *status.Status {
 		spiceerrors.ForReason(
 			v1.ErrorReason_ERROR_REASON_TOO_MANY_UPDATES_IN_REQUEST,
 			map[string]string{
-				"update_count":            strconv.Itoa(int(err.updateCount)),
-				"maximum_updates_allowed": strconv.Itoa(int(err.maxCountAllowed)),
+				"update_count":            strconv.FormatUint(err.updateCount, 10),
+				"maximum_updates_allowed": strconv.FormatUint(err.maxCountAllowed, 10),
 			},
 		),
 	)
 }
 
 // NewExceedsMaximumUpdatesErr creates a new error representing that too many updates were given to a WriteRelationships call.
-func NewExceedsMaximumUpdatesErr(updateCount uint16, maxCountAllowed uint16) ErrExceedsMaximumUpdates {
+func NewExceedsMaximumUpdatesErr(updateCount uint64, maxCountAllowed uint64) ErrExceedsMaximumUpdates {
 	return ErrExceedsMaximumUpdates{
 		error:           fmt.Errorf("update count of %d is greater than maximum allowed of %d", updateCount, maxCountAllowed),
 		updateCount:     updateCount,
@@ -54,13 +90,13 @@ func NewExceedsMaximumUpdatesErr(updateCount uint16, maxCountAllowed uint16) Err
 // ErrExceedsMaximumPreconditions occurs when too many preconditions are given to a call.
 type ErrExceedsMaximumPreconditions struct {
 	error
-	preconditionCount uint16
-	maxCountAllowed   uint16
+	preconditionCount uint64
+	maxCountAllowed   uint64
 }
 
 // MarshalZerologObject implements zerolog object marshalling.
 func (err ErrExceedsMaximumPreconditions) MarshalZerologObject(e *zerolog.Event) {
-	e.Err(err.error).Uint16("preconditionCount", err.preconditionCount).Uint16("maxCountAllowed", err.maxCountAllowed)
+	e.Err(err.error).Uint64("preconditionCount", err.preconditionCount).Uint64("maxCountAllowed", err.maxCountAllowed)
 }
 
 // GRPCStatus implements retrieving the gRPC status for the error.
@@ -71,15 +107,15 @@ func (err ErrExceedsMaximumPreconditions) GRPCStatus() *status.Status {
 		spiceerrors.ForReason(
 			v1.ErrorReason_ERROR_REASON_TOO_MANY_PRECONDITIONS_IN_REQUEST,
 			map[string]string{
-				"precondition_count":      strconv.Itoa(int(err.preconditionCount)),
-				"maximum_updates_allowed": strconv.Itoa(int(err.maxCountAllowed)),
+				"precondition_count":      strconv.FormatUint(err.preconditionCount, 10),
+				"maximum_updates_allowed": strconv.FormatUint(err.maxCountAllowed, 10),
 			},
 		),
 	)
 }
 
 // NewExceedsMaximumPreconditionsErr creates a new error representing that too many preconditions were given to a call.
-func NewExceedsMaximumPreconditionsErr(preconditionCount uint16, maxCountAllowed uint16) ErrExceedsMaximumPreconditions {
+func NewExceedsMaximumPreconditionsErr(preconditionCount uint64, maxCountAllowed uint64) ErrExceedsMaximumPreconditions {
 	return ErrExceedsMaximumPreconditions{
 		error: fmt.Errorf(
 			"precondition count of %d is greater than maximum allowed of %d",

internal/services/v1/experimental.go

@@ -34,6 +34,7 @@ import (
 	"github.com/authzed/spicedb/pkg/cursor"
 	"github.com/authzed/spicedb/pkg/datastore"
 	dsoptions "github.com/authzed/spicedb/pkg/datastore/options"
+	"github.com/authzed/spicedb/pkg/genutil"
 	"github.com/authzed/spicedb/pkg/genutil/mapz"
 	"github.com/authzed/spicedb/pkg/genutil/slicez"
 	"github.com/authzed/spicedb/pkg/middleware/consistency"
@@ -409,14 +410,25 @@ func (es *experimentalServer) BulkExportRelationships(
 	return nil
 }
 
+const maxBulkCheckCount = 10000
+
 func (es *experimentalServer) BulkCheckPermission(ctx context.Context, req *v1.BulkCheckPermissionRequest) (*v1.BulkCheckPermissionResponse, error) {
 	atRevision, checkedAt, err := consistency.RevisionFromContext(ctx)
 	if err != nil {
 		return nil, es.rewriteError(ctx, err)
 	}
 
+	if len(req.Items) > maxBulkCheckCount {
+		return nil, es.rewriteError(ctx, NewExceedsMaximumChecksErr(uint64(len(req.Items)), maxBulkCheckCount))
+	}
+
 	// Compute a hash for each requested item and record its index(es) for the items, to be used for sorting of results.
-	itemIndexByHash := mapz.NewMultiMapWithCap[string, int](uint32(len(req.Items)))
+	itemCount, err := genutil.EnsureUInt32(len(req.Items))
+	if err != nil {
+		return nil, es.rewriteError(ctx, err)
+	}
+
+	itemIndexByHash := mapz.NewMultiMapWithCap[string, int](itemCount)
 	for index, item := range req.Items {
 		itemHash, err := computeBulkCheckPermissionItemHash(item)
 		if err != nil {

internal/services/v1/experimental_test.go

@@ -204,7 +204,7 @@ func TestBulkExportRelationships(t *testing.T) {
 					}
 
 					require.NoError(err)
-					require.LessOrEqual(uint32(len(batch.Relationships)), tc.batchSize)
+					require.LessOrEqual(uint64(len(batch.Relationships)), uint64(tc.batchSize))
 					require.NotNil(batch.AfterResultCursor)
 					require.NotEmpty(batch.AfterResultCursor.Token)
 

internal/services/v1/relationships.go

@@ -26,6 +26,7 @@ import (
 	"github.com/authzed/spicedb/pkg/datastore"
 	"github.com/authzed/spicedb/pkg/datastore/options"
 	"github.com/authzed/spicedb/pkg/datastore/pagination"
+	"github.com/authzed/spicedb/pkg/genutil"
 	"github.com/authzed/spicedb/pkg/genutil/mapz"
 	"github.com/authzed/spicedb/pkg/middleware/consistency"
 	dispatchv1 "github.com/authzed/spicedb/pkg/proto/dispatch/v1"
@@ -254,14 +255,14 @@ func (ps *permissionServer) WriteRelationships(ctx context.Context, req *v1.Writ
 	if len(req.Updates) > int(ps.config.MaxUpdatesPerWrite) {
 		return nil, ps.rewriteError(
 			ctx,
-			NewExceedsMaximumUpdatesErr(uint16(len(req.Updates)), ps.config.MaxUpdatesPerWrite),
+			NewExceedsMaximumUpdatesErr(uint64(len(req.Updates)), uint64(ps.config.MaxUpdatesPerWrite)),
 		)
 	}
 
 	if len(req.OptionalPreconditions) > int(ps.config.MaxPreconditionsCount) {
 		return nil, ps.rewriteError(
 			ctx,
-			NewExceedsMaximumPreconditionsErr(uint16(len(req.OptionalPreconditions)), ps.config.MaxPreconditionsCount),
+			NewExceedsMaximumPreconditionsErr(uint64(len(req.OptionalPreconditions)), uint64(ps.config.MaxPreconditionsCount)),
 		)
 	}
 
@@ -302,9 +303,14 @@ func (ps *permissionServer) WriteRelationships(ctx context.Context, req *v1.Writ
 			return ps.rewriteError(ctx, err)
 		}
 
+		dispatchCount, err := genutil.EnsureUInt32(len(req.OptionalPreconditions) + 1)
+		if err != nil {
+			return ps.rewriteError(ctx, err)
+		}
+
 		usagemetrics.SetInContext(ctx, &dispatchv1.ResponseMeta{
 			// One request per precondition and one request for the actual writes.
-			DispatchCount: uint32(len(req.OptionalPreconditions)) + 1,
+			DispatchCount: dispatchCount,
 		})
 
 		span.AddEvent("preconditions")
@@ -338,7 +344,7 @@ func (ps *permissionServer) DeleteRelationships(ctx context.Context, req *v1.Del
 	if len(req.OptionalPreconditions) > int(ps.config.MaxPreconditionsCount) {
 		return nil, ps.rewriteError(
 			ctx,
-			NewExceedsMaximumPreconditionsErr(uint16(len(req.OptionalPreconditions)), ps.config.MaxPreconditionsCount),
+			NewExceedsMaximumPreconditionsErr(uint64(len(req.OptionalPreconditions)), uint64(ps.config.MaxPreconditionsCount)),
 		)
 	}
 
@@ -350,9 +356,14 @@ func (ps *permissionServer) DeleteRelationships(ctx context.Context, req *v1.Del
 			return err
 		}
 
+		dispatchCount, err := genutil.EnsureUInt32(len(req.OptionalPreconditions) + 1)
+		if err != nil {
+			return ps.rewriteError(ctx, err)
+		}
+
 		usagemetrics.SetInContext(ctx, &dispatchv1.ResponseMeta{
 			// One request per precondition and one request for the actual delete.
-			DispatchCount: uint32(len(req.OptionalPreconditions)) + 1,
+			DispatchCount: dispatchCount,
 		})
 
 		if err := checkPreconditions(ctx, rwt, req.OptionalPreconditions); err != nil {
@@ -403,7 +414,7 @@ func (ps *permissionServer) DeleteRelationships(ctx context.Context, req *v1.Del
 		}
 
 		// Otherwise, kick off an unlimited deletion.
-		_, err := rwt.DeleteRelationships(ctx, req.RelationshipFilter)
+		_, err = rwt.DeleteRelationships(ctx, req.RelationshipFilter)
 		return err
 	})
 	if err != nil {

internal/services/v1/schema.go

@@ -14,6 +14,7 @@ import (
 	"github.com/authzed/spicedb/internal/middleware/usagemetrics"
 	"github.com/authzed/spicedb/internal/services/shared"
 	"github.com/authzed/spicedb/pkg/datastore"
+	"github.com/authzed/spicedb/pkg/genutil"
 	dispatchv1 "github.com/authzed/spicedb/pkg/proto/dispatch/v1"
 	"github.com/authzed/spicedb/pkg/schemadsl/compiler"
 	"github.com/authzed/spicedb/pkg/schemadsl/generator"
@@ -87,8 +88,13 @@ func (ss *schemaServer) ReadSchema(ctx context.Context, _ *v1.ReadSchemaRequest)
 		return nil, ss.rewriteError(ctx, err)
 	}
 
+	dispatchCount, err := genutil.EnsureUInt32(len(nsDefs) + len(caveatDefs))
+	if err != nil {
+		return nil, ss.rewriteError(ctx, err)
+	}
+
 	usagemetrics.SetInContext(ctx, &dispatchv1.ResponseMeta{
-		DispatchCount: uint32(len(nsDefs) + len(caveatDefs)),
+		DispatchCount: dispatchCount,
 	})
 
 	return &v1.ReadSchemaResponse{
@@ -124,8 +130,14 @@ func (ss *schemaServer) WriteSchema(ctx context.Context, in *v1.WriteSchemaReque
 		if err != nil {
 			return err
 		}
+
+		dispatchCount, err := genutil.EnsureUInt32(applied.TotalOperationCount)
+		if err != nil {
+			return err
+		}
+
 		usagemetrics.SetInContext(ctx, &dispatchv1.ResponseMeta{
-			DispatchCount: applied.TotalOperationCount,
+			DispatchCount: dispatchCount,
 		})
 		return nil
 	})