Fix a couple streaming interceptor bugs

- The wrap function for server interceptors was working as if it were a
  client interceptor: expecting the stream to be the result of calling
  the endpoint or the next unary interceptor in the chain. It now
  properly wraps the stream in the raw payload (which is actually a
  special `EndpointStruct` with the stream and optionally the unary
  payload).
- The unary `Payload` methods of server interceptor info structs for a
  streaming method assumed that the raw payload would be the payload
  type rather than the special `EndpointStruct` type. Now, the `Payload`
  methods properly expect the raw payload to either be the `EndpointStruct`
  or the payload type in the case when the server interceptor info
  struct is also used for a client interceptor.

Author: douglaswth

codegen/service/interceptors.go

@@ -122,6 +122,7 @@ func interceptorFile(svc *Data, server bool) *codegen.File {
 			Data:   interceptors,
 			FuncMap: map[string]any{
 				"hasPrivateImplementationTypes": hasPrivateImplementationTypes,
+				"hasEndpointStruct":             hasEndpointStruct(server),
 			},
 		})
 	}
@@ -226,6 +227,17 @@ func hasPrivateImplementationTypes(interceptors []*InterceptorData) bool {
 	return false
 }
 
+// hasEndpointStruct returns a function that returns true if the method has an endpoint struct
+// if server is true, otherwise it returns false.
+func hasEndpointStruct(server bool) func(*MethodInterceptorData) bool {
+	if !server {
+		return func(*MethodInterceptorData) bool { return false }
+	}
+	return func(m *MethodInterceptorData) bool {
+		return m.ServerStream != nil && m.ServerStream.EndpointStruct != ""
+	}
+}
+
 // collectWrappedStreams returns a slice of streams to be wrapped by interceptor wrapper functions.
 func collectWrappedStreams(interceptors []*InterceptorData, server bool) []*StreamInterceptorData {
 	var (

codegen/service/service_data.go

@@ -227,8 +227,7 @@ type (
 		// function.
 		MustClose bool
 		// EndpointStruct is the name of the endpoint struct that holds a payload
-		// reference (if any) and the endpoint server stream. It is set only if the
-		// client sends a normal payload and server streams a result.
+		// reference (if any) and the endpoint server stream.
 		EndpointStruct string
 		// Kind is the kind of the stream (payload, result or bidirectional).
 		Kind expr.StreamKind
@@ -345,6 +344,9 @@ type (
 		// MustClose indicates whether the stream should implement the Close()
 		// function.
 		MustClose bool
+		// EndpointStruct is the name of the endpoint struct that holds a payload
+		// reference (if any) and the endpoint server stream.
+		EndpointStruct string
 	}
 
 	// AttributeData describes a single attribute.
@@ -1314,6 +1316,7 @@ func buildInterceptorMethodData(i *expr.InterceptorExpr, md *MethodData) *Method
 			RecvWithContextName: md.ServerStream.RecvWithContextName,
 			RecvTypeRef:         md.ServerStream.RecvTypeRef,
 			MustClose:           md.ServerStream.MustClose,
+			EndpointStruct:      md.ServerStream.EndpointStruct,
 		}
 	}
 	if md.ClientStream != nil {

codegen/service/templates/interceptors.go.tpl

@@ -29,13 +29,31 @@ func (info *{{ .Name }}Info) Payload() {{ .Name }}Payload {
 		switch info.Method() {
 			{{- range .Methods }}
 		case "{{ .MethodName }}":
+				{{- if hasEndpointStruct . }}
+			switch pay := info.RawPayload().(type) {
+			case *{{ .ServerStream.EndpointStruct }}:
+				return &{{ .PayloadAccess }}{payload: pay.Payload}
+			default:
+				return &{{ .PayloadAccess }}{payload: pay.({{ .PayloadRef }})}
+			}
+				{{- else }}
 			return &{{ .PayloadAccess }}{payload: info.RawPayload().({{ .PayloadRef }})}
+				{{- end }}
 			{{- end }}
 		default:
 			return nil
 		}
 		{{- else }}
+			{{- if hasEndpointStruct (index .Methods 0) }}
+	switch pay := info.RawPayload().(type) {
+	case *{{ (index .Methods 0).ServerStream.EndpointStruct }}:
+		return &{{ (index .Methods 0).PayloadAccess }}{payload: pay.Payload}
+	default:
+		return &{{ (index .Methods 0).PayloadAccess }}{payload: pay.({{ (index .Methods 0).PayloadRef }})}
+	}
+			{{- else }}
 	return &{{ (index .Methods 0).PayloadAccess }}{payload: info.RawPayload().({{ (index .Methods 0).PayloadRef }})}
+			{{- end }}
 		{{- end }}
 }
 	{{- end }}

codegen/service/templates/server_interceptor_wrappers.go.tpl

@@ -6,22 +6,8 @@
 func wrap{{ .MethodName }}{{ $interceptor.Name }}(endpoint goa.Endpoint, i ServerInterceptors) goa.Endpoint {
 	return func(ctx context.Context, req any) (any, error) {
 	{{- if or $interceptor.HasStreamingPayloadAccess $interceptor.HasStreamingResultAccess }}
-		{{- if $interceptor.HasPayloadAccess }}
-		info := &{{ $interceptor.Name }}Info{
-			service:    "{{ $.Service }}",
-			method:     "{{ .MethodName }}",
-			callType:   goa.InterceptorUnary,
-			rawPayload: req,
-		}
-		res, err := i.{{ $interceptor.Name }}(ctx, info, endpoint)
-		{{- else }}
-		res, err := endpoint(ctx, req)
-		{{- end }}
-		if err != nil {
-			return res, err
-		}
-		stream := res.({{ .ServerStream.Interface }})
-		return &wrapped{{ .ServerStream.Interface }}{
+		stream := req.(*{{ .ServerStream.EndpointStruct }}).Stream
+		req.(*{{ .ServerStream.EndpointStruct }}).Stream = &wrapped{{ .ServerStream.Interface }}{
 			ctx:     ctx,
 		{{- if $interceptor.HasStreamingResultAccess }}
 			sendWithContext: func(ctx context.Context, req {{ .ServerStream.SendTypeRef }}) error {
@@ -53,7 +39,18 @@ func wrap{{ .MethodName }}{{ $interceptor.Name }}(endpoint goa.Endpoint, i Serve
 			},
 		{{- end }}
 			stream: stream,
-		}, nil
+		}
+		{{- if $interceptor.HasPayloadAccess }}
+		info := &{{ $interceptor.Name }}Info{
+			service:    "{{ $.Service }}",
+			method:     "{{ .MethodName }}",
+			callType:   goa.InterceptorUnary,
+			rawPayload: req,
+		}
+		return i.{{ $interceptor.Name }}(ctx, info, endpoint)
+		{{- else }}
+		return endpoint(ctx, req)
+		{{- end }}
 	{{- else }}
 		info := &{{ $interceptor.Name }}Info{
 			service:    "{{ $.Service }}",

codegen/service/testdata/interceptors/streaming-interceptors-with-read-payload-and-read-streaming-payload_interceptor_wrappers.go.golden

@@ -19,18 +19,8 @@ type wrappedMethodClientStream struct {
 // wrapLoggingMethod applies the logging server interceptor to endpoints.
 func wrapMethodLogging(endpoint goa.Endpoint, i ServerInterceptors) goa.Endpoint {
 	return func(ctx context.Context, req any) (any, error) {
-		info := &LoggingInfo{
-			service:    "StreamingInterceptorsWithReadPayloadAndReadStreamingPayload",
-			method:     "Method",
-			callType:   goa.InterceptorUnary,
-			rawPayload: req,
-		}
-		res, err := i.Logging(ctx, info, endpoint)
-		if err != nil {
-			return res, err
-		}
-		stream := res.(MethodServerStream)
-		return &wrappedMethodServerStream{
+		stream := req.(*MethodEndpointInput).Stream
+		req.(*MethodEndpointInput).Stream = &wrappedMethodServerStream{
 			ctx: ctx,
 			recvWithContext: func(ctx context.Context) (*MethodStreamingPayload, error) {
 				info := &LoggingInfo{
@@ -45,7 +35,14 @@ func wrapMethodLogging(endpoint goa.Endpoint, i ServerInterceptors) goa.Endpoint
 				return castRes, err
 			},
 			stream: stream,
-		}, nil
+		}
+		info := &LoggingInfo{
+			service:    "StreamingInterceptorsWithReadPayloadAndReadStreamingPayload",
+			method:     "Method",
+			callType:   goa.InterceptorUnary,
+			rawPayload: req,
+		}
+		return i.Logging(ctx, info, endpoint)
 	}
 }
 

codegen/service/testdata/interceptors/streaming-interceptors-with-read-payload-and-read-streaming-payload_service_interceptors.go.golden

@@ -73,7 +73,12 @@ func (info *LoggingInfo) RawPayload() any {
 
 // Payload returns a type-safe accessor for the method payload.
 func (info *LoggingInfo) Payload() LoggingPayload {
-	return &loggingMethodPayload{payload: info.RawPayload().(*MethodPayload)}
+	switch pay := info.RawPayload().(type) {
+	case *MethodEndpointInput:
+		return &loggingMethodPayload{payload: pay.Payload}
+	default:
+		return &loggingMethodPayload{payload: pay.(*MethodPayload)}
+	}
 }
 
 // ClientStreamingPayload returns a type-safe accessor for the method streaming payload for a client-side interceptor.

codegen/service/testdata/interceptors/streaming-interceptors-with-read-payload_service_interceptors.go.golden

@@ -63,7 +63,12 @@ func (info *LoggingInfo) RawPayload() any {
 
 // Payload returns a type-safe accessor for the method payload.
 func (info *LoggingInfo) Payload() LoggingPayload {
-	return &loggingMethodPayload{payload: info.RawPayload().(*MethodPayload)}
+	switch pay := info.RawPayload().(type) {
+	case *MethodEndpointInput:
+		return &loggingMethodPayload{payload: pay.Payload}
+	default:
+		return &loggingMethodPayload{payload: pay.(*MethodPayload)}
+	}
 }
 
 // Private implementation methods

codegen/service/testdata/interceptors/streaming-interceptors-with-read-streaming-result_interceptor_wrappers.go.golden

@@ -19,12 +19,8 @@ type wrappedMethodClientStream struct {
 // wrapLoggingMethod applies the logging server interceptor to endpoints.
 func wrapMethodLogging(endpoint goa.Endpoint, i ServerInterceptors) goa.Endpoint {
 	return func(ctx context.Context, req any) (any, error) {
-		res, err := endpoint(ctx, req)
-		if err != nil {
-			return res, err
-		}
-		stream := res.(MethodServerStream)
-		return &wrappedMethodServerStream{
+		stream := req.(*MethodEndpointInput).Stream
+		req.(*MethodEndpointInput).Stream = &wrappedMethodServerStream{
 			ctx: ctx,
 			sendWithContext: func(ctx context.Context, req *MethodResult) error {
 				info := &LoggingInfo{
@@ -40,7 +36,8 @@ func wrapMethodLogging(endpoint goa.Endpoint, i ServerInterceptors) goa.Endpoint
 				return err
 			},
 			stream: stream,
-		}, nil
+		}
+		return endpoint(ctx, req)
 	}
 }
 

codegen/service/testdata/interceptors/streaming-interceptors_interceptor_wrappers.go.golden

@@ -21,12 +21,8 @@ type wrappedMethodClientStream struct {
 // wrapLoggingMethod applies the logging server interceptor to endpoints.
 func wrapMethodLogging(endpoint goa.Endpoint, i ServerInterceptors) goa.Endpoint {
 	return func(ctx context.Context, req any) (any, error) {
-		res, err := endpoint(ctx, req)
-		if err != nil {
-			return res, err
-		}
-		stream := res.(MethodServerStream)
-		return &wrappedMethodServerStream{
+		stream := req.(*MethodEndpointInput).Stream
+		req.(*MethodEndpointInput).Stream = &wrappedMethodServerStream{
 			ctx: ctx,
 			sendWithContext: func(ctx context.Context, req *MethodResult) error {
 				info := &LoggingInfo{
@@ -54,7 +50,8 @@ func wrapMethodLogging(endpoint goa.Endpoint, i ServerInterceptors) goa.Endpoint
 				return castRes, err
 			},
 			stream: stream,
-		}, nil
+		}
+		return endpoint(ctx, req)
 	}
 }