Add delimiter after response stream error message (grpc-ecosystem#2591)

stelcodes · Mar 20, 2022 · 768087f · 768087f
1 parent 9c4c61e
commit 768087f
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/runtime/handler.go b/runtime/handler.go
@@ -52,11 +52,11 @@ func ForwardResponseStream(ctx context.Context, mux *ServeMux, marshaler Marshal
 			return
 		}
 		if err != nil {
-			handleForwardResponseStreamError(ctx, wroteHeader, marshaler, w, req, mux, err)
+			handleForwardResponseStreamError(ctx, wroteHeader, marshaler, w, req, mux, err, delimiter)
 			return
 		}
 		if err := handleForwardResponseOptions(ctx, w, resp, opts); err != nil {
-			handleForwardResponseStreamError(ctx, wroteHeader, marshaler, w, req, mux, err)
+			handleForwardResponseStreamError(ctx, wroteHeader, marshaler, w, req, mux, err, delimiter)
 			return
 		}
 
@@ -82,7 +82,7 @@ func ForwardResponseStream(ctx context.Context, mux *ServeMux, marshaler Marshal
 
 		if err != nil {
 			grpclog.Infof("Failed to marshal response chunk: %v", err)
-			handleForwardResponseStreamError(ctx, wroteHeader, marshaler, w, req, mux, err)
+			handleForwardResponseStreamError(ctx, wroteHeader, marshaler, w, req, mux, err, delimiter)
 			return
 		}
 		if _, err = w.Write(buf); err != nil {
@@ -200,7 +200,7 @@ func handleForwardResponseOptions(ctx context.Context, w http.ResponseWriter, re
 	return nil
 }
 
-func handleForwardResponseStreamError(ctx context.Context, wroteHeader bool, marshaler Marshaler, w http.ResponseWriter, req *http.Request, mux *ServeMux, err error) {
+func handleForwardResponseStreamError(ctx context.Context, wroteHeader bool, marshaler Marshaler, w http.ResponseWriter, req *http.Request, mux *ServeMux, err error, delimiter []byte) {
 	st := mux.streamErrorHandler(ctx, err)
 	msg := errorChunk(st)
 	if !wroteHeader {
@@ -216,6 +216,10 @@ func handleForwardResponseStreamError(ctx context.Context, wroteHeader bool, mar
 		grpclog.Infof("Failed to notify error to client: %v", werr)
 		return
 	}
+	if _, derr := w.Write(delimiter); derr != nil {
+		grpclog.Infof("Failed to send delimiter chunk: %v", err)
+		return
+	}
 }
 
 func errorChunk(st *status.Status) map[string]proto.Message {

diff --git a/runtime/handler_test.go b/runtime/handler_test.go
@@ -120,6 +120,7 @@ func TestForwardResponseStream(t *testing.T) {
 						// Skip non-stream errors
 						t.Skip("checking error encodings")
 					}
+					delimiter := marshaler.Delimiter()
 					st := status.Convert(msg.err)
 					b, err := marshaler.Marshal(map[string]proto.Message{
 						"error": st.Proto(),
@@ -128,7 +129,7 @@ func TestForwardResponseStream(t *testing.T) {
 						t.Errorf("marshaler.Marshal() failed %v", err)
 					}
 					errBytes := body[len(want):]
-					if string(errBytes) != string(b) {
+					if string(errBytes) != string(b)+string(delimiter) {
 						t.Errorf("ForwardResponseStream() = \"%s\" want \"%s\"", errBytes, b)
 					}