Add support for rpc.Code annotations in Zanzibar

codegen/method.go

@@ -45,6 +45,8 @@ const (
 	// The annotated method should have one and only one argument.
 	AntHTTPReqDefBoxed = "%s.http.req.def"
 	antHTTPResNoBody   = "%s.http.res.body.disallow"
+	// RPCCode is a rpc code type that allows to read the annotation in thrift and convert to a corresponding header
+	RPCCode = "rpc.Code"
 )
 
 const queryAnnotationPrefix = "query."
@@ -65,6 +67,8 @@ type ExceptionSpec struct {
 
 	StatusCode       StatusCode
 	IsBodyDisallowed bool
+	IsRPCCodeSet     bool
+	RPCCodeValue     string
 }
 
 // HeaderFieldInfo contains information about where to store
@@ -165,6 +169,7 @@ type annotations struct {
 	Handler         string
 	HTTPReqDefBoxed string
 	HTTPResNoBody   string
+	RPCCode         string
 }
 
 // StructSpec specifies a Go struct to be generated.
@@ -210,6 +215,7 @@ func NewMethod(
 		Meta:            fmt.Sprintf(antMeta, ant),
 		Handler:         fmt.Sprintf(antHandler, ant),
 		HTTPReqDefBoxed: fmt.Sprintf(AntHTTPReqDefBoxed, ant),
+		RPCCode:         fmt.Sprintf(RPCCode, ant),
 		HTTPResNoBody:   fmt.Sprintf(antHTTPResNoBody, ant),
 	}
 
@@ -416,14 +422,19 @@ func (ms *MethodSpec) setExceptions(
 		}
 
 		bodyDisallowed := ms.isBodyDisallowed(e)
+		rpcSet := ms.isRPCHeaderSet(e)
+		rpcValue := ms.extractRPCCodeValue(e)
 		if !ms.WantAnnot {
 			exception := ExceptionSpec{
 				StructSpec: StructSpec{
 					Type: typeName,
 					Name: e.Name,
 				},
 				IsBodyDisallowed: bodyDisallowed,
+				RPCCodeValue:     rpcValue,
+				IsRPCCodeSet:     rpcSet,
 			}
+
 			ms.Exceptions[i] = exception
 			ms.ExceptionsIndex[e.Name] = exception
 			if _, exists := ms.ExceptionsByStatusCode[exception.StatusCode.Code]; !exists {
@@ -455,6 +466,8 @@ func (ms *MethodSpec) setExceptions(
 				Message: e.Name,
 			},
 			IsBodyDisallowed: bodyDisallowed,
+			IsRPCCodeSet:     rpcSet,
+			RPCCodeValue:     rpcValue,
 		}
 		ms.Exceptions[i] = exception
 		ms.ExceptionsIndex[e.Name] = exception
@@ -1580,6 +1593,15 @@ func (ms *MethodSpec) isBodyDisallowed(f *compile.FieldSpec) bool {
 	return ok && val == "true"
 }
 
+func (ms *MethodSpec) isRPCHeaderSet(f *compile.FieldSpec) bool {
+	_, ok := f.Annotations[RPCCode]
+	return ok
+}
+
+func (ms *MethodSpec) extractRPCCodeValue(f *compile.FieldSpec) string {
+	return f.Annotations[RPCCode]
+}
+
 func headers(annotation string) []string {
 	if annotation == "" {
 		return nil

codegen/templates/endpoint.tmpl

@@ -234,6 +234,9 @@ func (h *{{$handlerName}}) HandleRequest(
 		{{end -}}
 		{{range $idx, $exception := .Exceptions}}
 		case *{{$exception.Type}}:
+		    {{if $exception.IsRPCCodeSet -}}
+            cliRespHeaders.Add("$rpc$-application-error-code", "{{$exception.RPCCodeValue}}")
+            {{end -}}
 			{{if $exception.IsBodyDisallowed -}}
 			res.WriteJSONBytes({{$exception.StatusCode.Code}}, cliRespHeaders, nil)
 			{{else -}}

codegen/templates/tchannel_endpoint.tmpl

@@ -38,7 +38,7 @@ import (
 {{- $clientID := .ClientID }}
 
 {{with .Method -}}
-// New{{$handlerName}} creates a handler to be registered with a thrift server.
+// New{{$handlerName}} creates a simple handler to be registered with a thrift server.
 func New{{$handlerName}}(deps *module.Dependencies) *{{$handlerName}} {
 	handler := &{{$handlerName}}{
 		Deps: deps,
@@ -159,8 +159,11 @@ func (h *{{$handlerName}}) Handle(
 		if err != nil {
 			switch v := err.(type) {
 			{{$method := .Name -}}
-			{{range .Exceptions -}}
+			{{range $idx, $exception := .Exceptions}}
 				case *{{.Type}}:
+					{{if $exception.IsRPCCodeSet -}}
+          		    resHeaders["$rpc$-application-error-code"] =  "{{$exception.RPCCodeValue}}"
+           			{{end -}}
 					ctxWithError := zanzibar.WithScopeTags(ctx, map[string]string{
 						"app-error": "{{.Type}}",
 					})

examples/example-gateway/idl/endpoints-idl/endpoints/bar/bar.thrift

@@ -81,7 +81,7 @@ service Bar {
     string helloWorld(
     ) throws (
        1: BarException barException (zanzibar.http.status = "403")
-       2: SeeOthersRedirection seeOthersRedirection (zanzibar.http.status = "303", zanzibar.http.res.body.disallow = "true")
+       2: SeeOthersRedirection seeOthersRedirection (zanzibar.http.status = "303", zanzibar.http.res.body.disallow = "true", rpc.Code = "UNAUTHORIZED")
    ) (
        zanzibar.http.method = "GET"
        zanzibar.http.path = "/bar/hello"