Adds support for secure channels (#18)

Signed-off-by: Elena Kolevska <elena@kolevska.com>

Author: elena-kolevska

.gitignore

@@ -128,4 +128,7 @@ dmypy.json
 # Pyre type checker
 .pyre/
 
+# IDEs
+.idea
+
 coverage.lcov

durabletask/client.py

@@ -95,8 +95,9 @@ def __init__(self, *,
                  host_address: Union[str, None] = None,
                  metadata: Union[List[Tuple[str, str]], None] = None,
                  log_handler = None,
-                 log_formatter: Union[logging.Formatter, None] = None):
-        channel = shared.get_grpc_channel(host_address, metadata)
+                 log_formatter: Union[logging.Formatter, None] = None,
+                 secure_channel: bool = False):
+        channel = shared.get_grpc_channel(host_address, metadata, secure_channel=secure_channel)
         self._stub = stubs.TaskHubSidecarServiceStub(channel)
         self._logger = shared.get_logger("client", log_handler, log_formatter)
 

durabletask/internal/shared.py

@@ -20,10 +20,15 @@ def get_default_host_address() -> str:
     return "localhost:4001"
 
 
-def get_grpc_channel(host_address: Union[str, None], metadata: Union[List[Tuple[str, str]], None]) -> grpc.Channel:
+def get_grpc_channel(host_address: Union[str, None], metadata: Union[List[Tuple[str, str]], None], secure_channel: bool = False) -> grpc.Channel:
     if host_address is None:
         host_address = get_default_host_address()
-    channel = grpc.insecure_channel(host_address)
+
+    if secure_channel:
+        channel = grpc.secure_channel(host_address, grpc.ssl_channel_credentials())
+    else:
+        channel = grpc.insecure_channel(host_address)
+
     if metadata is not None and len(metadata) > 0:
         interceptors = [DefaultClientInterceptorImpl(metadata)]
         channel = grpc.intercept_channel(channel, *interceptors)

durabletask/worker.py

@@ -87,14 +87,16 @@ def __init__(self, *,
                  host_address: Union[str, None] = None,
                  metadata: Union[List[Tuple[str, str]], None] = None,
                  log_handler = None,
-                 log_formatter: Union[logging.Formatter, None] = None):
+                 log_formatter: Union[logging.Formatter, None] = None,
+                 secure_channel: bool = False):
         self._registry = _Registry()
         self._host_address = host_address if host_address else shared.get_default_host_address()
         self._metadata = metadata
         self._logger = shared.get_logger("worker", log_handler, log_formatter)
         self._shutdown = Event()
         self._response_stream = None
         self._is_running = False
+        self._secure_channel = secure_channel
 
     def __enter__(self):
         return self
@@ -116,7 +118,7 @@ def add_activity(self, fn: task.Activity) -> str:
 
     def start(self):
         """Starts the worker on a background thread and begins listening for work items."""
-        channel = shared.get_grpc_channel(self._host_address, self._metadata)
+        channel = shared.get_grpc_channel(self._host_address, self._metadata, self._secure_channel)
         stub = stubs.TaskHubSidecarServiceStub(channel)
 
         if self._is_running:

tests/test_client.py

@@ -0,0 +1,41 @@
+from unittest.mock import patch
+
+from durabletask.internal.shared import (DefaultClientInterceptorImpl,
+                                         get_default_host_address,
+                                         get_grpc_channel)
+
+HOST_ADDRESS = 'localhost:50051'
+METADATA = [('key1', 'value1'), ('key2', 'value2')]
+
+
+def test_get_grpc_channel_insecure():
+    with patch('grpc.insecure_channel') as mock_channel:
+        get_grpc_channel(HOST_ADDRESS, METADATA, False)
+        mock_channel.assert_called_once_with(HOST_ADDRESS)
+
+
+def test_get_grpc_channel_secure():
+    with patch('grpc.secure_channel') as mock_channel, patch(
+            'grpc.ssl_channel_credentials') as mock_credentials:
+        get_grpc_channel(HOST_ADDRESS, METADATA, True)
+        mock_channel.assert_called_once_with(HOST_ADDRESS, mock_credentials.return_value)
+
+
+def test_get_grpc_channel_default_host_address():
+    with patch('grpc.insecure_channel') as mock_channel:
+        get_grpc_channel(None, METADATA, False)
+        mock_channel.assert_called_once_with(get_default_host_address())
+
+
+def test_get_grpc_channel_with_metadata():
+    with patch('grpc.insecure_channel') as mock_channel, patch(
+            'grpc.intercept_channel') as mock_intercept_channel:
+        get_grpc_channel(HOST_ADDRESS, METADATA, False)
+        mock_channel.assert_called_once_with(HOST_ADDRESS)
+        mock_intercept_channel.assert_called_once()
+
+        # Capture and check the arguments passed to intercept_channel()
+        args, kwargs = mock_intercept_channel.call_args
+        assert args[0] == mock_channel.return_value
+        assert isinstance(args[1], DefaultClientInterceptorImpl)
+        assert args[1]._metadata == METADATA