From 1e2ddeca0cf63fe6d1e702afc55a1325f43c3936 Mon Sep 17 00:00:00 2001 From: Matteo Merli Date: Thu, 16 Apr 2026 17:24:30 -0700 Subject: [PATCH] Add TLS support Adds a tls parameter to Client that wraps the gRPC channel with ssl_channel_credentials() for encrypted transport. Two ways to enable: - Pass tls=True to Client() - Prefix the service address with "tls://" (matching the Java SDK convention) The system trust store is used for server certificate verification. Custom CA certs can be added in a follow-up if needed. Signed-off-by: Matteo Merli --- src/oxia/client.py | 13 +++- src/oxia/internal/connection_pool.py | 9 ++- tests/tls_test.py | 92 ++++++++++++++++++++++++++++ 3 files changed, 110 insertions(+), 4 deletions(-) create mode 100644 tests/tls_test.py diff --git a/src/oxia/client.py b/src/oxia/client.py index f4f0d4c..80325fa 100644 --- a/src/oxia/client.py +++ b/src/oxia/client.py @@ -151,18 +151,27 @@ def __init__(self, service_address: str, namespace: str = "default", session_timeout_ms: int = 30_000, client_identifier: str = None, + tls: bool = False, ): """Create a new Oxia client. - @param service_address: Oxia service address (``host:port``). + @param service_address: Oxia service address (``host:port`` or + ``tls://host:port``). Prefixing with ``tls://`` is + equivalent to passing ``tls=True``. @param namespace: Oxia namespace. Default is ``"default"``. @param session_timeout_ms: Session timeout in milliseconds for ephemeral records. Default is 30 000 ms. @param client_identifier: Optional client identity string. If ``None``, a random UUID is generated. + @param tls: If ``True``, use a TLS-encrypted gRPC channel. The + system trust store is used to verify server certificates. + Default is ``False`` (insecure channel). """ self._closed = False - self._connections = ConnectionPool() + if service_address.startswith("tls://"): + service_address = service_address[len("tls://"):] + tls = True + self._connections = ConnectionPool(tls=tls) self._service_discovery = ServiceDiscovery(service_address, self._connections, namespace) self._session_manager = SessionManager(self._service_discovery, session_timeout_ms, client_identifier) diff --git a/src/oxia/internal/connection_pool.py b/src/oxia/internal/connection_pool.py index 9699d7c..3c0574a 100644 --- a/src/oxia/internal/connection_pool.py +++ b/src/oxia/internal/connection_pool.py @@ -20,15 +20,20 @@ class ConnectionPool: - def __init__(self): + def __init__(self, tls: bool = False): self._lock = threading.Lock() + self._tls = tls self.connections = {} def get(self, address) -> OxiaClientStub: with self._lock: x = self.connections.get(address) if x is None: - channel = grpc.insecure_channel(address) + if self._tls: + channel = grpc.secure_channel( + address, grpc.ssl_channel_credentials()) + else: + channel = grpc.insecure_channel(address) stub = OxiaClientStub(channel) x = (channel, stub) self.connections[address] = x diff --git a/tests/tls_test.py b/tests/tls_test.py new file mode 100644 index 0000000..653059c --- /dev/null +++ b/tests/tls_test.py @@ -0,0 +1,92 @@ +# Copyright 2025 The Oxia Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for TLS channel creation.""" + +from unittest.mock import patch, MagicMock + +from oxia.internal.connection_pool import ConnectionPool + + +def test_insecure_channel_by_default(): + pool = ConnectionPool() + with patch("oxia.internal.connection_pool.grpc.insecure_channel") as insecure, \ + patch("oxia.internal.connection_pool.grpc.secure_channel") as secure: + insecure.return_value = MagicMock() + pool.get("localhost:6648") + insecure.assert_called_once_with("localhost:6648") + secure.assert_not_called() + + +def test_secure_channel_when_tls_true(): + pool = ConnectionPool(tls=True) + with patch("oxia.internal.connection_pool.grpc.insecure_channel") as insecure, \ + patch("oxia.internal.connection_pool.grpc.secure_channel") as secure, \ + patch("oxia.internal.connection_pool.grpc.ssl_channel_credentials") as creds: + secure.return_value = MagicMock() + creds.return_value = "fake-creds" + pool.get("oxia.example.com:6648") + + creds.assert_called_once() + secure.assert_called_once_with("oxia.example.com:6648", "fake-creds") + insecure.assert_not_called() + + +def test_tls_url_prefix_enables_tls(): + """service_address starting with tls:// must enable TLS and strip + the scheme, without needing an explicit tls=True.""" + # We can't actually construct a Client without a server, so test + # the prefix-handling logic at the boundary: ConnectionPool receives + # the stripped address and tls=True. + import oxia + + real_init = ConnectionPool.__init__ + init_calls = [] + + def capture_init(self, *args, **kwargs): + init_calls.append(kwargs) + real_init(self, *args, **kwargs) + + real_sd_init = None # will be patched to avoid starting a real thread + + with patch.object(ConnectionPool, "__init__", capture_init), \ + patch("oxia.client.ServiceDiscovery") as mock_sd, \ + patch("oxia.client.SessionManager"): + mock_sd.return_value = MagicMock() + oxia.Client("tls://example.com:6648") + + assert len(init_calls) == 1 + assert init_calls[0].get("tls") is True + # And the address passed to ServiceDiscovery must have the scheme stripped + _pos, sd_kwargs = mock_sd.call_args[0], mock_sd.call_args[1] + passed_address = mock_sd.call_args.args[0] + assert passed_address == "example.com:6648" + + +def test_explicit_tls_parameter_enables_tls(): + import oxia + + real_init = ConnectionPool.__init__ + init_calls = [] + + def capture_init(self, *args, **kwargs): + init_calls.append(kwargs) + real_init(self, *args, **kwargs) + + with patch.object(ConnectionPool, "__init__", capture_init), \ + patch("oxia.client.ServiceDiscovery"), \ + patch("oxia.client.SessionManager"): + oxia.Client("example.com:6648", tls=True) + + assert init_calls[0].get("tls") is True