From 0971bef5b5a3df6f778ac37d50dab317fb60cdb4 Mon Sep 17 00:00:00 2001
From: Alex Stapleton <alexs@prol.etari.at>
Date: Mon, 27 Dec 2021 10:49:00 +0000
Subject: [PATCH] Always send grpc-status in trailers

---
 sonora/asgi.py | 24 +++++++++++++-----------
 sonora/wsgi.py | 34 +++++++++++++++++-----------------
 2 files changed, 30 insertions(+), 28 deletions(-)

diff --git a/sonora/asgi.py b/sonora/asgi.py
index 46a91ed..057fbd2 100644
--- a/sonora/asgi.py
+++ b/sonora/asgi.py
@@ -177,6 +177,7 @@ async def _do_streaming_response(
                 recv_task.cancel()
 
         trailers = [("grpc-status", str(context.code.value[0]))]
+
         if context.details:
             trailers.append(("grpc-message", quote(context.details)))
 
@@ -199,12 +200,6 @@ async def _do_unary_response(
 
         status = 200
 
-        headers.append((b"grpc-status", str(context.code.value[0]).encode()))
-        if context.details:
-            headers.append(
-                (b"grpc-message", quote(context.details.encode("utf8")).encode("ascii"))
-            )
-
         if context._initial_metadata:
             headers.extend(context._initial_metadata)
 
@@ -215,13 +210,18 @@ async def _do_unary_response(
         else:
             message_data = b""
 
+        trailers = [(b"grpc-status", str(context.code.value[0]).encode())]
+
+        if context.details:
+            trailers.append(
+                (b"grpc-message", quote(context.details.encode("utf8")).encode("ascii"))
+            )
+
         if context._trailing_metadata:
-            trailers = context._trailing_metadata
+            trailers.extend(context._trailing_metadata)
 
-            trailer_message = protocol.pack_trailers(trailers)
-            trailer_data = wrap_message(True, False, trailer_message)
-        else:
-            trailer_data = b""
+        trailer_message = protocol.pack_trailers(trailers)
+        trailer_data = wrap_message(True, False, trailer_message)
 
         content_length = len(message_data) + len(trailer_data)
 
@@ -230,9 +230,11 @@ async def _do_unary_response(
         await send(
             {"type": "http.response.start", "status": status, "headers": headers}
         )
+
         await send(
             {"type": "http.response.body", "body": message_data, "more_body": True}
         )
+
         await send(
             {"type": "http.response.body", "body": trailer_data, "more_body": False}
         )
diff --git a/sonora/wsgi.py b/sonora/wsgi.py
index 099da3d..b54d219 100644
--- a/sonora/wsgi.py
+++ b/sonora/wsgi.py
@@ -138,7 +138,7 @@ def _do_streaming_response(
             pass
 
         if context._initial_metadata:
-            headers.extend(protocol.encode_headers(context._initial_metadata))
+            headers.extend(context._initial_metadata)
 
         start_response("200 OK", headers)
 
@@ -158,7 +158,7 @@ def _do_streaming_response(
             trailers.append(("grpc-message", quote(context.details.encode("utf8"))))
 
         if context._trailing_metadata:
-            trailers.extend(protocol.encode_headers(context._trailing_metadata))
+            trailers.extend(context._trailing_metadata)
 
         trailer_message = protocol.pack_trailers(trailers)
 
@@ -174,28 +174,28 @@ def _do_unary_response(
         else:
             message_data = b""
 
-        if context._trailing_metadata:
-            trailers = protocol.encode_headers(context._trailing_metadata)
-            trailer_message = protocol.pack_trailers(trailers)
-            trailer_data = wrap_message(True, False, trailer_message)
-        else:
-            trailer_data = b""
+        if context._initial_metadata:
+            headers.extend(context._initial_metadata)
 
-        content_length = len(message_data) + len(trailer_data)
+        trailers = [("grpc-status", str(context.code.value[0]))]
 
-        headers.append(("content-length", str(content_length)))
+        if context.details:
+            trailers.append(("grpc-message", quote(context.details.encode("utf8"))))
 
-        headers.append(("grpc-status", str(context.code.value[0])))
+        if context._trailing_metadata:
+            trailers.extend(context._trailing_metadata)
 
-        if context.details:
-            headers.append(("grpc-message", quote(context.details.encode("utf8"))))
+        trailer_message = protocol.pack_trailers(trailers)
+        trailer_data = wrap_message(True, False, trailer_message)
 
-        if context._initial_metadata:
-            headers.extend(protocol.encode_headers(context._initial_metadata))
+        content_length = len(message_data) + len(trailer_data)
+
+        headers.append(("content-length", str(content_length)))
 
         start_response("200 OK", headers)
 
         yield message_data
+
         yield trailer_data
 
     def _do_cors_preflight(self, environ, start_response):
@@ -342,10 +342,10 @@ def invocation_metadata(self):
         return self._invocation_metadata
 
     def send_initial_metadata(self, initial_metadata):
-        self._initial_metadata = initial_metadata
+        self._initial_metadata = protocol.encode_headers(initial_metadata)
 
     def set_trailing_metadata(self, trailing_metadata):
-        self._trailing_metadata = trailing_metadata
+        self._trailing_metadata = protocol.encode_headers(trailing_metadata)
 
     def peer(self):
         raise NotImplementedError()