Add first version of request headers

Author: standy66

src/purerpc/client.py

@@ -25,14 +25,17 @@ def __init__(self, service_name: str, channel: Channel):
         self.service_name = service_name
         self.channel = channel
 
-    async def rpc(self, method_name: str, request_type, response_type):
+    async def rpc(self, method_name: str, request_type, response_type, metadata=None):
         if self.channel.grpc_socket is None:
             await self.channel.connect()
         message_type = request_type.DESCRIPTOR.full_name
+        if metadata is None:
+            metadata = ()
         stream = await self.channel.grpc_socket.start_request("http", self.service_name,
                                                               method_name, message_type,
                                                               "{}:{}".format(self.channel.host,
-                                                                             self.channel.port))
+                                                                             self.channel.port),
+                                                              custom_metadata=metadata)
         stream.expect_message_type(response_type)
         return stream
 

src/purerpc/grpclib/connection.py

@@ -8,6 +8,7 @@
 import h2.exceptions
 from h2.settings import SettingCodes
 
+from .headers import HeaderDict, sanitize_headers
 from .status import Status
 from .config import GRPCConfiguration
 from .events import MessageReceived, RequestReceived, RequestEnded, ResponseReceived, ResponseEnded
@@ -75,13 +76,13 @@ def _request_received(self, event: h2.events.RequestReceived):
         if event.stream_ended:
             raise ProtocolError("Stream ended before data was sent")
         request = RequestReceived.parse_from_stream_id_and_headers_destructive(
-            event.stream_id, dict(event.headers))
+            event.stream_id, HeaderDict(event.headers))
         self.message_read_buffers[event.stream_id] = MessageReadBuffer(request.message_encoding,
             self.config.max_message_length)
         return [request]
 
     def _response_received(self, event: h2.events.ResponseReceived):
-        headers = dict(event.headers)
+        headers = HeaderDict(event.headers)
         response_received = ResponseReceived.parse_from_stream_id_and_headers_destructive(
             event.stream_id, headers)
         if event.stream_ended:
@@ -97,7 +98,7 @@ def _response_received(self, event: h2.events.ResponseReceived):
 
     def _trailers_received(self, event: h2.events.TrailersReceived):
         response_ended = ResponseEnded.parse_from_stream_id_and_headers_destructive(
-            event.stream_id, dict(event.headers))
+            event.stream_id, HeaderDict(event.headers))
         return [response_ended]
 
     def _informational_response_received(self, event: h2.events.InformationalResponseReceived):
@@ -228,7 +229,6 @@ def start_request(self, stream_id: int, scheme: str, service_name: str, method_n
             (":path", "/{}/{}".format(service_name, method_name)),
             ("te", "trailers"),
             ("content-type", "application/grpc" + content_type_suffix),
-            *custom_metadata
         ]
         if authority is not None:
             headers.insert(3, (":authority", authority))
@@ -250,6 +250,7 @@ def start_request(self, stream_id: int, scheme: str, service_name: str, method_n
             headers.append(("grpc-accept-encoding", self.config._message_accept_encoding))
         if self.config._user_agent is not None:
             headers.append(("user-agent", self.config._user_agent))
+        headers.extend(sanitize_headers(custom_metadata))
         self.h2_connection.send_headers(stream_id, headers, end_stream=False)
 
     def end_request(self, stream_id: int):
@@ -259,12 +260,12 @@ def start_response(self, stream_id: int, content_type_suffix="", custom_metadata
         headers = [
             (":status", "200"),
             ("content-type", "application/grpc" + content_type_suffix),
-            *custom_metadata,
         ]
         if self.config._message_encoding is not None:
             headers.append(("grpc-encoding", self.config._message_encoding))
         if self.config._message_accept_encoding is not None:
             headers.append(("grpc-accept-encoding", self.config._message_accept_encoding))
+        headers.extend(sanitize_headers(custom_metadata))
         self.h2_connection.send_headers(stream_id, headers, end_stream=False)
 
     def respond_status(self, stream_id: int, status: Status, content_type_suffix="",
@@ -273,19 +274,19 @@ def respond_status(self, stream_id: int, status: Status, content_type_suffix="",
             (":status", "200"),
             ("content-type", "application/grpc" + content_type_suffix),
             ("grpc-status", str(status.int_value)),
-            *custom_metadata,
         ]
         if status.status_message:
             # TODO: should be percent encoded
             trailers.append(("grpc-message", status.status_message))
+        trailers.extend(sanitize_headers(custom_metadata))
         self.h2_connection.send_headers(stream_id, trailers, end_stream=True)
 
     def end_response(self, stream_id: int, status: Status, custom_metadata=()):
         trailers = [
             ("grpc-status", str(status.int_value)),
-            *custom_metadata,
         ]
         if status.status_message:
             # TODO: should be percent encoded
             trailers.append(("grpc-message", status.status_message))
+        trailers.extend(sanitize_headers(custom_metadata))
         self.h2_connection.send_headers(stream_id, trailers, end_stream=True)

src/purerpc/grpclib/events.py

@@ -1,5 +1,6 @@
-import base64
 import datetime
+
+from .headers import HeaderDict
 from .exceptions import ProtocolError
 from .status import Status
 
@@ -22,10 +23,10 @@ def __init__(self, stream_id: int, scheme: str, service_name: str, method_name:
         self.message_encoding = None
         self.message_accept_encoding = None
         self.user_agent = None
-        self.custom_metadata = {}
+        self.custom_metadata = ()
 
     @staticmethod
-    def parse_from_stream_id_and_headers_destructive(stream_id: int, headers: dict):
+    def parse_from_stream_id_and_headers_destructive(stream_id: int, headers: HeaderDict):
         if headers.pop(":method") != "POST":
             raise ProtocolError("Unsupported method {}".format(headers[":method"]))
 
@@ -83,11 +84,8 @@ def parse_from_stream_id_and_headers_destructive(stream_id: int, headers: dict):
         if "grpc-message-type" in headers:
             event.message_type = headers.pop("grpc-message-type")
 
-        for header_name in list(headers.keys()):
-            if header_name.endswith("-bin"):
-                event.custom_metadata[header_name] = base64.b64decode(headers.pop(header_name))
-            else:
-                event.custom_metadata[header_name] = headers.pop(header_name)
+        event.custom_metadata = tuple(header for header_name in list(headers.keys())
+                                      for header in headers.extract_headers(header_name))
         return event
 
 
@@ -109,10 +107,10 @@ def __init__(self, stream_id: int, content_type: str):
         self.content_type = content_type
         self.message_encoding = None
         self.message_accept_encoding = None
-        self.custom_metadata = {}
+        self.custom_metadata = ()
 
     @staticmethod
-    def parse_from_stream_id_and_headers_destructive(stream_id: int, headers: dict):
+    def parse_from_stream_id_and_headers_destructive(stream_id: int, headers: HeaderDict):
         if int(headers.pop(":status")) != 200:
             raise ProtocolError("http status is not 200")
 
@@ -128,25 +126,19 @@ def parse_from_stream_id_and_headers_destructive(stream_id: int, headers: dict):
         if "grpc-accept-encoding" in headers:
             event.message_accept_encoding = headers.pop("grpc-accept-encoding").split(",")
 
-        for header_name in list(headers.keys()):
-            if header_name in ["grpc-status", "grpc-message"]:
-                # is not metadata, will be parsed in ResponseEnded
-                continue
-            if header_name.endswith("-bin"):
-                event.custom_metadata[header_name] = base64.b64decode(headers.pop(header_name))
-            else:
-                event.custom_metadata[header_name] = headers.pop(header_name)
+        event.custom_metadata = tuple(header for header_name in list(headers.keys())
+                                      for header in headers.extract_headers(header_name))
         return event
 
 
 class ResponseEnded(Event):
     def __init__(self, stream_id: int, status: Status):
         self.stream_id = stream_id
         self.status = status
-        self.custom_metadata = {}
+        self.custom_metadata = ()
 
     @staticmethod
-    def parse_from_stream_id_and_headers_destructive(stream_id: int, headers: dict):
+    def parse_from_stream_id_and_headers_destructive(stream_id: int, headers: HeaderDict):
         if "grpc-status" not in headers:
             raise ProtocolError("Expected grpc-status in trailers")
 
@@ -159,9 +151,6 @@ def parse_from_stream_id_and_headers_destructive(stream_id: int, headers: dict):
 
         event = ResponseEnded(stream_id, Status(status_code, status_message))
 
-        for header_name in list(headers.keys()):
-            if header_name.endswith("-bin"):
-                event.custom_metadata[header_name] = base64.b64decode(headers.pop(header_name))
-            else:
-                event.custom_metadata[header_name] = headers.pop(header_name)
+        event.custom_metadata = tuple(header for header_name in list(headers.keys())
+                                      for header in headers.extract_headers(header_name))
         return event

src/purerpc/grpclib/headers.py

@@ -0,0 +1,59 @@
+import base64
+
+
+class HeaderDict(dict):
+    def __init__(self, values):
+        super().__init__()
+        for key, value in values:
+            if key not in self:
+                self[key] = [value]
+            else:
+                self[key].append(value)
+        for key in self:
+            if len(self[key]) == 1:
+                self[key] = self[key][0]
+
+    def extract_headers(self, header_name: str):
+        """Returns all headers with name == header_name as list of tuples (name, value)"""
+        if header_name.startswith("grpc-"):
+            return ()
+        else:
+            value = self.pop(header_name)
+            is_binary = header_name.endswith("-bin")
+
+            if not isinstance(value, list):
+                value_list = [value]
+            else:
+                value_list = value
+
+            if is_binary:
+                return ((header_name, b64decode(value)) for value_sublist in value_list for value
+                        in value_sublist.split(","))
+            else:
+                return ((header_name, value) for value in value_list)
+
+
+def sanitize_headers(headers):
+    for name, value in headers:
+        if isinstance(value, bytes) and not name.endswith("-bin"):
+            raise ValueError(f"Got binary value for header name '{name}', but name does not end "
+                             f"with '-bin' suffix")
+        if name.startswith("grpc-"):
+            raise ValueError(f"Got header with name '{name}', but custom metadata headers should "
+                             f"not start with 'grpc-' prefix")
+        if name.endswith("-bin"):
+            yield name, b64encode(value)
+        else:
+            yield name, value
+
+
+def b64decode(data: str) -> bytes:
+    # Apply missing padding
+    missing_padding = len(data) % 4
+    if missing_padding:
+        data += "=" * (4 - missing_padding)
+    return base64.b64decode(data)
+
+
+def b64encode(data: bytes) -> str:
+    return base64.b64encode(data).rstrip(b"=").decode("utf-8")

src/purerpc/server.py

@@ -2,13 +2,15 @@
 import inspect
 import warnings
 import collections
+import functools
 from multiprocessing import Process
 
 import curio
 import curio.meta
 import typing
 import logging
 
+from .grpclib.events import RequestReceived
 from .grpclib.status import Status, StatusCode
 from .grpclib.exceptions import RpcFailedError
 from purerpc.grpc_proto import GRPCProtoStream, GRPCProtoSocket
@@ -29,24 +31,37 @@ def __init__(self, name):
         self.name = name
         self.methods = {}
 
-    def add_method(self, method_name: str, method_fn, rpc_signature: RPCSignature):
-        self.methods[method_name] = BoundRPCMethod(method_fn, rpc_signature)
+    def add_method(self, method_name: str, method_fn, rpc_signature: RPCSignature,
+                   method_signature: inspect.Signature = None):
+        if method_signature is None:
+            method_signature = inspect.signature(method_fn)
+        if len(method_signature.parameters) == 1:
+            def method_fn_with_headers(arg, request):
+                return method_fn(arg)
+        elif len(method_signature.parameters) == 2:
+            if list(method_signature.parameters.values())[1].name == "request":
+                method_fn_with_headers = method_fn
+            else:
+                raise ValueError("Expected second parameter 'request'")
+        else:
+            raise ValueError("Expected method_fn to have exactly one or two parameters")
+        self.methods[method_name] = BoundRPCMethod(method_fn_with_headers, rpc_signature)
 
     def rpc(self, method_name):
         def decorator(func):
             signature = inspect.signature(func)
             if signature.return_annotation == signature.empty:
                 raise ValueError("Only annotated methods can be used with Service.rpc() decorator")
-            if len(signature.parameters) != 1:
-                raise ValueError("Only functions with one parameter can be used with Service.rpc("
-                                 ") decorator")
+            if len(signature.parameters) not in (1, 2):
+                raise ValueError("Only functions with one or two parameters can be used with "
+                                 "Service.rpc() decorator")
             parameter = next(iter(signature.parameters.values()))
             if parameter.annotation == parameter.empty:
                 raise ValueError("Only annotated methods can be used with Service.rpc() decorator")
 
             rpc_signature = RPCSignature.from_annotations(parameter.annotation,
                                                           signature.return_annotation)
-            self.add_method(method_name, func, rpc_signature)
+            self.add_method(method_name, func, rpc_signature, method_signature=signature)
             return func
         return decorator
 
@@ -104,6 +119,10 @@ async def request_received(self, stream: GRPCProtoStream):
         await stream.start_response()
         event = await stream.receive_event()
 
+        if not isinstance(event, RequestReceived):
+            await stream.close(Status(StatusCode.INTERNAL, status_message="Expected headers"))
+            return
+
         try:
             service = self.server.services[event.service_name]
         except KeyError:
@@ -125,7 +144,7 @@ async def request_received(self, stream: GRPCProtoStream):
 
         # TODO: Should at least pass through GeneratorExit
         try:
-            method_fn = bound_rpc_method.method_fn
+            method_fn = functools.partial(bound_rpc_method.method_fn, request=event)
             cardinality = bound_rpc_method.signature.cardinality
             stream.expect_message_type(bound_rpc_method.signature.request_type)
             if cardinality == Cardinality.STREAM_STREAM:

src/purerpc/wrappers.py

@@ -80,40 +80,40 @@ def __init__(self, stream_fn):
 
 
 class ClientStubUnaryUnary(ClientStub):
-    async def __call__(self, message):
-        stream = await self._stream_fn()
+    async def __call__(self, message, *, metadata=None):
+        stream = await self._stream_fn(metadata=metadata)
         await send_single_message_client(stream, message)
         return await extract_message_from_singleton_stream(stream)
 
 
 class ClientStubUnaryStream(ClientStub):
-    async def __call__(self, message):
-        stream = await self._stream_fn()
+    async def __call__(self, message, *, metadata=None):
+        stream = await self._stream_fn(metadata=metadata)
         await send_single_message_client(stream, message)
         async for message in stream_to_async_iterator(stream):
             yield message
 
 
 class ClientStubStreamUnary(ClientStub):
-    async def __call__(self, message_aiter):
-        stream = await self._stream_fn()
+    async def __call__(self, message_aiter, *, metadata=None):
+        stream = await self._stream_fn(metadata=metadata)
         await curio.spawn(send_multiple_messages_client, stream, message_aiter, daemon=True)
         return await extract_message_from_singleton_stream(stream)
 
 
 class ClientStubStreamStream(ClientStub):
-    async def call_aiter(self, message_aiter):
-        stream = await self._stream_fn()
+    async def call_aiter(self, message_aiter, metadata):
+        stream = await self._stream_fn(metadata=metadata)
         if message_aiter is not None:
             await curio.spawn(send_multiple_messages_client, stream, message_aiter, daemon=True)
             async for message in stream_to_async_iterator(stream):
                 yield message
 
-    async def call_stream(self):
-        return await self._stream_fn()
+    async def call_stream(self, metadata):
+        return await self._stream_fn(metadata=metadata)
 
-    def __call__(self, message_aiter=None):
+    def __call__(self, message_aiter=None, *, metadata=None):
         if message_aiter is None:
-            return self.call_stream()
+            return self.call_stream(metadata)
         else:
-            return self.call_aiter(message_aiter)
+            return self.call_aiter(message_aiter, metadata)