Add interceptors for stream requests

This adds authentication and signing interceptors for stream requests.
Specifically for those where the client sends a single message and
receives a stream.

Signed-off-by: Florian Wagner <[email protected]>

Author: florian-wagner-frequenz

src/frequenz/client/base/authentication.py

@@ -4,13 +4,15 @@
 """An Interceptor that adds the API key to a gRPC call."""
 
 import dataclasses
-from typing import Callable
+from typing import AsyncIterable, Callable
 
 from grpc.aio import (
     ClientCallDetails,
     Metadata,
     UnaryUnaryCall,
     UnaryUnaryClientInterceptor,
+    UnaryStreamCall,
+    UnaryStreamClientInterceptor,
 )
 
 
@@ -22,8 +24,8 @@ class AuthenticationOptions:
     """The API key to authenticate with."""
 
 
-# There is an issue in gRPC that causes the type to be unspecifieable correctly here.
-class AuthenticationInterceptor(UnaryUnaryClientInterceptor):  # type: ignore[type-arg]
+# There is an issue in gRPC which means the type can not be specified correctly here.
+class AuthenticationInterceptor(UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor):  # type: ignore[type-arg]
     """An Interceptor that adds HMAC authentication of the metadata fields to a gRPC call."""
 
     def __init__(self, options: AuthenticationOptions):
@@ -54,11 +56,20 @@ async def intercept_unary_unary(
         Returns:
             The response object (this implementation does not modify the response).
         """
-        self.add_auth_header(
-            client_call_details,
-        )
+        self.add_auth_header(client_call_details)
         return await continuation(client_call_details, request)
 
+    async def intercept_unary_stream(
+        self,
+        continuation: Callable[[ClientCallDetails, object], UnaryStreamCall],
+        client_call_details: ClientCallDetails,
+        request: object,
+    ) -> AsyncIterable | UnaryStreamCall:
+
+        self.add_auth_header(client_call_details)
+
+        return continuation(client_call_details, request)
+
     def add_auth_header(
         self,
         client_call_details: ClientCallDetails,

src/frequenz/client/base/channel.py

@@ -17,7 +17,7 @@
     secure_channel,
 )
 
-from .authentication import AuthenticationInterceptor, AuthenticationOptions
+from .authentication import AuthenticationInterceptorUnaryUnary, AuthenticationOptions
 from .signing import SigningInterceptor, SigningOptions
 
 
@@ -194,7 +194,7 @@ def parse_grpc_uri(
     interceptors: list[ClientInterceptor] = []
     if defaults.auth is not None:
         interceptors.append(
-            AuthenticationInterceptor(options=defaults.auth)  # type: ignore[arg-type]
+            AuthenticationInterceptorUnaryUnary(options=defaults.auth)  # type: ignore[arg-type]
         )
 
     if defaults.sign is not None:

src/frequenz/client/base/signing.py

@@ -9,12 +9,14 @@
 import secrets
 import time
 from base64 import urlsafe_b64encode
-from typing import Any, Callable
+from typing import Any, AsyncIterable, Callable
 
 from grpc.aio import (
     ClientCallDetails,
     UnaryUnaryCall,
     UnaryUnaryClientInterceptor,
+    UnaryStreamCall,
+    UnaryStreamClientInterceptor,
 )
 
 _logger = logging.getLogger(__name__)
@@ -28,8 +30,8 @@ class SigningOptions:
     """The secret to sign the message with."""
 
 
-# There is an issue in gRPC that causes the type to be unspecifieable correctly here.
-class SigningInterceptor(UnaryUnaryClientInterceptor):  # type: ignore[type-arg]
+# There is an issue in gRPC which means the type can not be specified correctly here.
+class SigningInterceptor(UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor):  # type: ignore[type-arg]
     """An Interceptor that adds HMAC authentication of the metadata fields to a gRPC call."""
 
     def __init__(self, options: SigningOptions):
@@ -67,6 +69,21 @@ async def intercept_unary_unary(
         )
         return await continuation(client_call_details, request)
 
+
+    async def intercept_unary_stream(
+        self,
+        continuation: Callable[[ClientCallDetails, Any], UnaryStreamCall],
+        client_call_details: ClientCallDetails,
+        request: Any,
+    ) -> AsyncIterable | UnaryStreamCall:
+        self.add_hmac(
+            client_call_details,
+            int(time.time()).to_bytes(8, "big"),
+            secrets.token_bytes(16),
+        )
+        return continuation(client_call_details, request)
+
+
     def add_hmac(
         self, client_call_details: ClientCallDetails, ts: bytes, nonce: bytes
     ) -> None: