support tls for grpc websocket (#378)

Support tls 

E2E Production test:
1. to server which does not  support tls
a. mode = disabled (default)
b. enable tls or mtls will error --> trigger retry and fallback logic

2. to server which requires tls
   1. tls with valid certs -> connect
   2. mtls with valid client certs -> connect
other cases fallback to use insecure channel which ultimately will fail
and trigger retry and fallback logic

Added unit test for error handling behavior
more e2e integration test to be found here
statsig-io/long-running-sdk#29

Author: xinlili-statsig

src/main/kotlin/com/statsig/sdk/SpecUpdater.kt

@@ -35,6 +35,11 @@ internal class SpecUpdater(
     private inline fun <reified T> Gson.fromJson(json: String) = fromJson<T>(json, object : TypeToken<T>() {}.type)
 
     fun initialize() {
+        transport.setStreamingFallback(NetworkEndpoint.DOWNLOAD_CONFIG_SPECS) {
+            this.transport.downloadConfigSpecsFromStatsig(
+                this.lastUpdateTime,
+            ).first
+        }
         transport.initialize()
     }
 
@@ -72,11 +77,6 @@ internal class SpecUpdater(
                 idListFlow.collect { idListCallback(it) }
             }
         }
-        transport.setStreamingFallback(NetworkEndpoint.DOWNLOAD_CONFIG_SPECS) {
-            this.transport.downloadConfigSpecsFromStatsig(
-                this.lastUpdateTime,
-            ).first
-        }
     }
 
     fun getInitializeOrder(): List<DataSource> {

src/main/kotlin/com/statsig/sdk/StatsigOptions.kt

@@ -5,6 +5,7 @@ import com.statsig.sdk.datastore.IDataStore
 import com.statsig.sdk.network.STATSIG_API_URL_BASE
 import com.statsig.sdk.persistent_storage.IUserPersistentStorage
 import com.statsig.sdk.persistent_storage.UserPersistedValues
+import java.io.InputStream
 import java.time.Instant
 import java.time.format.DateTimeFormatter
 import java.time.temporal.ChronoUnit
@@ -208,6 +209,17 @@ enum class NetworkProtocol {
     GRPC_WEBSOCKET,
 }
 
+enum class AuthenticationMode {
+    @SerializedName("disabled")
+    DISABLED,
+
+    @SerializedName("tls")
+    TLS,
+
+    @SerializedName("mTls")
+    MTLS,
+}
+
 data class ForwardProxyConfig(
     @SerializedName("proxyAddress") var proxyAddress: String,
     @SerializedName("protocol") val proxyProtocol: NetworkProtocol,
@@ -216,6 +228,12 @@ data class ForwardProxyConfig(
     @SerializedName("retry_backoff_multiplier") val retryBackoffMultiplier: Int? = null,
     @SerializedName("retry_backoff_base_ms") val retryBackoffBaseMs: Long? = null,
     @SerializedName("push_worker_failover_threshold") val pushWorkerFailoverThreshold: Int? = null,
+
+    // TLS Certification
+    @SerializedName("authentication_mode") var authenticationMode: AuthenticationMode = AuthenticationMode.DISABLED,
+    @SerializedName("tls_cert_chain") var tlsCertChain: InputStream? = null,
+    @SerializedName("tls_private_key") var tlsPrivateKey: InputStream? = null,
+    @SerializedName("tls_private_key_password") var tlsPrivateKeyPassword: InputStream? = null,
 )
 
 data class ProxyConfig @JvmOverloads constructor(

src/main/kotlin/com/statsig/sdk/StatsigServer.kt

@@ -286,19 +286,25 @@ private class StatsigServerImpl() :
     private var setupStartTime = 0L
 
     override fun setup(serverSecret: String, options: StatsigOptions) {
-        Thread.setDefaultUncaughtExceptionHandler(MainThreadExceptionHandler(this, Thread.currentThread()))
-        setupStartTime = System.currentTimeMillis()
-        errorBoundary = ErrorBoundary(serverSecret, options, statsigMetadata)
-        coroutineExceptionHandler = CoroutineExceptionHandler { _, ex ->
-            // no-op - supervisor job should not throw when a child fails
-            errorBoundary.logException("coroutineExceptionHandler", ex)
-        }
-        statsigJob = SupervisorJob()
-        statsigScope = CoroutineScope(statsigJob + coroutineExceptionHandler)
-        transport = StatsigTransport(serverSecret, options, statsigMetadata, statsigScope, errorBoundary, sdkConfigs)
-        logger = StatsigLogger(statsigScope, transport, statsigMetadata, options, sdkConfigs)
-        options.customLogger.also { outputLogger = it }
-        this.options = options
+        try {
+            Thread.setDefaultUncaughtExceptionHandler(MainThreadExceptionHandler(this, Thread.currentThread()))
+            setupStartTime = System.currentTimeMillis()
+            errorBoundary = ErrorBoundary(serverSecret, options, statsigMetadata)
+            coroutineExceptionHandler = CoroutineExceptionHandler { _, ex ->
+                // no-op - supervisor job should not throw when a child fails
+                errorBoundary.logException("coroutineExceptionHandler", ex)
+            }
+            statsigJob = SupervisorJob()
+            statsigScope = CoroutineScope(statsigJob + coroutineExceptionHandler)
+            transport = StatsigTransport(serverSecret, options, statsigMetadata, statsigScope, errorBoundary, sdkConfigs)
+            logger = StatsigLogger(statsigScope, transport, statsigMetadata, options, sdkConfigs)
+            options.customLogger.also { outputLogger = it }
+            this.options = options
+        } catch (e: Throwable) {
+            // noop swallow and let other part handle error
+            options.customLogger.warn("[STATSIG]Failed to setup sdk")
+            options.customLogger.warn(e.stackTraceToString())
+        }
     }
 
     override suspend fun initialize(serverSecret: String, options: StatsigOptions): InitializationDetails? {

src/main/kotlin/com/statsig/sdk/network/GRPCWebsocketWorker.kt

@@ -1,18 +1,34 @@
 package com.statsig.sdk.network
 
-import com.statsig.sdk.*
+import com.statsig.sdk.AuthenticationMode
+import com.statsig.sdk.Diagnostics
+import com.statsig.sdk.ErrorBoundary
+import com.statsig.sdk.FailureDetails
+import com.statsig.sdk.FailureReason
+import com.statsig.sdk.ForwardProxyConfig
+import com.statsig.sdk.KeyType
+import com.statsig.sdk.NetworkProtocol
+import com.statsig.sdk.StatsigEvent
+import com.statsig.sdk.StatsigOptions
 import grpc.generated.statsig_forward_proxy.StatsigForwardProxyGrpc
 import grpc.generated.statsig_forward_proxy.StatsigForwardProxyOuterClass.ConfigSpecRequest
 import grpc.generated.statsig_forward_proxy.StatsigForwardProxyOuterClass.ConfigSpecResponse
 import io.grpc.Channel
+import io.grpc.Grpc
 import io.grpc.ManagedChannelBuilder
+import io.grpc.TlsChannelCredentials
 import io.grpc.stub.StreamObserver
-import kotlinx.coroutines.*
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.Job
 import kotlinx.coroutines.channels.BufferOverflow
+import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.asSharedFlow
 import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.withTimeoutOrNull
 
 private const val RETRY_LIMIT = 10
 private const val INITIAL_RETRY_BACKOFF_MS: Long = 10 * 1000
@@ -30,8 +46,6 @@ internal class GRPCWebsocketWorker(
     override val isPullWorker: Boolean = false
 
     private var diagnostics: Diagnostics? = null
-    private val channel: Channel = ManagedChannelBuilder.forTarget(proxyConfig.proxyAddress).usePlaintext().build()
-    private val stub = StatsigForwardProxyGrpc.newStub(channel)
     private val logger = options.customLogger
 
     private val observer = object : StreamObserver<ConfigSpecResponse> {
@@ -64,6 +78,59 @@ internal class GRPCWebsocketWorker(
     private var connected: Boolean = true
     private var downloadConfigsJob: Job? = null
     var streamingFallback: StreamingFallback? = null
+    private val stub: StatsigForwardProxyGrpc.StatsigForwardProxyStub
+
+    private var channel: Channel
+    init {
+        try {
+            channel = setupChannel()
+        } catch (e: Exception) {
+            options.customLogger.warn("Failed setup channel falling back to insecure channel")
+            channel = ManagedChannelBuilder.forTarget(proxyConfig.proxyAddress).usePlaintext().build()
+        }
+
+        stub = StatsigForwardProxyGrpc.newStub(channel)
+    }
+
+    private fun setupChannel(): Channel {
+        when (proxyConfig.authenticationMode) {
+            AuthenticationMode.DISABLED ->
+                return ManagedChannelBuilder.forTarget(proxyConfig.proxyAddress).usePlaintext().build()
+            AuthenticationMode.MTLS -> {
+                val tlsBuilder = TlsChannelCredentials.newBuilder()
+                proxyConfig.tlsCertChain.let {
+                    if (it != null) {
+                        tlsBuilder.trustManager(it)
+                    } else {
+                        options.customLogger.warn("Failed to get cert chain for tls")
+                    }
+                    proxyConfig.tlsPrivateKey.let { privateKey ->
+                        val privateKeyPassword = proxyConfig.tlsPrivateKeyPassword
+                        if (privateKey != null && privateKeyPassword != null) {
+                            tlsBuilder.keyManager(privateKey, privateKeyPassword)
+                            return Grpc.newChannelBuilder(proxyConfig.proxyAddress, tlsBuilder.build()).build()
+                        } else {
+                            options.customLogger.warn("Failed to get private key and password for tls")
+                            return Grpc.newChannelBuilder(proxyConfig.proxyAddress, tlsBuilder.build()).build()
+                        }
+                    }
+                }
+            }
+            AuthenticationMode.TLS -> {
+                val tlsBuilder = TlsChannelCredentials.newBuilder()
+                proxyConfig.tlsCertChain.let {
+                    if (it != null) {
+                        tlsBuilder.trustManager(it)
+                        return Grpc.newChannelBuilder(proxyConfig.proxyAddress, tlsBuilder.build()).build()
+                    } else {
+                        options.customLogger.warn("Failed to get cert chain for tls")
+                    }
+                }
+            }
+        }
+        options.customLogger.warn("Falling back to insecure channel")
+        return ManagedChannelBuilder.forTarget(proxyConfig.proxyAddress).usePlaintext().build()
+    }
 
     private val dcsFlowBacker = MutableSharedFlow<String>(
         replay = 1,

src/test/java/com/statsig/sdk/GRPCTest.kt

@@ -1,16 +1,21 @@
 package com.statsig.sdk
 
+import com.statsig.sdk.network.GRPCWebsocketWorker
 import com.statsig.sdk.network.GRPCWorker
 import grpc.generated.statsig_forward_proxy.StatsigForwardProxyGrpc
 import grpc.generated.statsig_forward_proxy.StatsigForwardProxyOuterClass.ConfigSpecRequest
 import grpc.generated.statsig_forward_proxy.StatsigForwardProxyOuterClass.ConfigSpecResponse
 import io.grpc.stub.StreamObserver
 import io.grpc.testing.GrpcServerRule
 import io.mockk.mockk
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.delay
 import kotlinx.coroutines.runBlocking
 import org.junit.Assert.assertEquals
 import org.junit.Rule
 import org.junit.Test
+import java.io.ByteArrayInputStream
 
 class GRPCTest {
     @Rule
@@ -21,12 +26,28 @@ class GRPCTest {
     @Rule
     val retry = RetryRule(3)
 
-    private fun setupGRPC(spec: ConfigSpecResponse) {
+    private fun setupGRPC(spec: ConfigSpecResponse, shouldThrowOnStreaming: Boolean = false, blockStreamTime: Long? = null) {
         val service = object : StatsigForwardProxyGrpc.StatsigForwardProxyImplBase() {
             override fun getConfigSpec(request: ConfigSpecRequest?, responseObserver: StreamObserver<ConfigSpecResponse>?) {
                 responseObserver?.onNext(spec)
                 responseObserver?.onCompleted()
             }
+
+            override fun streamConfigSpec(
+                request: ConfigSpecRequest?,
+                responseObserver: StreamObserver<ConfigSpecResponse>?,
+            ) {
+                if (shouldThrowOnStreaming) {
+                    responseObserver?.onError(Exception("io exception"))
+                }
+                if (blockStreamTime != null) {
+                    runBlocking {
+                        delay(blockStreamTime)
+                    }
+                }
+                responseObserver?.onNext(spec)
+                responseObserver?.onCompleted()
+            }
         }
         grpcServerRule.serviceRegistry.addService(service)
     }
@@ -46,4 +67,59 @@ class GRPCTest {
 
         assertEquals(Pair("spec-1", null), worker.downloadConfigSpecs(0))
     }
+
+    @Test
+    fun testTLSConfigurationError() = runBlocking {
+        setupGRPC(ConfigSpecResponse.newBuilder().setSpec("spec-1").setLastUpdated(123).build())
+        val options = StatsigOptions()
+        val boundary = mockk<ErrorBoundary>()
+        val scope = CoroutineScope(Dispatchers.Default)
+        val proxyConfig = ForwardProxyConfig(
+            "proxy:8000",
+            NetworkProtocol.GRPC_WEBSOCKET,
+            retryBackoffMultiplier = 1,
+            retryBackoffBaseMs = 10,
+            authenticationMode = AuthenticationMode.TLS,
+            // We will try catch this invalid cert chain and fallback to insecure channel
+            tlsCertChain = ByteArrayInputStream("invalid".toByteArray()),
+        )
+        val worker = GRPCWebsocketWorker("sdk", options, scope, boundary, proxyConfig)
+        val stub = StatsigForwardProxyGrpc.newStub(grpcServerRule.channel)
+
+        val stubField = GRPCWebsocketWorker::class.java.getDeclaredField("stub")
+        stubField.isAccessible = true
+        stubField.set(worker, stub)
+        worker.initializeFlows()
+        // In this test, stream will be successful, because we mock
+        assertEquals(Pair("spec-1", null), worker.downloadConfigSpecs(0))
+    }
+
+    /*
+* Test when server returns error we start retry
+* This error can be because of error internal unavailability, authentication failed
+* */
+    @Test
+    fun testTimeout() = runBlocking {
+        setupGRPC(ConfigSpecResponse.newBuilder().setSpec("spec-1").setLastUpdated(123).build(), blockStreamTime = 500)
+        val options = StatsigOptions(initTimeoutMs = 100)
+        val boundary = mockk<ErrorBoundary>()
+        val scope = CoroutineScope(Dispatchers.Default)
+        val proxyConfig = ForwardProxyConfig(
+            "proxy:8000",
+            NetworkProtocol.GRPC_WEBSOCKET,
+            retryBackoffMultiplier = 1,
+            retryBackoffBaseMs = 10,
+            authenticationMode = AuthenticationMode.DISABLED,
+        )
+        val worker = GRPCWebsocketWorker("sdk", options, scope, boundary, proxyConfig)
+        val stub = StatsigForwardProxyGrpc.newStub(grpcServerRule.channel)
+
+        val stubField = GRPCWebsocketWorker::class.java.getDeclaredField("stub")
+        stubField.isAccessible = true
+        stubField.set(worker, stub)
+        worker.initializeFlows()
+        val details = worker.downloadConfigSpecs(0)
+        assert(details.first == null)
+        assert(details.second?.exception?.message?.contains("failed to receive config spec within init timeout") == true)
+    }
 }