From 1fa8ab0ed5293e1d77b8145272f4b80911a18c87 Mon Sep 17 00:00:00 2001
From: Damien Neil <dneil@google.com>
Date: Thu, 27 Sep 2018 15:51:05 -0700
Subject: [PATCH] protogen: add an option to rewrite import paths

This allows us to implement the import_prefix parameter in the v1
protoc-gen-go.

Drop support for import_prefix in protogen, and explicitly produce an
error if it is used in the v2 protoc-gen-go.

Change-Id: I66136b6b3affa3c0e9a93dc565619c90c42c0ecc
Reviewed-on: https://go-review.googlesource.com/138257
Reviewed-by: Joe Tsai <thebrokentoaster@gmail.com>
---
 cmd/protoc-gen-go/main.go | 16 +++++++++++-----
 protogen/protogen.go      | 19 ++++++++++++++++---
 protogen/protogen_test.go | 36 ++++++++++++++++++++++++++++++++++++
 3 files changed, 63 insertions(+), 8 deletions(-)

diff --git a/cmd/protoc-gen-go/main.go b/cmd/protoc-gen-go/main.go
index efd294b3..8c72be80 100644
--- a/cmd/protoc-gen-go/main.go
+++ b/cmd/protoc-gen-go/main.go
@@ -15,15 +15,21 @@ import (
 )
 
 func main() {
-	var flags flag.FlagSet
-	plugins := flags.String("plugins", "", "deprecated option")
-	opts := &protogen.Options{
-		ParamFunc: flags.Set,
-	}
+	var (
+		flags        flag.FlagSet
+		plugins      = flags.String("plugins", "", "deprecated option")
+		importPrefix = flags.String("import_prefix", "", "deprecated option")
+		opts         = &protogen.Options{
+			ParamFunc: flags.Set,
+		}
+	)
 	protogen.Run(opts, func(gen *protogen.Plugin) error {
 		if *plugins != "" {
 			return errors.New("protoc-gen-go: plugins are not supported; use 'protoc --go-grpc_out=...' to generate gRPC")
 		}
+		if *importPrefix != "" {
+			return errors.New("protoc-gen-go: import_prefix is not supported")
+		}
 		for _, f := range gen.Files {
 			if !f.Generate {
 				continue
diff --git a/protogen/protogen.go b/protogen/protogen.go
index 87e643cf..80e7dc13 100644
--- a/protogen/protogen.go
+++ b/protogen/protogen.go
@@ -100,6 +100,7 @@ type Plugin struct {
 	enumsByName    map[protoreflect.FullName]*Enum
 	pathType       pathType
 	genFiles       []*GeneratedFile
+	opts           *Options
 	err            error
 }
 
@@ -129,6 +130,11 @@ type Options struct {
 	//     if *value { ... }
 	//   })
 	ParamFunc func(name, value string) error
+
+	// ImportRewriteFunc is called with the import path of each package
+	// imported by a generated file. It returns the import path to use
+	// for this package.
+	ImportRewriteFunc func(GoImportPath) GoImportPath
 }
 
 // New returns a new Plugin.
@@ -144,6 +150,7 @@ func New(req *pluginpb.CodeGeneratorRequest, opts *Options) (*Plugin, error) {
 		fileReg:        protoregistry.NewFiles(),
 		messagesByName: make(map[protoreflect.FullName]*Message),
 		enumsByName:    make(map[protoreflect.FullName]*Enum),
+		opts:           opts,
 	}
 
 	packageNames := make(map[string]GoPackageName) // filename -> package name
@@ -158,8 +165,6 @@ func New(req *pluginpb.CodeGeneratorRequest, opts *Options) (*Plugin, error) {
 		switch param {
 		case "":
 			// Ignore.
-		case "import_prefix":
-			// TODO
 		case "import_path":
 			packageImportPath = GoImportPath(value)
 		case "paths":
@@ -712,6 +717,7 @@ func newEnumValue(gen *Plugin, f *File, message *Message, enum *Enum, desc proto
 
 // A GeneratedFile is a generated file.
 type GeneratedFile struct {
+	gen              *Plugin
 	filename         string
 	goImportPath     GoImportPath
 	buf              bytes.Buffer
@@ -724,6 +730,7 @@ type GeneratedFile struct {
 // and import path.
 func (gen *Plugin) NewGeneratedFile(filename string, goImportPath GoImportPath) *GeneratedFile {
 	g := &GeneratedFile{
+		gen:              gen,
 		filename:         filename,
 		goImportPath:     goImportPath,
 		packageNames:     make(map[GoImportPath]GoPackageName),
@@ -876,8 +883,14 @@ func (g *GeneratedFile) Content() ([]byte, error) {
 		importPaths = append(importPaths, string(importPath))
 	}
 	sort.Strings(importPaths)
+	rewriteImport := func(importPath string) string {
+		if f := g.gen.opts.ImportRewriteFunc; f != nil {
+			return string(f(GoImportPath(importPath)))
+		}
+		return importPath
+	}
 	for _, importPath := range importPaths {
-		astutil.AddNamedImport(fset, file, string(g.packageNames[GoImportPath(importPath)]), importPath)
+		astutil.AddNamedImport(fset, file, string(g.packageNames[GoImportPath(importPath)]), rewriteImport(importPath))
 	}
 	for importPath := range g.manualImports {
 		if _, ok := g.packageNames[importPath]; ok {
diff --git a/protogen/protogen_test.go b/protogen/protogen_test.go
index 237226c0..6d8e5af4 100644
--- a/protogen/protogen_test.go
+++ b/protogen/protogen_test.go
@@ -307,6 +307,42 @@ got:
 	}
 }
 
+func TestImportRewrites(t *testing.T) {
+	gen, err := New(&pluginpb.CodeGeneratorRequest{}, &Options{
+		ImportRewriteFunc: func(i GoImportPath) GoImportPath {
+			return "prefix/" + i
+		},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	g := gen.NewGeneratedFile("foo.go", "golang.org/x/foo")
+	g.P("package foo")
+	g.P("var _ = ", GoIdent{GoName: "X", GoImportPath: "golang.org/x/bar"})
+	want := `package foo
+
+import bar "prefix/golang.org/x/bar"
+
+var _ = bar.X
+`
+	got, err := g.Content()
+	if err != nil {
+		t.Fatalf("g.Content() = %v", err)
+	}
+	if want != string(got) {
+		t.Fatalf(`want:
+==========
+%v
+==========
+
+got:
+==========
+%v
+==========`,
+			want, string(got))
+	}
+}
+
 // makeRequest returns a CodeGeneratorRequest for the given protoc inputs.
 //
 // It does this by running protoc with the current binary as the protoc-gen-go