feat: support tls:// schema to auto enable tls (#192)

* feat: support tls:// schema to auto enable tls * add test * fix spotless
streamnative · Nov 16, 2024 · f45c4b2 · f45c4b2
1 parent 7ffdddd
commit f45c4b2
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 9 deletions.
diff --git a/client/src/main/java/io/streamnative/oxia/client/grpc/OxiaStub.java b/client/src/main/java/io/streamnative/oxia/client/grpc/OxiaStub.java
@@ -19,19 +19,16 @@
 
 import com.google.common.base.Throwables;
 import io.grpc.CallCredentials;
+import io.grpc.ChannelCredentials;
 import io.grpc.Grpc;
 import io.grpc.InsecureChannelCredentials;
 import io.grpc.ManagedChannel;
-import io.grpc.Metadata;
 import io.grpc.TlsChannelCredentials;
 import io.grpc.internal.BackoffPolicy;
-import io.grpc.stub.MetadataUtils;
 import io.streamnative.oxia.client.api.Authentication;
 import io.streamnative.oxia.proto.OxiaClientGrpc;
 
 import java.lang.reflect.Field;
-import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executor;
 import javax.annotation.Nullable;
 
@@ -40,19 +37,30 @@
 
 @Slf4j
 public class OxiaStub implements AutoCloseable {
+    public static String TLS_SCHEMA = "tls://";
     private final ManagedChannel channel;
     private final @NonNull OxiaClientGrpc.OxiaClientStub asyncStub;
 
+    static String getAddress(String address) {
+        if (address.startsWith(TLS_SCHEMA)) {
+            return address.substring(TLS_SCHEMA.length());
+        }
+        return address;
+    }
+
+    static ChannelCredentials getChannelCredential(String address, boolean tlsEnabled) {
+        return tlsEnabled || address.startsWith(TLS_SCHEMA)
+                ? TlsChannelCredentials.newBuilder().build()
+                : InsecureChannelCredentials.create();
+    }
+
     public OxiaStub(
             String address,
             @Nullable Authentication authentication,
             boolean enableTls,
             @Nullable BackoffPolicy.Provider backoffProvider) {
-        this(Grpc.newChannelBuilder(
-                                address,
-                                enableTls
-                                        ? TlsChannelCredentials.newBuilder().build()
-                                        : InsecureChannelCredentials.create())
+
+        this(Grpc.newChannelBuilder(getAddress(address), getChannelCredential(address, enableTls))
                         .directExecutor()
                         .build(),
                 authentication, backoffProvider);

diff --git a/client/src/test/java/io/streamnative/oxia/client/grpc/OxiaStubTest.java b/client/src/test/java/io/streamnative/oxia/client/grpc/OxiaStubTest.java
@@ -15,6 +15,8 @@
  */
 package io.streamnative.oxia.client.grpc;
 
+import io.grpc.InsecureChannelCredentials;
+import io.grpc.TlsChannelCredentials;
 import io.grpc.stub.StreamObserver;
 import io.streamnative.oxia.proto.GetRequest;
 import io.streamnative.oxia.proto.ReadRequest;
@@ -137,4 +139,30 @@ public void testMaxConnectionPerNode() {
         }
         Assertions.assertEquals(maxConnectionPerNode, stubManager.stubs.size());
     }
+
+    @Test
+    public void testAddressTrim() {
+        final var tlsAddress = "tls://localhost:6648";
+        Assertions.assertEquals("localhost:6648", OxiaStub.getAddress(tlsAddress));
+
+        final var planTxtAddress = "localhost:6648";
+        Assertions.assertEquals("localhost:6648", OxiaStub.getAddress(planTxtAddress));
+    }
+
+    @Test
+    public void testTlsCredential() {
+        final var tlsAddress = "tls://localhost:6648";
+        var channelCredential = OxiaStub.getChannelCredential(tlsAddress, false);
+        Assertions.assertInstanceOf(TlsChannelCredentials.class, channelCredential);
+
+        channelCredential = OxiaStub.getChannelCredential(tlsAddress, true);
+        Assertions.assertInstanceOf(TlsChannelCredentials.class, channelCredential);
+
+        final var planTxtAddress = "localhost:6648";
+        channelCredential = OxiaStub.getChannelCredential(planTxtAddress, false);
+        Assertions.assertInstanceOf(InsecureChannelCredentials.class, channelCredential);
+
+        channelCredential = OxiaStub.getChannelCredential(planTxtAddress, true);
+        Assertions.assertInstanceOf(TlsChannelCredentials.class, channelCredential);
+    }
 }