refactor: use new protoc plugin for api v2 (#5798)

* refactor: use new protoc plugin for api v2

* simplify code

build/zitadel/generate-grpc.sh

@@ -18,8 +18,9 @@ protoc \
   --validate_out=lang=go:${GOPATH}/src \
   $(find ${PROTO_PATH} -iname *.proto)
 
-# install authoption proto compiler
+# install authoption and zitadel proto compiler
 go install ${ZITADEL_PATH}/internal/protoc/protoc-gen-auth
+go install ${ZITADEL_PATH}/internal/protoc/protoc-gen-zitadel
 
 # output folder for openapi v2
 mkdir -p ${OPENAPI_PATH}
@@ -79,7 +80,7 @@ protoc \
   --openapiv2_out ${OPENAPI_PATH} \
   --openapiv2_opt logtostderr=true \
   --openapiv2_opt allow_delete_body=true \
-  --auth_out=${GOPATH}/src \
+  --zitadel_out=${GOPATH}/src \
   --validate_out=lang=go:${GOPATH}/src \
   ${PROTO_PATH}/user/v2alpha/user_service.proto
 
@@ -91,7 +92,7 @@ protoc \
   --openapiv2_out ${OPENAPI_PATH} \
   --openapiv2_opt logtostderr=true \
   --openapiv2_opt allow_delete_body=true \
-  --auth_out=${GOPATH}/src \
+  --zitadel_out=${GOPATH}/src \
   --validate_out=lang=go:${GOPATH}/src \
   ${PROTO_PATH}/session/v2alpha/session_service.proto
 

internal/api/grpc/server/gateway.go

@@ -12,6 +12,7 @@ import (
 	"google.golang.org/grpc/credentials/insecure"
 	healthpb "google.golang.org/grpc/health/grpc_health_v1"
 	"google.golang.org/protobuf/encoding/protojson"
+	"google.golang.org/protobuf/proto"
 
 	client_middleware "github.com/zitadel/zitadel/internal/api/grpc/client/middleware"
 	"github.com/zitadel/zitadel/internal/api/grpc/server/middleware"
@@ -38,6 +39,7 @@ var (
 		runtime.WithMarshalerOption(runtime.MIMEWildcard, jsonMarshaler),
 		runtime.WithIncomingHeaderMatcher(headerMatcher),
 		runtime.WithOutgoingHeaderMatcher(runtime.DefaultHeaderMatcher),
+		runtime.WithForwardResponseOption(responseForwarder),
 	}
 
 	headerMatcher = runtime.HeaderMatcherFunc(
@@ -50,6 +52,15 @@ var (
 			return runtime.DefaultHeaderMatcher(header)
 		},
 	)
+
+	responseForwarder = func(ctx context.Context, w http.ResponseWriter, resp proto.Message) error {
+		t, ok := resp.(CustomHTTPResponse)
+		if ok {
+			// TODO: find a way to return a location header if needed w.Header().Set("location", t.Location())
+			w.WriteHeader(t.CustomHTTPCode())
+		}
+		return nil
+	}
 )
 
 type Gateway struct {
@@ -62,6 +73,10 @@ func (g *Gateway) Handler() http.Handler {
 	return addInterceptors(g.mux, g.http1HostName)
 }
 
+type CustomHTTPResponse interface {
+	CustomHTTPCode() int
+}
+
 type RegisterGatewayFunc func(ctx context.Context, mux *runtime.ServeMux, conn *grpc.ClientConn) error
 
 func CreateGatewayWithPrefix(ctx context.Context, g WithGatewayPrefix, port uint16, http1HostName string) (http.Handler, string, error) {

internal/api/grpc/server/middleware/auth_interceptor.go

@@ -11,6 +11,7 @@ import (
 	grpc_util "github.com/zitadel/zitadel/internal/api/grpc"
 	"github.com/zitadel/zitadel/internal/api/http"
 	"github.com/zitadel/zitadel/internal/telemetry/tracing"
+	object "github.com/zitadel/zitadel/pkg/grpc/object/v2alpha"
 )
 
 func AuthorizationInterceptor(verifier *authz.TokenVerifier, authConfig authz.Config) grpc.UnaryServerInterceptor {
@@ -34,8 +35,8 @@ func authorize(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo,
 	}
 
 	orgID := grpc_util.GetHeader(authCtx, http.ZitadelOrgID)
-	if o, ok := req.(AuthContext); ok {
-		orgID = o.AuthContext()
+	if o, ok := req.(OrganisationFromRequest); ok {
+		orgID = o.OrganisationFromRequest().GetOrgId()
 	}
 
 	ctxSetter, err := authz.CheckUserAuthorization(authCtx, req, authToken, orgID, verifier, authConfig, authOpt, info.FullMethod)
@@ -46,6 +47,6 @@ func authorize(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo,
 	return handler(ctxSetter(ctx), req)
 }
 
-type AuthContext interface {
-	AuthContext() string
+type OrganisationFromRequest interface {
+	OrganisationFromRequest() *object.Organisation
 }

internal/protoc/protoc-gen-zitadel/main.go

@@ -0,0 +1,147 @@
+package main
+
+import (
+	"bytes"
+	_ "embed"
+	"fmt"
+	"io"
+	"os"
+	"text/template"
+
+	"google.golang.org/protobuf/compiler/protogen"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/descriptorpb"
+	"google.golang.org/protobuf/types/pluginpb"
+
+	protoc_gen_zitadel "github.com/zitadel/zitadel/pkg/grpc/protoc/v2"
+)
+
+var (
+	//go:embed zitadel.pb.go.tmpl
+	zitadelTemplate []byte
+)
+
+type authMethods struct {
+	GoPackageName       string
+	ProtoPackageName    string
+	ServiceName         string
+	AuthOptions         []authOption
+	AuthContext         []authContext
+	CustomHTTPResponses []httpResponse
+}
+
+type authOption struct {
+	Name           string
+	Permission     string
+	CheckFieldName string
+}
+
+type authContext struct {
+	Name      string
+	OrgMethod string
+}
+
+type httpResponse struct {
+	Name string
+	Code int32
+}
+
+func main() {
+	input, _ := io.ReadAll(os.Stdin)
+	var req pluginpb.CodeGeneratorRequest
+	err := proto.Unmarshal(input, &req)
+	if err != nil {
+		panic(err)
+	}
+
+	opts := protogen.Options{}
+	plugin, err := opts.New(&req)
+	if err != nil {
+		panic(err)
+	}
+	plugin.SupportedFeatures = uint64(pluginpb.CodeGeneratorResponse_FEATURE_PROTO3_OPTIONAL)
+
+	tmpl := loadTemplate(zitadelTemplate)
+
+	for _, file := range plugin.Files {
+		methods := new(authMethods)
+		for _, service := range file.Services {
+			methods.ServiceName = service.GoName
+			methods.GoPackageName = string(file.GoPackageName)
+			methods.ProtoPackageName = *file.Proto.Package
+			for _, method := range service.Methods {
+				options := method.Desc.Options().(*descriptorpb.MethodOptions)
+				if options == nil {
+					continue
+				}
+				ext := proto.GetExtension(options, protoc_gen_zitadel.E_Options).(*protoc_gen_zitadel.Options)
+				if ext == nil {
+					continue
+				}
+				if ext.AuthOption != nil {
+					generateAuthOption(methods, ext.AuthOption, method)
+				}
+				if ext.HttpResponse != nil {
+					methods.CustomHTTPResponses = append(methods.CustomHTTPResponses, httpResponse{Name: string(method.Output.Desc.Name()), Code: ext.HttpResponse.SuccessCode})
+				}
+			}
+		}
+		if len(methods.AuthOptions) > 0 {
+			generateFile(tmpl, methods, file, plugin)
+		}
+	}
+
+	// Generate a response from our plugin and marshall as protobuf
+	stdout := plugin.Response()
+	out, err := proto.Marshal(stdout)
+	if err != nil {
+		panic(err)
+	}
+
+	// Write the response to stdout, to be picked up by protoc
+	_, err = fmt.Fprint(os.Stdout, string(out))
+	if err != nil {
+		panic(err)
+	}
+}
+
+func generateAuthOption(methods *authMethods, protoAuthOption *protoc_gen_zitadel.AuthOption, method *protogen.Method) {
+	methods.AuthOptions = append(methods.AuthOptions, authOption{Name: string(method.Desc.Name()), Permission: protoAuthOption.Permission})
+	if protoAuthOption.OrgField == "" {
+		return
+	}
+	orgMethod := buildAuthContextField(method.Input.Fields, protoAuthOption.OrgField)
+	if orgMethod != "" {
+		methods.AuthContext = append(methods.AuthContext, authContext{Name: string(method.Input.Desc.Name()), OrgMethod: orgMethod})
+	}
+}
+
+func generateFile(tmpl *template.Template, methods *authMethods, protoFile *protogen.File, plugin *protogen.Plugin) {
+	var buffer bytes.Buffer
+	err := tmpl.Execute(&buffer, &methods)
+	if err != nil {
+		panic(err)
+	}
+
+	filename := protoFile.GeneratedFilenamePrefix + ".pb.zitadel.go"
+	file := plugin.NewGeneratedFile(filename, ".")
+
+	_, err = file.Write(buffer.Bytes())
+	if err != nil {
+		panic(err)
+	}
+}
+
+func loadTemplate(templateData []byte) *template.Template {
+	return template.Must(template.New("").
+		Parse(string(templateData)))
+}
+
+func buildAuthContextField(fields []*protogen.Field, fieldName string) string {
+	for _, field := range fields {
+		if string(field.Desc.Name()) == fieldName {
+			return ".Get" + field.GoName + "()"
+		}
+	}
+	return ""
+}

internal/protoc/protoc-gen-zitadel/zitadel.pb.go.tmpl

@@ -0,0 +1,29 @@
+// Code generated by protoc-gen-zitadel. DO NOT EDIT.
+
+package {{.GoPackageName}}
+
+import (
+	"github.com/zitadel/zitadel/internal/api/authz"
+    {{if .AuthContext}}"github.com/zitadel/zitadel/pkg/grpc/object/v2alpha"{{end}}
+)
+
+var {{.ServiceName}}_AuthMethods = authz.MethodMapping {
+    {{ range $m := .AuthOptions}}
+        {{$.ServiceName}}_{{$m.Name}}_FullMethodName: authz.Option{
+            Permission: "{{$m.Permission}}",
+            CheckParam: "{{$m.CheckFieldName}}",
+        },
+    {{ end}}
+}
+
+{{ range $m := .AuthContext}}
+func (r *{{ $m.Name }}) OrganisationFromRequest() *object.Organisation {
+    return r{{$m.OrgMethod}}
+}
+{{ end }}
+
+{{ range $resp := .CustomHTTPResponses}}
+func (r *{{ $resp.Name }}) CustomHTTPCode() int {
+    return {{$resp.Code}}
+}
+{{ end }}

pkg/grpc/protoc/v2/doc.go

@@ -0,0 +1,3 @@
+// Package protoc contains the generated protobuf structs
+// the folder will be empty until the files are generated
+package protoc

proto/zitadel/protoc_gen_zitadel/v2/options.proto

@@ -0,0 +1,26 @@
+syntax = "proto3";
+
+package zitadel.protoc_gen_zitadel.v2;
+
+import "google/protobuf/descriptor.proto";
+
+option go_package = "github.com/zitadel/zitadel/pkg/grpc/protoc/v2;protoc";
+
+extend google.protobuf.MethodOptions {
+  Options options = 50001;
+}
+
+message Options {
+  AuthOption auth_option = 1;
+  CustomHTTPResponse http_response = 2;
+}
+
+message AuthOption {
+  reserved 2;
+  string permission = 1;
+  string org_field = 3;
+}
+
+message CustomHTTPResponse {
+  int32 success_code = 1;
+}

proto/zitadel/session/v2alpha/session_service.proto

@@ -2,7 +2,8 @@ syntax = "proto3";
 
 package zitadel.session.v2alpha;
 
-import "zitadel/options.proto";
+
+import "zitadel/protoc_gen_zitadel/v2/options.proto";
 import "zitadel/session/v2alpha/session.proto";
 import "google/api/annotations.proto";
 import "validate/validate.proto";
@@ -19,8 +20,10 @@ service SessionService {
       get: "/v2alpha/sessions/{id}"
     };
 
-    option (zitadel.v1.auth_option) = {
-      permission: "authenticated"
+    option (zitadel.protoc_gen_zitadel.v2.options) = {
+      auth_option: {
+        permission: "authenticated"
+      }
     };
   }
 }

proto/zitadel/user/v2alpha/user_service.proto

@@ -2,8 +2,8 @@ syntax = "proto3";
 
 package zitadel.user.v2alpha;
 
-import "zitadel/options.proto";
 import "zitadel/object/v2alpha/object.proto";
+import "zitadel/protoc_gen_zitadel/v2/options.proto";
 import "zitadel/user/v2alpha/email.proto";
 import "zitadel/user/v2alpha/password.proto";
 import "zitadel/user/v2alpha/user.proto";
@@ -82,8 +82,14 @@ service UserService {
       body: "*"
     };
 
-    option (zitadel.v1.auth_option) = {
-      permission: "user.write"
+    option (zitadel.protoc_gen_zitadel.v2.options) = {
+      auth_option: {
+        permission: "user.write"
+        org_field: "organisation"
+      }
+      http_response: {
+        success_code: 201
+      }
     };
 
     option (grpc.gateway.protoc_gen_openapiv2.options.openapiv2_operation) = {
@@ -105,8 +111,10 @@ service UserService {
       body: "*"
     };
 
-    option (zitadel.v1.auth_option) = {
-      permission: "authenticated"
+    option (zitadel.protoc_gen_zitadel.v2.options) = {
+      auth_option: {
+        permission: "authenticated"
+      }
     };
 
     option (grpc.gateway.protoc_gen_openapiv2.options.openapiv2_operation) = {
@@ -128,8 +136,10 @@ service UserService {
       body: "*"
     };
 
-    option (zitadel.v1.auth_option) = {
-      permission: "authenticated"
+    option (zitadel.protoc_gen_zitadel.v2.options) = {
+      auth_option: {
+        permission: "authenticated"
+      }
     };
 
     option (grpc.gateway.protoc_gen_openapiv2.options.openapiv2_operation) = {