proto-socket/kotlin/crosstest/kotlin_python.kt

@file:JvmName("KotlinPythonKt")

package com.tokilabs.proto_socket.crosstest

import com.tokilabs.proto_socket.TcpClient
import com.tokilabs.proto_socket.TcpServer
import com.tokilabs.proto_socket.ParserMap
import com.tokilabs.proto_socket.WsClient
import com.tokilabs.proto_socket.WsServer
import com.tokilabs.proto_socket.addListenerTyped
import com.tokilabs.proto_socket.addRequestListenerTyped
import com.tokilabs.proto_socket.typeNameOf
import com.tokilabs.proto_socket.packets.TestData
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.delay
import kotlinx.coroutines.launch
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.withContext
import kotlinx.coroutines.withTimeoutOrNull
import java.io.File
import java.security.KeyFactory
import java.security.KeyStore
import java.security.cert.CertificateFactory
import java.security.cert.X509Certificate
import java.security.spec.PKCS8EncodedKeySpec
import java.util.Base64
import java.util.concurrent.TimeUnit
import javax.net.ssl.KeyManagerFactory
import javax.net.ssl.SSLContext
import kotlin.system.exitProcess

private const val HOST = "127.0.0.1"
private const val TCP_PORT = 29394
private const val WS_PORT = 29396
private const val TLS_TCP_PORT = 29398
private const val WSS_PORT = 29400
private const val WS_PATH = "/"
private const val PROCESS_TIMEOUT_MS = 20_000L
private const val SERVER_OBSERVATION_MS = 200L
private const val STREAM_JOIN_TIMEOUT_MS = 5_000L

fun main() = runBlocking {
    println("INFO typeName kotlin=${typeNameOf<TestData>()}")
    val selectedTransport = System.getenv("PROTO_SOCKET_CROSSTEST_TRANSPORT")
    try {
        if (selectedTransport == null || selectedTransport == "tcp") {
            runTcpSendPush()
            delay(150)
            runTcpRequests()
            delay(150)
            runTcpLegacySimpleReceive()
            delay(150)
        }
        if (selectedTransport == null || selectedTransport == "ws") {
            runWs()
            delay(150)
        }
        if (selectedTransport == null || selectedTransport == "tls") {
            runTlsTcp()
            delay(150)
        }
        if (selectedTransport == null || selectedTransport == "wss") {
            runWss()
        }
        println("PASS all kotlin-server/python-client crosstests passed")
    } catch (error: Throwable) {
        System.err.println("FAIL crosstest error=${error.message ?: error}")
        exitProcess(1)
    }
}

private suspend fun runTcpSendPush() = coroutineScope {
    val received = CompletableDeferred<Boolean>()
    val server = TcpServer(HOST, TCP_PORT) { socket ->
        TcpClient.fromSocket(socket, 0, 0, parserMap())
    }
    server.onClientConnected = { client ->
        addListenerTyped<TestData>(client.communicator) { data ->
            println("SERVER_RECEIVED index=${data.index} message=${data.message}")
            val valid = data.index == 101 && data.message == "fire from python client"
            received.complete(valid)
            if (valid) {
                launch {
                    client.send(
                        TestData.newBuilder()
                            .setIndex(200)
                            .setMessage("push from kotlin server")
                            .build(),
                    )
                }
            }
        }
    }
    withServer(server::start, server::stop) {
        runPythonClient("tcp", TCP_PORT, "send-push", setOf("1", "2"))
        val ok = withTimeoutOrNull(SERVER_OBSERVATION_MS) { received.await() } ?: false
        check(ok) { "TCP send-push server did not receive expected data" }
    }
}

private suspend fun runTcpRequests() {
    val server = TcpServer(HOST, TCP_PORT) { socket ->
        TcpClient.fromSocket(socket, 0, 0, parserMap())
    }
    server.onClientConnected = { client ->
        addRequestListenerTyped<TestData, TestData>(client.communicator) { req ->
            TestData.newBuilder()
                .setIndex(req.index * 2)
                .setMessage("echo: ${req.message}")
                .build()
        }
    }
    withServer(server::start, server::stop) {
        runPythonClient("tcp", TCP_PORT, "requests", setOf("3", "4"))
    }
}

private suspend fun runTcpLegacySimpleReceive() = coroutineScope {
    val received = CompletableDeferred<Boolean>()
    val server = TcpServer(HOST, TCP_PORT) { socket ->
        TcpClient.fromSocket(socket, 0, 0, parserMap())
    }
    server.onClientConnected = { client ->
        addListenerTyped<TestData>(client.communicator) { data ->
            if (received.isActive) received.complete(data.index == 101)
        }
    }
    withServer(server::start, server::stop) {
        runPythonClient("tcp", TCP_PORT, "legacy-receive", emptySet())
        val ok = withTimeoutOrNull(SERVER_OBSERVATION_MS) { received.await() } ?: false
        check(ok) { "TCP legacy simple-name server received unexpected data" }
        println("PASS scenario=tcp-legacy-simple-receive detail=server-received-legacy-alias")
    }
}

private suspend fun runWs() = coroutineScope {
    val server = WsServer(HOST, WS_PORT, WS_PATH) { conn ->
        WsClient.forServer(conn, 0, 0, parserMap())
    }
    server.start()
    delay(100)
    try {
        runWsSendPush(server)
        delay(150)
        runWsRequests(server)
    } finally {
        server.stop()
    }
}

private suspend fun runWsSendPush(server: WsServer) = coroutineScope {
    val received = CompletableDeferred<Boolean>()
    server.onClientConnected = { client ->
        addListenerTyped<TestData>(client.communicator) { data ->
            println("SERVER_RECEIVED index=${data.index} message=${data.message}")
            val valid = data.index == 101 && data.message == "fire from python client"
            received.complete(valid)
            if (valid) {
                launch {
                    client.send(
                        TestData.newBuilder()
                            .setIndex(200)
                            .setMessage("push from kotlin server")
                            .build(),
                    )
                }
            }
        }
    }
    runPythonClient("ws", WS_PORT, "send-push", setOf("1", "2"))
    val ok = withTimeoutOrNull(SERVER_OBSERVATION_MS) { received.await() } ?: false
    check(ok) { "WS send-push server did not receive expected data" }
}

private suspend fun runWsRequests(server: WsServer) {
    server.onClientConnected = { client ->
        addRequestListenerTyped<TestData, TestData>(client.communicator) { req ->
            TestData.newBuilder()
                .setIndex(req.index * 2)
                .setMessage("echo: ${req.message}")
                .build()
        }
    }
    runPythonClient("ws", WS_PORT, "requests", setOf("3", "4"))
}

private suspend fun runTlsTcp() {
    val certFile = kotlinResourceFile("server.crt")
    val keyFile = kotlinResourceFile("server.key")
    val serverCtx = buildServerSslContext(certFile.path, keyFile.path)
    runTlsTcpSendPush(serverCtx, certFile.path)
    delay(150)
    runTlsTcpRequests(serverCtx, certFile.path)
}

private suspend fun runTlsTcpSendPush(serverCtx: SSLContext, certPath: String) = coroutineScope {
    val received = CompletableDeferred<Boolean>()
    val server = TcpServer(HOST, TLS_TCP_PORT, serverCtx) { socket ->
        TcpClient.fromSocket(socket, 0, 0, parserMap())
    }
    server.onClientConnected = { client ->
        addListenerTyped<TestData>(client.communicator) { data ->
            println("SERVER_RECEIVED index=${data.index} message=${data.message}")
            val valid = data.index == 101 && data.message == "fire from python client"
            received.complete(valid)
            if (valid) {
                launch {
                    client.send(
                        TestData.newBuilder()
                            .setIndex(200)
                            .setMessage("push from kotlin server")
                            .build(),
                    )
                }
            }
        }
    }
    withServer(server::start, server::stop) {
        runPythonClient("tls", TLS_TCP_PORT, "send-push", setOf("1", "2"), certPath)
        val ok = withTimeoutOrNull(SERVER_OBSERVATION_MS) { received.await() } ?: false
        check(ok) { "TLS TCP send-push server did not receive expected data" }
    }
}

private suspend fun runTlsTcpRequests(serverCtx: SSLContext, certPath: String) {
    val server = TcpServer(HOST, TLS_TCP_PORT, serverCtx) { socket ->
        TcpClient.fromSocket(socket, 0, 0, parserMap())
    }
    server.onClientConnected = { client ->
        addRequestListenerTyped<TestData, TestData>(client.communicator) { req ->
            TestData.newBuilder()
                .setIndex(req.index * 2)
                .setMessage("echo: ${req.message}")
                .build()
        }
    }
    withServer(server::start, server::stop) {
        runPythonClient("tls", TLS_TCP_PORT, "requests", setOf("3", "4"), certPath)
    }
}

private suspend fun runWss() = coroutineScope {
    val certFile = kotlinResourceFile("server.crt")
    val keyFile = kotlinResourceFile("server.key")
    val serverCtx = buildServerSslContext(certFile.path, keyFile.path)
    val server = WsServer(HOST, WSS_PORT, WS_PATH, sslContext = serverCtx) { conn ->
        WsClient.forServer(conn, 0, 0, parserMap())
    }
    server.start()
    delay(100)
    try {
        runWssSendPush(server, certFile.path)
        delay(150)
        runWssRequests(server, certFile.path)
    } finally {
        server.stop()
    }
}

private suspend fun runWssSendPush(server: WsServer, certPath: String) = coroutineScope {
    val received = CompletableDeferred<Boolean>()
    server.onClientConnected = { client ->
        addListenerTyped<TestData>(client.communicator) { data ->
            println("SERVER_RECEIVED index=${data.index} message=${data.message}")
            val valid = data.index == 101 && data.message == "fire from python client"
            received.complete(valid)
            if (valid) {
                launch {
                    client.send(
                        TestData.newBuilder()
                            .setIndex(200)
                            .setMessage("push from kotlin server")
                            .build(),
                    )
                }
            }
        }
    }
    runPythonClient("wss", WSS_PORT, "send-push", setOf("1", "2"), certPath)
    val ok = withTimeoutOrNull(SERVER_OBSERVATION_MS) { received.await() } ?: false
    check(ok) { "WSS send-push server did not receive expected data" }
}

private suspend fun runWssRequests(server: WsServer, certPath: String) {
    server.onClientConnected = { client ->
        addRequestListenerTyped<TestData, TestData>(client.communicator) { req ->
            TestData.newBuilder()
                .setIndex(req.index * 2)
                .setMessage("echo: ${req.message}")
                .build()
        }
    }
    runPythonClient("wss", WSS_PORT, "requests", setOf("3", "4"), certPath)
}

private suspend fun withServer(
    start: () -> Unit,
    stop: () -> Unit,
    body: suspend () -> Unit,
) {
    start()
    delay(100)
    try {
        body()
    } finally {
        stop()
    }
}

private fun destroyProcessTree(process: Process) {
    process.toHandle().descendants().forEach { it.destroyForcibly() }
    process.destroyForcibly()
}

private suspend fun runPythonClient(
    mode: String,
    port: Int,
    phase: String,
    expectedScenarios: Set<String>,
    certPath: String? = null,
) = withContext(Dispatchers.IO) {
    val args = mutableListOf(
        "python3",
        "crosstest/kotlin_python_client.py",
        "--mode=$mode",
        "--port=$port",
        "--phase=$phase",
    )
    if (certPath != null) {
        args += "--cert=$certPath"
    }
    val process = ProcessBuilder(args)
        .directory(pythonDir())
        .redirectErrorStream(false)
        .start()

    val stdoutLines = mutableListOf<String>()
    val stdoutThread = Thread {
        process.inputStream.bufferedReader().forEachLine { line ->
            println(line)
            if (line.startsWith("PASS ") || line.startsWith("FAIL ")) {
                stdoutLines.add(line)
            }
        }
    }
    val stderrThread = Thread {
        process.errorStream.bufferedReader().forEachLine { line ->
            System.err.println(line)
        }
    }
    stdoutThread.start()
    stderrThread.start()

    val finished = process.waitFor(PROCESS_TIMEOUT_MS, TimeUnit.MILLISECONDS)
    if (!finished) {
        destroyProcessTree(process)
    }
    stdoutThread.join(STREAM_JOIN_TIMEOUT_MS)
    stderrThread.join(STREAM_JOIN_TIMEOUT_MS)
    val label = "python-client $mode/$phase"
    check(!stdoutThread.isAlive && !stderrThread.isAlive) {
        "$label stream readers did not finish"
    }

    validateResultLines(
        label,
        if (finished) process.exitValue() else -1,
        stdoutLines,
        expectedScenarios,
    )
}

private fun validateResultLines(
    label: String,
    exitCode: Int,
    lines: List<String>,
    expectedScenarios: Set<String>,
) {
    val failed = lines.filter { it.startsWith("FAIL ") }
    val passed = lines
        .filter { it.startsWith("PASS ") }
        .mapNotNull { Regex("""scenario=([^ ]+)""").find(it)?.groupValues?.get(1) }
        .toSet()
    val missing = expectedScenarios - passed
    check(exitCode == 0 && failed.isEmpty() && missing.isEmpty()) {
        "$label failed exitCode=$exitCode failed=$failed missing=$missing"
    }
}

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

private fun buildServerSslContext(certPath: String, keyPath: String): SSLContext {
    val certFactory = CertificateFactory.getInstance("X.509")
    val cert = File(certPath).inputStream().use {
        certFactory.generateCertificate(it) as X509Certificate
    }
    val keyPem = File(keyPath).readText()
        .replace("-----BEGIN PRIVATE KEY-----", "")
        .replace("-----END PRIVATE KEY-----", "")
        .replace("\\s+".toRegex(), "")
    val privateKey = KeyFactory.getInstance("RSA")
        .generatePrivate(PKCS8EncodedKeySpec(Base64.getDecoder().decode(keyPem)))
    val keyStore = KeyStore.getInstance("PKCS12")
    keyStore.load(null, null)
    keyStore.setKeyEntry("toki", privateKey, CharArray(0), arrayOf(cert))
    val keyManagerFactory = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm())
    keyManagerFactory.init(keyStore, CharArray(0))
    val ctx = SSLContext.getInstance("TLS")
    ctx.init(keyManagerFactory.keyManagers, null, null)
    return ctx
}

private fun kotlinResourceFile(name: String): File {
    var dir = File(System.getProperty("user.dir")).absoluteFile
    while (dir.parentFile != null) {
        val direct = File(dir, "src/test/resources/$name").canonicalFile
        if (direct.isFile) return direct
        val nested = File(dir, "kotlin/src/test/resources/$name").canonicalFile
        if (nested.isFile) return nested
        dir = dir.parentFile
    }
    error("cannot resolve kotlin test resource $name")
}

private fun pythonDir(): File {
    var dir = File(System.getProperty("user.dir")).absoluteFile
    while (dir.parentFile != null) {
        val candidate = File(dir, "../python").canonicalFile
        if (File(candidate, "pyproject.toml").isFile) return candidate
        val direct = File(dir, "python").canonicalFile
        if (File(direct, "pyproject.toml").isFile) return direct
        dir = dir.parentFile
    }
    error("cannot resolve python directory")
}