diff --git a/staging/src/k8s.io/code-generator/cmd/go-to-protobuf/protobuf/cmd.go b/staging/src/k8s.io/code-generator/cmd/go-to-protobuf/protobuf/cmd.go index ca21c76c344..51b75fc2d61 100644 --- a/staging/src/k8s.io/code-generator/cmd/go-to-protobuf/protobuf/cmd.go +++ b/staging/src/k8s.io/code-generator/cmd/go-to-protobuf/protobuf/cmd.go @@ -46,6 +46,7 @@ type Generator struct { Clean bool OnlyIDL bool KeepGogoproto bool + DropGogoGo bool SkipGeneratedRewrite bool DropEmbeddedFields string } @@ -65,6 +66,7 @@ func New() *Generator { }, ","), Packages: "", DropEmbeddedFields: "k8s.io/apimachinery/pkg/apis/meta/v1.TypeMeta", + DropGogoGo: true, } } @@ -79,6 +81,7 @@ func (g *Generator) BindFlags(flag *flag.FlagSet) { flag.BoolVar(&g.OnlyIDL, "only-idl", g.OnlyIDL, "If true, only generate the IDL for each package.") flag.BoolVar(&g.KeepGogoproto, "keep-gogoproto", g.KeepGogoproto, "If true, the generated IDL will contain gogoprotobuf extensions which are normally removed") flag.BoolVar(&g.SkipGeneratedRewrite, "skip-generated-rewrite", g.SkipGeneratedRewrite, "If true, skip fixing up the generated.pb.go file (debugging only).") + flag.BoolVar(&g.DropGogoGo, "drop-gogo-go", g.DropGogoGo, "Drop all references to gogo packages in generated code") flag.StringVar(&g.DropEmbeddedFields, "drop-embedded-fields", g.DropEmbeddedFields, "Comma-delimited list of embedded Go types to omit from generated protobufs") } @@ -99,6 +102,10 @@ func Run(g *Generator) { log.Fatalf("Both apimachinery-packages and packages are empty. At least one package must be specified.") } + if g.DropGogoGo && g.SkipGeneratedRewrite { + log.Fatalf("--drop-gogo-go=true and --skip-generated-rewrite=true are mutually exclusive") + } + // Build up a list of packages to load from all the inputs. Track the // special modifiers for each. NOTE: This does not support pkg/... syntax. type modifier struct { @@ -288,7 +295,7 @@ func Run(g *Generator) { // alter the generated protobuf file to remove the generated types (but leave the serializers) and rewrite the // package statement to match the desired package name - if err := RewriteGeneratedGogoProtobufFile(outputPath, p.ExtractGeneratedType, p.OptionalTypeName, buf.Bytes()); err != nil { + if err := RewriteGeneratedGogoProtobufFile(outputPath, p.ExtractGeneratedType, p.OptionalTypeName, buf.Bytes(), g.DropGogoGo); err != nil { log.Fatalf("Unable to rewrite generated %s: %v", outputPath, err) } diff --git a/staging/src/k8s.io/code-generator/cmd/go-to-protobuf/protobuf/parser.go b/staging/src/k8s.io/code-generator/cmd/go-to-protobuf/protobuf/parser.go index 32f9e0da274..d0c0d69a4f2 100644 --- a/staging/src/k8s.io/code-generator/cmd/go-to-protobuf/protobuf/parser.go +++ b/staging/src/k8s.io/code-generator/cmd/go-to-protobuf/protobuf/parser.go @@ -77,7 +77,7 @@ type ExtractFunc func(*ast.TypeSpec) bool // and should have its marshal functions adjusted to remove the 'Items' accessor. type OptionalFunc func(name string) bool -func RewriteGeneratedGogoProtobufFile(name string, extractFn ExtractFunc, optionalFn OptionalFunc, header []byte) error { +func RewriteGeneratedGogoProtobufFile(name string, extractFn ExtractFunc, optionalFn OptionalFunc, header []byte, dropGogo bool) error { return rewriteFile(name, header, func(fset *token.FileSet, file *ast.File) error { cmap := ast.NewCommentMap(fset, file, file.Comments) @@ -85,6 +85,22 @@ func RewriteGeneratedGogoProtobufFile(name string, extractFn ExtractFunc, option for _, d := range file.Decls { rewriteOptionalMethods(d, optionalFn) } + if dropGogo { + // transform references to gogo sort util + var oldSortImport string + var usedSort bool + for _, d := range file.Decls { + oldSortImport, usedSort = rewriteGogoSortImport(d) + if usedSort { + break + } + } + if usedSort { + for _, d := range file.Decls { + rewriteGogoSort(d, oldSortImport, "sort") + } + } + } // remove types that are already declared decls := []ast.Decl{} @@ -95,6 +111,10 @@ func RewriteGeneratedGogoProtobufFile(name string, extractFn ExtractFunc, option if dropEmptyImportDeclarations(d) { continue } + // remove all but required functions + if dropGogo && dropUnusedGo(d) { + continue + } decls = append(decls, d) } file.Decls = decls @@ -105,6 +125,133 @@ func RewriteGeneratedGogoProtobufFile(name string, extractFn ExtractFunc, option }) } +// rewriteGogoSortImport rewrites an import of "github.com/gogo/protobuf/sortkeys" to "sort", +// and returns the original package alias and true if the rewrite occurred. +// Returns "", false if the decl is not an import decl, or does not contain an import of "github.com/gogo/protobuf/sortkeys", +func rewriteGogoSortImport(decl ast.Decl) (string, bool) { + t, ok := decl.(*ast.GenDecl) + if !ok { + return "", false + } + if t.Tok != token.IMPORT { + return "", false + } + for _, s := range t.Specs { + if spec, ok := s.(*ast.ImportSpec); ok { + if spec.Path != nil && spec.Path.Value == `"github.com/gogo/protobuf/sortkeys"` { + // switch gogo sort to stdlib sort + spec.Path.Value = `"sort"` + oldName := "sortkeys" + if spec.Name != nil { + oldName = spec.Name.Name + } + spec.Name = nil + return oldName, true + } + } + } + return "", false +} + +// rewriteGogoSort walks the AST, replacing use of the oldSortImport package with newSortImport +func rewriteGogoSort(decl ast.Decl, oldSortImport, newSortImport string) { + t, ok := decl.(*ast.FuncDecl) + if !ok { + return + } + ast.Walk(replacePackageVisitor{oldPackage: oldSortImport, newPackage: newSortImport}, t.Body) +} + +// keepFuncs is an allowlist of top-level func decls we should keep +var keepFuncs = map[string]bool{ + // generated helpers + "sovGenerated": true, + "sozGenerated": true, + "skipGenerated": true, + "encodeVarintGenerated": true, + "valueToStringGenerated": true, + + // unmarshal + "Reset": true, + "ProtoMessage": true, + "Unmarshal": true, + + // marshal + "Size": true, + "Marshal": true, + "MarshalTo": true, + "MarshalToSizedBuffer": true, + + // other widely used methods + "String": true, +} + +// keepVars is an allowlist of top-level var decls we should keep +var keepVars = map[string]bool{ + "ErrInvalidLengthGenerated": true, + "ErrIntOverflowGenerated": true, + "ErrUnexpectedEndOfGroupGenerated": true, +} + +// dropUnusedGo returns true if the top-level decl should be dropped. +// Has the following behavior for different decl types: +// * import: decl is rewritten to drop gogo package imports. Returns true if all imports in the decl were gogo imports, false if non-gogo imports remain. +// * var: decl is rewritten to drop vars not in the keepVars allowlist. Returns true if all vars in the decl were removed, false if allowlisted vars remain. +// * const: returns true +// * type: returns true +// * func: returns true if the func is not in the keepFuncs allowlist and should be dropped. +// * other: returns false +func dropUnusedGo(decl ast.Decl) bool { + switch t := decl.(type) { + case *ast.GenDecl: + switch t.Tok { + case token.IMPORT: + specs := []ast.Spec{} + for _, s := range t.Specs { + if spec, ok := s.(*ast.ImportSpec); ok { + if spec.Path == nil || !strings.HasPrefix(spec.Path.Value, `"github.com/gogo/protobuf/`) { + specs = append(specs, spec) + } + } + } + if len(specs) == 0 { + return true + } + t.Specs = specs + return false + case token.CONST: + // drop all const declarations + return true + case token.VAR: + specs := []ast.Spec{} + for _, s := range t.Specs { + if spec, ok := s.(*ast.ValueSpec); ok { + if keepVars[spec.Names[0].Name] { + specs = append(specs, spec) + } + } + } + if len(specs) == 0 { + return true + } + t.Specs = specs + return false + case token.TYPE: + // drop all type declarations + return true + } + case *ast.FuncDecl: + name := "" + if t.Name != nil { + name = t.Name.Name + } + return !keepFuncs[name] + default: + return false + } + return false +} + // rewriteOptionalMethods makes specific mutations to marshaller methods that belong to types identified // as being "optional" (they may be nil on the wire). This allows protobuf to serialize a map or slice and // properly discriminate between empty and nil (which is not possible in protobuf). @@ -451,3 +598,19 @@ func updateStructTags(decl ast.Decl, structTags map[string]map[string]string, to } return errs } + +type replacePackageVisitor struct { + oldPackage string + newPackage string +} + +// Visit walks the provided node, transforming references to the old package to the new package. +func (v replacePackageVisitor) Visit(n ast.Node) ast.Visitor { + if e, ok := n.(*ast.SelectorExpr); ok { + if i, ok := e.X.(*ast.Ident); ok && i.Name == v.oldPackage { + i.Name = v.newPackage + } + return nil + } + return v +}