protovalidate: avoid pointer comparisons

When deciding whether the user has bypassed validation for a message, we
should compare the message's fully-qualified name rather than doing a
fragile pointer comparison.

Along the way, this commit removes some puzzling and unused code to
plumb a context through the wrappedServerStream.

Signed-off-by: Akshay Shah <[email protected]>

Author: akshayjshah

interceptors/protovalidate/options.go

@@ -12,7 +12,7 @@ import (
 )
 
 type options struct {
-	ignoreMessages []protoreflect.MessageType
+	ignoreMessages []protoreflect.FullName
 }
 
 // An Option lets you add options to protovalidate interceptors using With* funcs.
@@ -29,13 +29,17 @@ func evaluateOpts(opts []Option) *options {
 // WithIgnoreMessages sets the messages that should be ignored by the validator. Use with
 // caution and ensure validation is performed elsewhere.
 func WithIgnoreMessages(msgs ...protoreflect.MessageType) Option {
+	names := make([]protoreflect.FullName, 0, len(msgs))
+	for _, msg := range msgs {
+		names = append(names, msg.Descriptor().FullName())
+	}
+	slices.Sort(names)
 	return func(o *options) {
-		o.ignoreMessages = msgs
+		o.ignoreMessages = names
 	}
 }
 
-func (o *options) shouldIgnoreMessage(m protoreflect.MessageType) bool {
-	return slices.ContainsFunc(o.ignoreMessages, func(t protoreflect.MessageType) bool {
-		return m == t
-	})
+func (o *options) shouldIgnoreMessage(fqn protoreflect.FullName) bool {
+	_, found := slices.BinarySearch(o.ignoreMessages, fqn)
+	return found
 }

interceptors/protovalidate/protovalidate.go

@@ -17,25 +17,17 @@ import (
 // UnaryServerInterceptor returns a new unary server interceptor that validates incoming messages.
 // If the request is invalid, clients may access a structured representation of the validation failure as an error detail.
 func UnaryServerInterceptor(validator *protovalidate.Validator, opts ...Option) grpc.UnaryServerInterceptor {
+	o := evaluateOpts(opts)
+
 	return func(
 		ctx context.Context,
 		req interface{},
 		info *grpc.UnaryServerInfo,
 		handler grpc.UnaryHandler,
 	) (resp interface{}, err error) {
-		o := evaluateOpts(opts)
-		switch msg := req.(type) {
-		case proto.Message:
-			if o.shouldIgnoreMessage(msg.ProtoReflect().Type()) {
-				break
-			}
-			if err = validator.Validate(msg); err != nil {
-				return nil, validationErrToStatus(err).Err()
-			}
-		default:
-			return nil, errors.New("unsupported message type")
+		if err := validateMsg(req, validator, o); err != nil {
+			return nil, err
 		}
-
 		return handler(ctx, req)
 	}
 }
@@ -49,66 +41,50 @@ func StreamServerInterceptor(validator *protovalidate.Validator, opts ...Option)
 		info *grpc.StreamServerInfo,
 		handler grpc.StreamHandler,
 	) error {
-		ctx := stream.Context()
+		return handler(srv, &wrappedServerStream{
+			ServerStream: stream,
+			validator:    validator,
+			options:      evaluateOpts(opts),
+		})
+	}
+}
 
-		wrapped := wrapServerStream(stream)
-		wrapped.wrappedContext = ctx
-		wrapped.validator = validator
-		wrapped.options = evaluateOpts(opts)
+// wrappedServerStream is a thin wrapper around grpc.ServerStream that allows modifying context.
+type wrappedServerStream struct {
+	grpc.ServerStream
 
-		return handler(srv, wrapped)
-	}
+	validator *protovalidate.Validator
+	options   *options
 }
 
 func (w *wrappedServerStream) RecvMsg(m interface{}) error {
 	if err := w.ServerStream.RecvMsg(m); err != nil {
 		return err
 	}
+	return validateMsg(m, w.validator, w.options)
+}
 
+func validateMsg(m interface{}, validator *protovalidate.Validator, opts *options) error {
 	msg, ok := m.(proto.Message)
 	if !ok {
 		return errors.New("unsupported message type")
 	}
-	if w.options.shouldIgnoreMessage(msg.ProtoReflect().Type()) {
+	if opts.shouldIgnoreMessage(msg.ProtoReflect().Descriptor().FullName()) {
 		return nil
 	}
-	if err := w.validator.Validate(msg); err != nil {
-		return validationErrToStatus(err).Err()
+	err := validator.Validate(msg)
+	if err == nil {
+		return nil
 	}
-
-	return nil
-}
-
-// wrappedServerStream is a thin wrapper around grpc.ServerStream that allows modifying context.
-type wrappedServerStream struct {
-	grpc.ServerStream
-	// wrappedContext is the wrapper's own Context. You can assign it.
-	wrappedContext context.Context
-
-	validator *protovalidate.Validator
-	options   *options
-}
-
-// Context returns the wrapper's WrappedContext, overwriting the nested grpc.ServerStream.Context()
-func (w *wrappedServerStream) Context() context.Context {
-	return w.wrappedContext
-}
-
-// wrapServerStream returns a ServerStream that has the ability to overwrite context.
-func wrapServerStream(stream grpc.ServerStream) *wrappedServerStream {
-	return &wrappedServerStream{ServerStream: stream, wrappedContext: stream.Context()}
-}
-
-func validationErrToStatus(err error) *status.Status {
-	// Message is invalid.
 	if valErr := new(protovalidate.ValidationError); errors.As(err, &valErr) {
+		// Message is invalid.
 		st := status.New(codes.InvalidArgument, err.Error())
 		ds, detErr := st.WithDetails(valErr.ToProto())
 		if detErr != nil {
-			return st
+			return st.Err()
 		}
-		return ds
+		return ds.Err()
 	}
 	// CEL expression doesn't compile or type-check.
-	return status.New(codes.Unknown, err.Error())
+	return status.Error(codes.Unknown, err.Error())
 }