Add sourceinfo plugin

Makefile

@@ -35,9 +35,11 @@ test: build $(BIN)/jv ## Run unit tests
 golden: generate
 	rm -rf internal/testdata/pubsub
 	rm -rf internal/testdata/jsonschema
+	rm -rf internal/testdata/sourceinfo
 	buf build ./internal/proto -o -#format=json > ./internal/testdata/codegenrequest/input.json
 	go run internal/cmd/pubsub-generate-testdata/main.go internal/testdata/pubsub
 	go run internal/cmd/jsonschema-generate-testdata/main.go internal/testdata/jsonschema
+	go run internal/cmd/sourceinfo-generate-testdata/main.go internal/testdata/sourceinfo
 
 .PHONY: build
 build: generate ## Build all packages

cmd/protoc-gen-jsonschema/main.go

@@ -16,10 +16,10 @@ package main
 
 import (
 	"github.com/bufbuild/protoplugin"
-	"github.com/bufbuild/protoschema-plugins/internal/protoschema"
+	"github.com/bufbuild/protoschema-plugins/cmd"
 	"github.com/bufbuild/protoschema-plugins/internal/protoschema/plugin/pluginjsonschema"
 )
 
 func main() {
-	protoplugin.Main(protoplugin.HandlerFunc(pluginjsonschema.Handle), protoplugin.WithVersion(protoschema.Version()))
+	protoplugin.Main(protoplugin.HandlerFunc(pluginjsonschema.Handle), protoplugin.WithVersion(cmd.Version()))
 }

cmd/protoc-gen-pubsub/main.go

@@ -16,10 +16,10 @@ package main
 
 import (
 	"github.com/bufbuild/protoplugin"
-	"github.com/bufbuild/protoschema-plugins/internal/protoschema"
+	"github.com/bufbuild/protoschema-plugins/cmd"
 	"github.com/bufbuild/protoschema-plugins/internal/protoschema/plugin/pluginpubsub"
 )
 
 func main() {
-	protoplugin.Main(protoplugin.HandlerFunc(pluginpubsub.Handle), protoplugin.WithVersion(protoschema.Version()))
+	protoplugin.Main(protoplugin.HandlerFunc(pluginpubsub.Handle), protoplugin.WithVersion(cmd.Version()))
 }

cmd/protoc-gen-sourceinfo/main.go

@@ -0,0 +1,25 @@
+// Copyright 2024 Buf Technologies, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"github.com/bufbuild/protoplugin"
+	"github.com/bufbuild/protoschema-plugins/cmd"
+	"github.com/bufbuild/protoschema-plugins/internal/protoschema/plugin/pluginsourceinfo"
+)
+
+func main() {
+	protoplugin.Main(protoplugin.HandlerFunc(pluginsourceinfo.Handle), protoplugin.WithVersion(cmd.Version()))
+}

internal/protoschema/protoschema.go → cmd/version.go

@@ -12,13 +12,14 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package protoschema
+package cmd
 
 import (
 	"runtime/debug"
 	"strings"
 )
 
+// Version returns the version of the protoschema library.
 func Version() string {
 	buildInfo, ok := debug.ReadBuildInfo()
 	if ok && buildInfo != nil && buildInfo.Main.Version != "" {

internal/cmd/jsonschema-generate-testdata/main.go

@@ -1,4 +1,4 @@
-// Copyright 2023 Buf Technologies, Inc.
+// Copyright 2024 Buf Technologies, Inc.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.

internal/cmd/sourceinfo-generate-testdata/main.go

@@ -0,0 +1,91 @@
+// Copyright 2024 Buf Technologies, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+	"strings"
+
+	_ "github.com/bufbuild/protoschema-plugins/internal/gen/proto/bufext/cel/expr/conformance/proto3"
+	"github.com/bufbuild/protoschema-plugins/internal/protoschema/golden"
+	"github.com/bufbuild/protoschema-plugins/internal/protoschema/plugin/pluginsourceinfo"
+	"google.golang.org/protobuf/reflect/protoreflect"
+)
+
+func main() {
+	if err := run(); err != nil {
+		if errString := err.Error(); errString != "" {
+			_, _ = fmt.Fprintln(os.Stderr, errString)
+		}
+		os.Exit(1)
+	}
+}
+
+func run() error {
+	if len(os.Args) != 2 {
+		return fmt.Errorf("usage: %s [output dir]", os.Args[0])
+	}
+	outputDir := os.Args[1]
+	// Make sure the directory exists
+	if err := os.MkdirAll(outputDir, 0755); err != nil {
+		return err
+	}
+
+	testFiles, err := golden.GetTestFiles("./internal/testdata")
+	if err != nil {
+		return err
+	}
+	// TODO: Use normal the plugin to generate golden files
+	includePrefixes := []string{
+		filepath.FromSlash("buf/protoschema/test/"),
+		filepath.FromSlash("bufext/cel/expr/conformance/proto3/"),
+	}
+	err = nil
+	testFiles.RangeFiles(func(testDesc protoreflect.FileDescriptor) bool {
+		if !shouldIncludeFile(testDesc, includePrefixes) {
+			return true
+		}
+		fileName := pluginsourceinfo.GetSourceInfoPath(testDesc)
+		filePath := filepath.Join(outputDir, fileName)
+		// Create any missing directories
+		if err = os.MkdirAll(filepath.Dir(filePath), 0755); err != nil {
+			err = fmt.Errorf("failed to create directory for %s: %w", filePath, err)
+			return false
+		}
+		var data []byte
+		data, err = pluginsourceinfo.GenFileContents(testDesc)
+		if err != nil {
+			err = fmt.Errorf("failed to generate source info for %s: %w", testDesc.FullName(), err)
+			return false
+		}
+		if err = os.WriteFile(filePath, data, 0600); err != nil {
+			err = fmt.Errorf("failed to write source info to %s: %w", filePath, err)
+			return false
+		}
+		return true
+	})
+	return err
+}
+
+func shouldIncludeFile(fileDesc protoreflect.FileDescriptor, includePrefixes []string) bool {
+	for _, prefix := range includePrefixes {
+		if strings.HasPrefix(fileDesc.Path(), prefix) {
+			return true
+		}
+	}
+	return false
+}

internal/proto/buf/protoschema/test/v1/test_cases.proto

@@ -19,7 +19,9 @@ package buf.protoschema.test.v1;
 import "buf/validate/validate.proto";
 import "bufext/cel/expr/conformance/proto3/test_all_types.proto";
 
+// A message comment.
 message NestedReference {
+  // A field comment.
   bufext.cel.expr.conformance.proto3.TestAllTypes.NestedMessage nested_message = 1;
 }
 

internal/protoschema/golden/golden.go

@@ -23,13 +23,14 @@ import (
 	"google.golang.org/protobuf/encoding/protojson"
 	"google.golang.org/protobuf/reflect/protodesc"
 	"google.golang.org/protobuf/reflect/protoreflect"
+	"google.golang.org/protobuf/reflect/protoregistry"
 	"google.golang.org/protobuf/types/descriptorpb"
 	"google.golang.org/protobuf/types/dynamicpb"
 )
 
 // GetTestDescriptors returns the test descriptors that were generated from the ./internal/proto
 // directory.
-func GetTestDescriptors(testdataPath string) ([]protoreflect.MessageDescriptor, error) {
+func GetTestFileDescriptorSet(testdataPath string) (*descriptorpb.FileDescriptorSet, error) {
 	inputPath := filepath.Join(filepath.FromSlash(testdataPath), "codegenrequest", "input.json")
 	input, err := os.ReadFile(inputPath)
 	if err != nil {
@@ -39,9 +40,25 @@ func GetTestDescriptors(testdataPath string) ([]protoreflect.MessageDescriptor,
 	if err = (&protojson.UnmarshalOptions{DiscardUnknown: true}).Unmarshal(input, fdset); err != nil {
 		return nil, fmt.Errorf("failed to parse file descriptor set at %q: %w", inputPath, err)
 	}
+	return fdset, nil
+}
+
+func GetTestFiles(testdataPath string) (*protoregistry.Files, error) {
+	fdset, err := GetTestFileDescriptorSet(testdataPath)
+	if err != nil {
+		return nil, err
+	}
 	files, err := protodesc.NewFiles(fdset)
 	if err != nil {
-		return nil, fmt.Errorf("failed to link file descriptor set at %q: %w", inputPath, err)
+		return nil, fmt.Errorf("failed to link file descriptor set at %q: %w", testdataPath, err)
+	}
+	return files, nil
+}
+
+func GetTestDescriptors(testdataPath string) ([]protoreflect.MessageDescriptor, error) {
+	files, err := GetTestFiles(testdataPath)
+	if err != nil {
+		return nil, err
 	}
 	types := dynamicpb.NewTypes(files)
 

internal/protoschema/plugin/pluginsourceinfo/pluginsourceinfo.go

@@ -0,0 +1,66 @@
+// Copyright 2024 Buf Technologies, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package pluginsourceinfo
+
+import (
+	"context"
+	"fmt"
+	"strings"
+
+	"github.com/bufbuild/protoplugin"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/reflect/protodesc"
+	"google.golang.org/protobuf/reflect/protoreflect"
+)
+
+const FileExtension = ".sourceinfo.binpb"
+
+// GetSourceInfoPath returns the path to the source info file for the given file descriptor.
+func GetSourceInfoPath(fileDescriptor protoreflect.FileDescriptor) string {
+	path := fileDescriptor.Path()
+	path = strings.TrimSuffix(path, ".proto")
+	return fmt.Sprintf("%s%s", path, FileExtension)
+}
+
+// Handle implements protoplugin.Handler and is the main entry point for the plugin.
+func Handle(
+	_ context.Context,
+	_ protoplugin.PluginEnv,
+	responseWriter protoplugin.ResponseWriter,
+	request protoplugin.Request,
+) error {
+	fileDescriptors, err := request.FileDescriptorsToGenerate()
+	if err != nil {
+		return err
+	}
+	for _, fileDescriptor := range fileDescriptors {
+		// Write the YAML string to the response.
+		data, err := GenFileContents(fileDescriptor)
+		if err != nil {
+			return err
+		}
+		name := GetSourceInfoPath(fileDescriptor)
+		responseWriter.AddFile(name, string(data))
+	}
+
+	responseWriter.SetFeatureProto3Optional()
+	return nil
+}
+
+func GenFileContents(fileDescriptor protoreflect.FileDescriptor) ([]byte, error) {
+	// Convert the file descriptor to a descriptorpb.FileDescriptorProto.
+	fileDescProto := protodesc.ToFileDescriptorProto(fileDescriptor)
+	return proto.Marshal(fileDescProto.SourceCodeInfo)
+}