From 66bb80c77433c8c9ea064da8df2239f8bd0fdb9c Mon Sep 17 00:00:00 2001
From: Jason Vigil <jasonvigil@google.com>
Date: Tue, 10 Dec 2024 23:55:06 +0000
Subject: [PATCH] feat: Add capability to replace proto contents in patch-proto

Use this new capability to codify the Apigee proto hack that was
originally made (manually) in
https://github.com/GoogleCloudPlatform/k8s-config-connector/pull/3183.

This addresses the suggested change in
https://github.com/GoogleCloudPlatform/k8s-config-connector/pull/3183#pullrequestreview-2447281185.
---
 mockgcp/apply-proto-patches.sh    |  11 +++
 mockgcp/tools/patch-proto/go.mod  |   2 +
 mockgcp/tools/patch-proto/go.sum  |   4 +
 mockgcp/tools/patch-proto/main.go | 136 +++++++++++++++++++++++++-----
 4 files changed, 130 insertions(+), 23 deletions(-)

diff --git a/mockgcp/apply-proto-patches.sh b/mockgcp/apply-proto-patches.sh
index 8773da7530..d571295ff0 100755
--- a/mockgcp/apply-proto-patches.sh
+++ b/mockgcp/apply-proto-patches.sh
@@ -67,3 +67,14 @@ cd tools/patch-proto
 #   // Optional. The configuration for Private Service Connect (PSC) for the cluster.
 #   PscConfig psc_config = 31 [(google.api.field_behavior) = OPTIONAL];
 # EOF
+
+go run . --file ${REPO_ROOT}/mockgcp/apis/mockgcp/cloud/apigee/v1/service.proto --service "ProjectsServer" --mode "replace" <<EOF
+
+  // Provisions a new Apigee organization with a functioning runtime. This is the standard way to create trial organizations for a free Apigee trial.
+  rpc ProvisionOrganizationProject(ProvisionOrganizationProjectRequest) returns (.google.longrunning.Operation) {
+    option (google.api.http) = {
+      post: "/v1/{name=projects/*}:provisionOrganization"
+      body: "project"
+    };
+  };
+EOF
\ No newline at end of file
diff --git a/mockgcp/tools/patch-proto/go.mod b/mockgcp/tools/patch-proto/go.mod
index 72b9654590..25f31c1bf7 100644
--- a/mockgcp/tools/patch-proto/go.mod
+++ b/mockgcp/tools/patch-proto/go.mod
@@ -14,4 +14,6 @@ require (
 	github.com/go-logr/logr v1.4.1 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
+	github.com/thediveo/enumflag/v2 v2.0.5 // indirect
+	golang.org/x/exp v0.0.0-20231006140011-7918f672742d // indirect
 )
diff --git a/mockgcp/tools/patch-proto/go.sum b/mockgcp/tools/patch-proto/go.sum
index e2d86f2270..f662c91e72 100644
--- a/mockgcp/tools/patch-proto/go.sum
+++ b/mockgcp/tools/patch-proto/go.sum
@@ -20,6 +20,10 @@ github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSS
 github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.7.4 h1:wZRexSlwd7ZXfKINDLsO4r7WBt3gTKONc6K/VesHvHM=
 github.com/stretchr/testify v1.7.4/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
+github.com/thediveo/enumflag/v2 v2.0.5 h1:VJjvlAqUb6m6mxOrB/0tfBJI0Kvi9wJ8ulh38xK87i8=
+github.com/thediveo/enumflag/v2 v2.0.5/go.mod h1:0NcG67nYgwwFsAvoQCmezG0J0KaIxZ0f7skg9eLq1DA=
+golang.org/x/exp v0.0.0-20231006140011-7918f672742d h1:jtJma62tbqLibJ5sFQz8bKtEM8rJBtfilJ2qTU199MI=
+golang.org/x/exp v0.0.0-20231006140011-7918f672742d/go.mod h1:ldy0pHrwJyGW56pPQzzkH36rKxoZW1tw7ZJpeKx+hdo=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
diff --git a/mockgcp/tools/patch-proto/main.go b/mockgcp/tools/patch-proto/main.go
index c85b61a46d..653ccf3bc5 100644
--- a/mockgcp/tools/patch-proto/main.go
+++ b/mockgcp/tools/patch-proto/main.go
@@ -28,6 +28,7 @@ import (
 	sitter "github.com/smacker/go-tree-sitter"
 	"github.com/smacker/go-tree-sitter/protobuf"
 	"github.com/spf13/cobra"
+	"github.com/thediveo/enumflag/v2"
 )
 
 func main() {
@@ -43,35 +44,67 @@ func run(ctx context.Context) error {
 
 	cmd := cobra.Command{
 		Use:   "proto-patch",
-		Short: "patches a message in a proto file",
-		Long:  `Inserts the contents of stdin at the end of the message in the proto file.`,
+		Short: "patches a proto file",
+		Long:  `Patches the contents of stdin into a proto file.`,
 		RunE: func(cmd *cobra.Command, args []string) error {
-			insert, err := io.ReadAll(os.Stdin)
+			patch, err := io.ReadAll(os.Stdin)
 			if err != nil {
 				return fmt.Errorf("reading stdin: %w", err)
 			}
-			opt.Insertion = insert
+			opt.Patch = patch
 
 			return RunProtoPatch(ctx, opt)
 		},
 	}
 	cmd.Flags().StringVar(&opt.ProtoPath, "file", opt.ProtoPath, "path to proto file to patch")
-	cmd.Flags().StringVar(&opt.Message, "message", opt.Message, "message file to patch")
+	cmd.Flags().VarP(
+		enumflag.New(&opt.Mode, "mode", ProtoPatchModeIDs, enumflag.EnumCaseInsensitive),
+		"mode", "m",
+		"patch mode; can be 'append' or 'replace'")
+
+	cmd.Flags().StringVar(&opt.Message, "message", opt.Message, "message name to patch")
+	cmd.Flags().StringVar(&opt.Service, "service", opt.Service, "service name to patch")
+
+	cmd.MarkFlagsOneRequired("message", "service")
+	cmd.MarkFlagsMutuallyExclusive("message", "service")
 
 	return cmd.Execute()
 }
 
+type ProtoPatchMode enumflag.Flag
+
+const (
+	ProtoPatchModeAppend ProtoPatchMode = iota
+	ProtoPatchModeReplace
+)
+
+var ProtoPatchModeIDs = map[ProtoPatchMode][]string{
+	ProtoPatchModeAppend:  {"append"},
+	ProtoPatchModeReplace: {"replace"},
+}
+
 type ProtoPatchOptions struct {
 	ProtoPath string
 	Message   string
-	Insertion []byte
+	Service   string
+	Mode      ProtoPatchMode
+	Patch     []byte
 }
 
 func RunProtoPatch(ctx context.Context, opt ProtoPatchOptions) error {
-	x := &insertPatchIntoMessage{
-		Message:   opt.Message,
-		Insertion: opt.Insertion,
+	x := &patchProto{
+		Patch: opt.Patch,
+		Mode:  opt.Mode,
 	}
+
+	if opt.Message != "" {
+		x.Id = ProtoIdentifierMessage
+		x.Name = opt.Message
+	} else if opt.Service != "" {
+		x.Id = ProtoIdentifierService
+		x.Name = opt.Service
+	}
+
 	protoPath := opt.ProtoPath
 
 	srcBytes, err := os.ReadFile(protoPath)
@@ -96,7 +129,7 @@ func RunProtoPatch(ctx context.Context, opt ProtoPatchOptions) error {
 	}
 
 	if x.Out == nil {
-		return fmt.Errorf("message %q not found in file %q", x.Message, protoPath)
+		return fmt.Errorf("identifier %q not found in file %q", x.Name, protoPath)
 	}
 	if err := os.WriteFile(protoPath, x.Out, 0644); err != nil {
 		return fmt.Errorf("writing to file %q: %w", protoPath, err)
@@ -105,16 +138,27 @@ func RunProtoPatch(ctx context.Context, opt ProtoPatchOptions) error {
 	return nil
 }
 
-type insertPatchIntoMessage struct {
-	Source    []byte
-	Insertion []byte
-	Message   string
-	Errors    []error
+type ProtoIdentifier string
+
+const (
+	ProtoIdentifierMessage ProtoIdentifier = "message"
+	ProtoIdentifierService ProtoIdentifier = "service"
+)
+
+type patchProto struct {
+	Mode ProtoPatchMode
+	Id   ProtoIdentifier
+
+	Source []byte
+	Patch  []byte
+	Name   string
+
+	Errors []error
 
 	Out []byte
 }
 
-func (x *insertPatchIntoMessage) VisitNode(depth int, cursor *sitter.TreeCursor) {
+func (x *patchProto) VisitNode(depth int, cursor *sitter.TreeCursor) {
 	node := cursor.CurrentNode()
 
 	// fmt.Printf("%s[%d:%s] %s\n", strings.Repeat("  ", depth), node.Symbol(), node.Type(), node.Content(x.Source))
@@ -145,7 +189,16 @@ func (x *insertPatchIntoMessage) VisitNode(depth int, cursor *sitter.TreeCursor)
 	case "message":
 		// e.g. message MyMessage { ... }
 		descend = false
-		x.VisitMessage(depth, cursor.CurrentNode())
+		if x.Id == ProtoIdentifierMessage {
+			x.VisitMessage(depth, cursor.CurrentNode())
+		}
+
+	case "service":
+		// e.g. service MyService { ... }
+		descend = false
+		if x.Id == ProtoIdentifierService {
+			x.VisitService(depth, cursor.CurrentNode())
+		}
 
 	default:
 		x.Errors = append(x.Errors, fmt.Errorf("unknown top-level node %q", protobuf.GetLanguage().SymbolName(node.Symbol())))
@@ -162,7 +215,7 @@ func (x *insertPatchIntoMessage) VisitNode(depth int, cursor *sitter.TreeCursor)
 	}
 }
 
-func (x *insertPatchIntoMessage) VisitMessage(depth int, node *sitter.Node) {
+func (x *patchProto) VisitMessage(depth int, node *sitter.Node) {
 	klog.V(2).Infof("%s[%d:%s] %s\n", strings.Repeat("  ", depth), node.Symbol(), node.Type(), node.Content(x.Source))
 
 	messageName := ""
@@ -182,7 +235,7 @@ func (x *insertPatchIntoMessage) VisitMessage(depth int, node *sitter.Node) {
 		}
 	}
 
-	if messageName == x.Message {
+	if messageName == x.Name {
 		if messageBody == nil {
 			x.Errors = append(x.Errors, fmt.Errorf("could not find message definition for message %q", messageName))
 			return
@@ -190,13 +243,50 @@ func (x *insertPatchIntoMessage) VisitMessage(depth int, node *sitter.Node) {
 
 		var out bytes.Buffer
 		out.Write(x.Source[:messageBody.StartByte()])
-		messageBodyContents := string(x.Source[messageBody.StartByte():messageBody.EndByte()])
-		messageBodyContents = strings.TrimSuffix(messageBodyContents, "}")
-		out.WriteString(messageBodyContents)
-		out.Write(x.Insertion)
+		if x.Mode == ProtoPatchModeAppend {
+			messageBodyContents := string(x.Source[messageBody.StartByte():messageBody.EndByte()])
+			messageBodyContents = strings.TrimSuffix(messageBodyContents, "}")
+			out.WriteString(messageBodyContents)
+		} else if x.Mode == ProtoPatchModeReplace {
+			out.WriteString("{\n")
+		}
+		out.Write(x.Patch)
 		out.WriteString("\n}")
 		out.Write(x.Source[messageBody.EndByte():])
 
 		x.Out = out.Bytes()
 	}
 }
+
+func (x *patchProto) VisitService(depth int, node *sitter.Node) {
+	klog.V(2).Infof("%s[%d:%s] %s\n", strings.Repeat("  ", depth), node.Symbol(), node.Type(), node.Content(x.Source))
+
+	childCount := int(node.ChildCount())
+	var serviceName *sitter.Node
+	var serviceBodyIdx int
+	for i := 0; i < childCount; i++ {
+		child := node.Child(i)
+		if protobuf.GetLanguage().SymbolName(child.Symbol()) == "service_name" {
+			serviceName = child
+			serviceBodyIdx = i + 1
+			break
+		}
+	}
+
+	if serviceName.Content(x.Source) == x.Name {
+		serviceBodyStart := node.Child(serviceBodyIdx)
+		lastChild := node.Child(childCount - 1)
+
+		var out bytes.Buffer
+		out.Write(x.Source[:serviceBodyStart.EndByte()])
+		if x.Mode == ProtoPatchModeAppend {
+			out.Write(x.Source[serviceBodyStart.EndByte():lastChild.StartByte()])
+		}
+		out.WriteString("\n")
+		out.Write(x.Patch)
+		out.WriteString(string(x.Source[lastChild.StartByte():lastChild.EndByte()]))
+		out.Write(x.Source[lastChild.EndByte():])
+
+		x.Out = out.Bytes()
+	}
+}