proto-socket/kotlin/crosstest/python_kotlin_client.kt

@file:JvmName("PythonKotlinClientKt")

package com.tokilabs.proto_socket.crosstest

import com.tokilabs.proto_socket.DialTcp
import com.tokilabs.proto_socket.DialTcpTls
import com.tokilabs.proto_socket.DialWs
import com.tokilabs.proto_socket.DialWss
import com.tokilabs.proto_socket.ParserMap
import com.tokilabs.proto_socket.WsClient
import com.tokilabs.proto_socket.TcpClient
import com.tokilabs.proto_socket.addListenerTyped
import com.tokilabs.proto_socket.sendRequestTyped
import com.tokilabs.proto_socket.typeNameOf
import com.tokilabs.proto_socket.packets.PacketBase
import com.tokilabs.proto_socket.packets.TestData
import com.google.protobuf.ByteString
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.async
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.delay
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.withTimeout
import kotlin.system.exitProcess

private const val HOST = "127.0.0.1"
private const val WS_PATH = "/"
private const val CONNECT_WINDOW_MS = 3_000L
private const val REQUEST_WINDOW_MS = 10_000L

private interface PythonClientHandle {
    suspend fun send(data: TestData)
    suspend fun close()
    val communicator: com.tokilabs.proto_socket.Communicator
}

private class PythonTcpHandle(
    private val client: TcpClient,
) : PythonClientHandle {
    override val communicator = client.communicator
    override suspend fun send(data: TestData) { client.send(data) }
    override suspend fun close() { client.close() }
}

private class PythonWsHandle(
    private val client: WsClient,
) : PythonClientHandle {
    override val communicator = client.communicator
    override suspend fun send(data: TestData) { client.send(data) }
    override suspend fun close() { client.close() }
}

fun main(args: Array<String>) = runBlocking {
    val mode = argValue(args, "mode") ?: "tcp"
    val phase = argValue(args, "phase") ?: "send-push"
    val port = argValue(args, "port")?.toIntOrNull()
    val cert = argValue(args, "cert")

    println("INFO typeName kotlin=${typeNameOf<TestData>()}")

    if (port == null) {
        fail("setup", "port is required")
        exitProcess(1)
    }

    val client = try {
        connectWithRetry(mode, port, cert)
    } catch (error: Throwable) {
        fail("setup", error.message ?: error.toString())
        exitProcess(1)
    }

    val ok: Boolean = try {
        when (phase) {
            "send-push" -> runSendPush(client)
            "requests" -> runRequests(client)
            "legacy-receive" -> { runLegacyReceive(client.communicator); true }
            else -> {
                fail("setup", "unknown phase $phase")
                false
            }
        }
    } finally {
        client.close()
    }

    if (!ok) exitProcess(1)
}

private suspend fun connectWithRetry(mode: String, port: Int, cert: String?): PythonClientHandle {
    val deadline = System.nanoTime() + CONNECT_WINDOW_MS * 1_000_000L
    var lastError: Throwable? = null
    while (System.nanoTime() < deadline) {
        try {
            return when (mode) {
                "tcp" -> PythonTcpHandle(DialTcp(HOST, port, 0, 0, parserMap()))
                "ws" -> PythonWsHandle(DialWs(HOST, port, WS_PATH, 0, 0, parserMap()))
                "tls" -> {
                    val sslCtx = buildClientSslContext(cert ?: error("--cert required for tls"))
                    PythonTcpHandle(DialTcpTls(HOST, port, sslCtx, 0, 0, parserMap()))
                }
                "wss" -> {
                    val sslCtx = buildClientSslContext(cert ?: error("--cert required for wss"))
                    PythonWsHandle(DialWss(HOST, port, WS_PATH, sslCtx, 0, 0, parserMap()))
                }
                else -> error("unknown mode $mode")
            }
        } catch (error: Throwable) {
            lastError = error
            delay(100)
        }
    }
    error("connect $mode:$port timed out: $lastError")
}

private fun buildClientSslContext(certPath: String): javax.net.ssl.SSLContext {
    val certFactory = java.security.cert.CertificateFactory.getInstance("X.509")
    val cert = java.io.File(certPath).inputStream().use {
        certFactory.generateCertificate(it) as java.security.cert.X509Certificate
    }
    val trustStore = java.security.KeyStore.getInstance("PKCS12")
    trustStore.load(null, null)
    trustStore.setCertificateEntry("toki", cert)
    val trustManagerFactory = javax.net.ssl.TrustManagerFactory.getInstance(
        javax.net.ssl.TrustManagerFactory.getDefaultAlgorithm(),
    )
    trustManagerFactory.init(trustStore)
    val ctx = javax.net.ssl.SSLContext.getInstance("TLS")
    ctx.init(null, trustManagerFactory.trustManagers, null)
    return ctx
}

private suspend fun runLegacyReceive(communicator: com.tokilabs.proto_socket.Communicator) {
    val data = TestData.newBuilder().setIndex(101).setMessage("legacy fire from kotlin").build().toByteArray()
    communicator.queuePacket(
        PacketBase.newBuilder()
            .setTypeName("TestData")
            .setNonce(1)
            .setData(ByteString.copyFrom(data))
            .build()
    )
    delay(150)
}

private suspend fun runSendPush(client: PythonClientHandle): Boolean {
    val push = CompletableDeferred<TestData>()
    addListenerTyped<TestData>(client.communicator) {
        push.complete(it)
    }

    return try {
        client.send(
            TestData.newBuilder()
                .setIndex(101)
                .setMessage("fire from kotlin client")
                .build(),
        )
        pass("1", "fire-and-forget sent")

        val msg = withTimeout(REQUEST_WINDOW_MS) { push.await() }
        if (msg.index != 200 || msg.message != "push from python server") {
            fail("2", "unexpected push index=${msg.index} message=${msg.message}")
            false
        } else {
            pass("2", "received push from python server")
            true
        }
    } catch (error: Throwable) {
        fail("2", error.message ?: error.toString())
        false
    }
}

private suspend fun runRequests(client: PythonClientHandle): Boolean =
    try {
        val single = sendRequestTyped<TestData, TestData>(
            client.communicator,
            TestData.newBuilder()
                .setIndex(21)
                .setMessage("single request from kotlin")
                .build(),
            REQUEST_WINDOW_MS,
        )
        if (single.index != 42 || single.message != "echo: single request from kotlin") {
            fail("3", "unexpected response index=${single.index} message=${single.message}")
            false
        } else {
            pass("3", "single request response matched")
            runConcurrentRequests(client)
        }
    } catch (error: Throwable) {
        fail("3", error.message ?: error.toString())
        false
    }

private suspend fun runConcurrentRequests(client: PythonClientHandle): Boolean =
    coroutineScope {
        val jobs = (0 until concurrencyCount()).map { i ->
            async {
                val index = 30 + i
                val message = "multi request $i from kotlin"
                val res = sendRequestTyped<TestData, TestData>(
                    client.communicator,
                    TestData.newBuilder()
                        .setIndex(index)
                        .setMessage(message)
                        .build(),
                    REQUEST_WINDOW_MS,
                )
                check(res.index == index * 2 && res.message == "echo: $message") {
                    "request $i got index=${res.index} message=${res.message}"
                }
            }
        }
        try {
            jobs.forEach { it.await() }
            pass("4", "concurrent request responses matched")
            true
        } catch (error: Throwable) {
            fail("4", error.message ?: error.toString())
            false
        }
    }

private fun concurrencyCount(): Int {
    val raw = System.getenv("PROTO_SOCKET_CROSSTEST_CONCURRENCY") ?: return 5
    val value = raw.toIntOrNull() ?: return 5
    return if (value > 0) value else 5
}

private fun parserMap(): ParserMap =
    mapOf(typeNameOf<TestData>() to { TestData.parseFrom(it) })

private fun argValue(args: Array<String>, name: String): String? {
    val prefix = "--$name="
    return args.firstOrNull { it.startsWith(prefix) }?.substring(prefix.length)
}

private fun pass(scenario: String, detail: String) {
    println("PASS scenario=$scenario detail=$detail")
}

private fun fail(scenario: String, error: String) {
    println("FAIL scenario=$scenario error=$error")
}