feat: add tracing to request validation

Author: miparnisari

CHANGELOG.md

@@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 ## [Unreleased]
 ### Changed
 - Updated CI so that Postgres tests run against v18 which is GA and not against v13 which is EOL (https://github.com/authzed/spicedb/pull/2926)
+- Added tracing to request validation (https://github.com/authzed/spicedb/pull/2950)
 
 ### Fixed
 - Regression introduced in 1.49.2: missing spans in ReadSchema calls (https://github.com/authzed/spicedb/pull/2947)

internal/middleware/handwrittenvalidation/doc.go

@@ -1,2 +1,2 @@
-// Package handwrittenvalidation defines middleware that runs custom-made validations on incoming requests.
+// Package handwrittenvalidation defines middleware that runs validations on incoming requests.
 package handwrittenvalidation

internal/middleware/handwrittenvalidation/handwrittenvalidation.go

@@ -2,51 +2,99 @@ package handwrittenvalidation
 
 import (
 	"context"
+	"fmt"
 
+	"buf.build/go/protovalidate"
+	"go.opentelemetry.io/otel"
+	otelcodes "go.opentelemetry.io/otel/codes"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/proto"
 )
 
+var tracer = otel.Tracer("spicedb/internal/middleware")
+
+// mustNewProtoValidator wraps protovalidate.New() to panic
+// if the validator can't be constructed.
+func mustNewProtoValidator(opts ...protovalidate.ValidatorOption) protovalidate.Validator {
+	validator, err := protovalidate.New(opts...)
+	if err != nil {
+		wrappedErr := fmt.Errorf("could not construct validator: %w", err)
+		panic(wrappedErr)
+	}
+	return validator
+}
+
 type handwrittenValidator interface {
 	HandwrittenValidate() error
 }
 
-// UnaryServerInterceptor returns a new unary server interceptor that runs the handwritten validation
-// on the incoming request, if any.
-func UnaryServerInterceptor(ctx context.Context, req any, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
-	validator, ok := req.(handwrittenValidator)
-	if ok {
-		err := validator.HandwrittenValidate()
+// UnaryServerInterceptor returns a function that performs proto and handwritten validation (if any) on the incoming request.
+func UnaryServerInterceptor() grpc.UnaryServerInterceptor {
+	protovalidator := mustNewProtoValidator()
+	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
+		_, span := tracer.Start(ctx, "protovalidate")
+		err := protovalidator.Validate(req.(proto.Message))
 		if err != nil {
+			span.SetStatus(otelcodes.Error, err.Error())
+			span.RecordError(err)
+			span.End()
 			return nil, status.Errorf(codes.InvalidArgument, "%s", err)
 		}
-	}
 
-	return handler(ctx, req)
+		validator, ok := req.(handwrittenValidator)
+		if ok {
+			_, span := tracer.Start(ctx, "handwritten.validation")
+			err := validator.HandwrittenValidate()
+			if err != nil {
+				span.SetStatus(otelcodes.Error, err.Error())
+				span.RecordError(err)
+				span.End()
+				return nil, status.Errorf(codes.InvalidArgument, "%s", err)
+			}
+		}
+
+		span.End()
+
+		return handler(ctx, req)
+	}
 }
 
-// StreamServerInterceptor returns a new stream server interceptor that runs the handwritten validation
-// on the incoming request messages, if any.
+// StreamServerInterceptor returns a function that performs proto and handwritten validation (if any) on the incoming request.
 func StreamServerInterceptor(srv any, stream grpc.ServerStream, _ *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
-	wrapper := &recvWrapper{stream}
+	wrapper := &recvWrapper{ServerStream: stream, protovalidator: mustNewProtoValidator()}
 	return handler(srv, wrapper)
 }
 
 type recvWrapper struct {
 	grpc.ServerStream
+	protovalidator protovalidate.Validator
 }
 
 func (s *recvWrapper) RecvMsg(m any) error {
 	if err := s.ServerStream.RecvMsg(m); err != nil {
 		return err
 	}
 
+	_, span := tracer.Start(s.Context(), "protovalidate")
+
+	err := s.protovalidator.Validate(m.(proto.Message))
+	if err != nil {
+		span.SetStatus(otelcodes.Error, err.Error())
+		span.RecordError(err)
+		span.End()
+		return status.Errorf(codes.InvalidArgument, "%s", err)
+	}
+
 	validator, ok := m.(handwrittenValidator)
 	if ok {
-		err := validator.HandwrittenValidate()
+		err = validator.HandwrittenValidate()
 		if err != nil {
-			return err
+			span.SetStatus(otelcodes.Error, err.Error())
+			span.RecordError(err)
+			span.End()
+			return status.Errorf(codes.InvalidArgument, "%s", err)
 		}
 	}
 

internal/services/v1/experimental.go

@@ -13,7 +13,6 @@ import (
 	"time"
 
 	"github.com/ccoveille/go-safecast/v2"
-	grpcvalidate "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/protovalidate"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/protobuf/types/known/timestamppb"
@@ -36,7 +35,6 @@ import (
 	"github.com/authzed/spicedb/pkg/datastore"
 	dsoptions "github.com/authzed/spicedb/pkg/datastore/options"
 	"github.com/authzed/spicedb/pkg/datastore/queryshape"
-	"github.com/authzed/spicedb/pkg/genutil"
 	"github.com/authzed/spicedb/pkg/middleware/consistency"
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
 	dispatchv1 "github.com/authzed/spicedb/pkg/proto/dispatch/v1"
@@ -94,18 +92,14 @@ func NewExperimentalServer(dispatch dispatch.Dispatcher, permServerConfig Permis
 		chunkSize = 100
 	}
 
-	validator := genutil.MustNewProtoValidator()
-
 	return &experimentalServer{
 		WithServiceSpecificInterceptors: shared.WithServiceSpecificInterceptors{
 			Unary: middleware.ChainUnaryServer(
-				grpcvalidate.UnaryServerInterceptor(validator),
-				handwrittenvalidation.UnaryServerInterceptor,
+				handwrittenvalidation.UnaryServerInterceptor(),
 				usagemetrics.UnaryServerInterceptor(),
 				perfinsights.UnaryServerInterceptor(permServerConfig.PerformanceInsightMetricsEnabled),
 			),
 			Stream: middleware.ChainStreamServer(
-				grpcvalidate.StreamServerInterceptor(validator),
 				handwrittenvalidation.StreamServerInterceptor,
 				usagemetrics.StreamServerInterceptor(),
 				streamtimeout.MustStreamServerInterceptor(config.StreamReadTimeout),

internal/services/v1/relationships.go

@@ -7,7 +7,6 @@ import (
 	"fmt"
 	"time"
 
-	grpcvalidate "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/protovalidate"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
 	"go.opentelemetry.io/otel/trace"
@@ -146,20 +145,16 @@ func NewPermissionsServer(
 		ExperimentalQueryPlan:              config.ExperimentalQueryPlan,
 	}
 
-	validator := genutil.MustNewProtoValidator()
-
 	return &permissionServer{
 		dispatch: dispatch,
 		config:   configWithDefaults,
 		WithServiceSpecificInterceptors: shared.WithServiceSpecificInterceptors{
 			Unary: middleware.ChainUnaryServer(
-				grpcvalidate.UnaryServerInterceptor(validator),
-				handwrittenvalidation.UnaryServerInterceptor,
+				handwrittenvalidation.UnaryServerInterceptor(),
 				usagemetrics.UnaryServerInterceptor(),
 				perfinsights.UnaryServerInterceptor(configWithDefaults.PerformanceInsightMetricsEnabled),
 			),
 			Stream: middleware.ChainStreamServer(
-				grpcvalidate.StreamServerInterceptor(validator),
 				handwrittenvalidation.StreamServerInterceptor,
 				usagemetrics.StreamServerInterceptor(),
 				streamtimeout.MustStreamServerInterceptor(configWithDefaults.StreamingAPITimeout),

internal/services/v1/schema.go

@@ -5,12 +5,11 @@ import (
 	"sort"
 	"strings"
 
-	grpcvalidate "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/protovalidate"
-
 	v1 "github.com/authzed/authzed-go/proto/authzed/api/v1"
 
 	log "github.com/authzed/spicedb/internal/logging"
 	"github.com/authzed/spicedb/internal/middleware"
+	"github.com/authzed/spicedb/internal/middleware/handwrittenvalidation"
 	"github.com/authzed/spicedb/internal/middleware/perfinsights"
 	"github.com/authzed/spicedb/internal/middleware/usagemetrics"
 	"github.com/authzed/spicedb/internal/services/shared"
@@ -45,17 +44,15 @@ type SchemaServerConfig struct {
 func NewSchemaServer(config SchemaServerConfig) v1.SchemaServiceServer {
 	cts := caveattypes.TypeSetOrDefault(config.CaveatTypeSet)
 
-	validator := genutil.MustNewProtoValidator()
-
 	return &schemaServer{
 		WithServiceSpecificInterceptors: shared.WithServiceSpecificInterceptors{
 			Unary: middleware.ChainUnaryServer(
-				grpcvalidate.UnaryServerInterceptor(validator),
+				handwrittenvalidation.UnaryServerInterceptor(),
 				usagemetrics.UnaryServerInterceptor(),
 				perfinsights.UnaryServerInterceptor(config.PerformanceInsightMetricsEnabled),
 			),
 			Stream: middleware.ChainStreamServer(
-				grpcvalidate.StreamServerInterceptor(validator),
+				handwrittenvalidation.StreamServerInterceptor,
 				usagemetrics.StreamServerInterceptor(),
 				perfinsights.StreamServerInterceptor(config.PerformanceInsightMetricsEnabled),
 			),

internal/services/v1/watch.go

@@ -5,18 +5,17 @@ import (
 	"slices"
 	"time"
 
-	grpcvalidate "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/protovalidate"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/types/known/structpb"
 
 	v1 "github.com/authzed/authzed-go/proto/authzed/api/v1"
 
+	"github.com/authzed/spicedb/internal/middleware/handwrittenvalidation"
 	"github.com/authzed/spicedb/internal/middleware/usagemetrics"
 	"github.com/authzed/spicedb/internal/services/shared"
 	"github.com/authzed/spicedb/pkg/datalayer"
 	"github.com/authzed/spicedb/pkg/datastore"
-	"github.com/authzed/spicedb/pkg/genutil"
 	"github.com/authzed/spicedb/pkg/genutil/mapz"
 	dispatchv1 "github.com/authzed/spicedb/pkg/proto/dispatch/v1"
 	"github.com/authzed/spicedb/pkg/tuple"
@@ -32,11 +31,9 @@ type watchServer struct {
 
 // NewWatchServer creates an instance of the watch server.
 func NewWatchServer(heartbeatDuration time.Duration) v1.WatchServiceServer {
-	validator := genutil.MustNewProtoValidator()
-
 	s := &watchServer{
 		WithStreamServiceSpecificInterceptor: shared.WithStreamServiceSpecificInterceptor{
-			Stream: grpcvalidate.StreamServerInterceptor(validator),
+			Stream: handwrittenvalidation.StreamServerInterceptor,
 		},
 		heartbeatDuration: heartbeatDuration,
 	}