Refactor again

private/bufpkg/bufcheck/bufcheckserver/internal/bufcheckserverhandle/lint.go

@@ -963,6 +963,9 @@ func handleLintProtovalidate(
 			args...,
 		)
 	}
+	// We create a new extension resolver using all of the files from the request, including
+	// import files. This is because there can be a case where a non-import file uses a predefined
+	// rule from an imported file.
 	extensionResolver, err := protoencoding.NewResolver(
 		slicesext.Map(
 			request.ProtosourceFiles(),
@@ -974,68 +977,35 @@ func handleLintProtovalidate(
 	if err != nil {
 		return err
 	}
-	// This for-loop checks that predefined rules have cel expressions that compile and adds
-	// the ones that compile to the extension resolver, as a side effect. These types are relied
-	// on to check the example values for fields.
-	for _, file := range request.ProtosourceFiles() {
-		// We check all predefined rules for all files and add them to the extension resolver
-		// if they compile, regardless if the file is an import or not. This is because a non-import
-		// file may use a predefined rule from an import file.
-		// However, we only add check annotations for non-import files.
-		if err := bufprotosource.ForEachMessage(
-			func(message bufprotosource.Message) error {
-				for _, extension := range message.Extensions() {
-					if err := buflintvalidate.CheckAndRegisterPredefinedRuleExtension(
-						addAnnotationFunc,
-						extension,
-						file.IsImport(),
-						extensionResolver,
-					); err != nil {
-						return err
-					}
-				}
-				return nil
-			},
-			file,
-		); err != nil {
-			return err
-		}
-		for _, extension := range file.Extensions() {
-			if err := buflintvalidate.CheckAndRegisterPredefinedRuleExtension(
-				addAnnotationFunc,
-				extension,
-				file.IsImport(),
-				extensionResolver,
-			); err != nil {
-				return err
-			}
-		}
-	}
+	// However, we only want to check non-import files, so we can use NewLintMessageRuleHandler
+	// and NewLintFieldRuleHandler utils to check messages and fields respectively.
 	if err := bufcheckserverutil.NewLintMessageRuleHandler(
 		func(
-			// The responseWriter is being passed in through the shared addAnnotationFunc, so we
-			// do not pass in responseWriter and request again. This should be addressed in a refactor.
 			_ bufcheckserverutil.ResponseWriter,
 			_ bufcheckserverutil.Request,
 			message bufprotosource.Message,
 		) error {
 			return buflintvalidate.CheckMessage(addAnnotationFunc, message)
-		}).Handle(ctx, nil, nil); err != nil {
+		},
+		// The responseWriter is being passed in through the shared addAnnotationFunc, so we
+		// do not pass in responseWriter again. This should be addressed in a refactor.
+	).Handle(ctx, nil, request); err != nil {
 		return err
 	}
-	if err := bufcheckserverutil.NewLintFieldRuleHandler(
+	return bufcheckserverutil.NewLintFieldRuleHandler(
 		func(
-			// The responseWriter is being passed in through the shared addAnnotationFunc, so we
-			// do not pass in responseWriter and request again. This should be addressed in a refactor.
 			_ bufcheckserverutil.ResponseWriter,
 			_ bufcheckserverutil.Request,
 			field bufprotosource.Field,
 		) error {
+			if err := buflintvalidate.CheckPredefinedRuleExtension(addAnnotationFunc, field, extensionResolver); err != nil {
+				return err
+			}
 			return buflintvalidate.CheckField(addAnnotationFunc, field, extensionResolver)
-		}).Handle(ctx, nil, nil); err != nil {
-		return err
-	}
-	return nil
+		},
+		// The responseWriter is being passed in through the shared addAnnotationFunc, so we
+		// do not pass in responseWriter again. This should be addressed in a refactor.
+	).Handle(ctx, nil, request)
 }
 
 // HandleLintRPCNoClientStreaming is a handle function.

private/bufpkg/bufcheck/bufcheckserver/internal/buflintvalidate/buflintvalidate.go

@@ -24,22 +24,10 @@ import (
 // https://buf.build/bufbuild/protovalidate/docs/v0.5.1:buf.validate#buf.validate.MessageConstraints
 const disabledFieldNumberInMesageConstraints = 1
 
-// CheckAndRegisterPredefinedRuleExtension checks whether an extension extending a protovalidate rule
-// is valid, checking that all of its CEL expressionus compile. If so, the extension type is appended to
-// the extension types passed in.
-func CheckAndRegisterPredefinedRuleExtension(
-	addAnnotationFunc func(bufprotosource.Descriptor, bufprotosource.Location, []bufprotosource.Location, string, ...interface{}),
-	field bufprotosource.Field,
-	fileIsImport bool,
-	extensionResolver protoencoding.Resolver,
-) error {
-	return checkAndRegisterPredefinedRuleExtension(addAnnotationFunc, field, fileIsImport, extensionResolver)
-}
-
 // CheckMessage validates that all rules on the message are valid, and any CEL expressions compile.
-//
-// addAnnotationFunc adds an annotation with the descriptor and location for check results.
+// It also checks all predefined rule extensions on the messages.
 func CheckMessage(
+	// addAnnotationFunc adds an annotation with the descriptor and location for check results.
 	addAnnotationFunc func(bufprotosource.Descriptor, bufprotosource.Location, []bufprotosource.Location, string, ...interface{}),
 	message bufprotosource.Message,
 ) error {
@@ -71,11 +59,21 @@ func CheckMessage(
 //  1. permit _some_ value and all example values, if any
 //  2. have a type compatible with the field it validates.
 //
-// addAnnotationFunc adds an annotation with the descriptor and location for check results.
+// This also checks all predefined rule extensions fields to ensure they compile.
 func CheckField(
+	// addAnnotationFunc adds an annotation with the descriptor and location for check results.
 	addAnnotationFunc func(bufprotosource.Descriptor, bufprotosource.Location, []bufprotosource.Location, string, ...interface{}),
 	field bufprotosource.Field,
 	extensionTypeResolver protoencoding.Resolver,
 ) error {
 	return checkField(addAnnotationFunc, field, extensionTypeResolver)
 }
+
+func CheckPredefinedRuleExtension(
+	// addAnnotationFunc adds an annotation with the descriptor and location for check results.
+	addAnnotationFunc func(bufprotosource.Descriptor, bufprotosource.Location, []bufprotosource.Location, string, ...interface{}),
+	field bufprotosource.Field,
+	extensionResolver protoencoding.Resolver,
+) error {
+	return checkPredefinedRuleExtension(addAnnotationFunc, field, extensionResolver)
+}

private/bufpkg/bufcheck/bufcheckserver/internal/buflintvalidate/predefined_rules.go

@@ -29,10 +29,9 @@ const (
 	celFieldNumberPath = int32(1)
 )
 
-func checkAndRegisterPredefinedRuleExtension(
+func checkPredefinedRuleExtension(
 	addAnnotationFunc func(bufprotosource.Descriptor, bufprotosource.Location, []bufprotosource.Location, string, ...interface{}),
 	extension bufprotosource.Field,
-	fileIsImport bool,
 	extensionResolver protoencoding.Resolver,
 ) error {
 	extensionDescriptor, err := extension.AsDescriptor()
@@ -94,16 +93,6 @@ func checkAndRegisterPredefinedRuleExtension(
 	if err != nil {
 		return err
 	}
-	// If the file is an import file, we only want to check that the CEL expression compiles,
-	// but we do not want to produce file annotations, so we set the addAnnotationFunc to a nop.
-	if fileIsImport {
-		addAnnotationFunc = func(
-			_ bufprotosource.Descriptor,
-			_ bufprotosource.Location,
-			_ []bufprotosource.Location,
-			_ string, _ ...interface{}) {
-		}
-	}
 	checkCEL(
 		celEnv,
 		predefinedConstraints.GetCel(),