KTor TCP Proxy Improvements - kotlin

I'm trying to build a performant tcp proxy to bi-directionally forward connections from client A to server B via a proxy server C.
The connection would look like so:
A <——> C <——> B
My initial plan was to use KTor raw TCP Sockets, but the initial example on the website was pretty lack luster and didn’t really point me in the direction of accomplishing my task.
Some things I wanted to be able to do:
register callbacks for when the client is closed
register callbacks for when the server is closed
clean up resources when either the client or server close
point the data from one tcp buffer to another without copying it over
currently I have a very naive implementation as follows:
fun main() {
runBlocking {
launch(Dispatchers.IO) {
val proxySelectorManager = SelectorManager(Dispatchers.IO)
val proxyServer = aSocket(proxySelectorManager).tcp().bind(HOST, PORT)
println("Listening on: ${HOST}:${PORT}")
while (true) {
val clientSocket = proxyServer.accept()
val serverSelectorManager = SelectorManager(Dispatchers.IO)
val serverSocket = aSocket(serverSelectorManager).tcp().connect(REMOTE_HOST, REMOTE_PORT)
println("Client joined at: ${clientSocket.remoteAddress} -> ${clientSocket.localAddress}")
println("Connected to Server at: ${serverSocket.localAddress} -> ${serverSocket.remoteAddress}")
launch(Dispatchers.IO) {
val clientReadChannel = clientSocket.openReadChannel()
val clientWriteChannel = clientSocket.openWriteChannel(true)
val serverReadChannel = serverSocket.openReadChannel()
val serverWriteChannel = serverSocket.openWriteChannel(true)
launch(Dispatchers.IO) {
while (true) {
try { serverWriteChannel.writeByte(clientReadChannel.readByte()) }
catch (e: Exception) { break }
}
}
launch(Dispatchers.IO) {
while (true) {
try { clientWriteChannel.writeByte(serverReadChannel.readByte()) }
catch (e: Exception) { break }
}
}
}
}
}
}
}
My main issues currently are that I have no way of really knowing when a connection was closed (either the server or the client) and my current method of copying over the buffers is extremely slow and has lots of overhead.
Any thoughts or ideas on how I can accomplish the following are greatly appreciated.

Related

What is Kotlin's functional equivalent for finally?

This example is from the documentation of HttpUrlConnection:
URL url = new URL("http://www.android.com/");
HttpURLConnection urlConnection = (HttpURLConnection) url.openConnection();
try {
InputStream in = new BufferedInputStream(urlConnection.getInputStream());
readStream(in);
}
finally {
urlConnection.disconnect();
}
The documentation says:
Once the response body has been read, the HttpURLConnection should be closed by calling disconnect().
I tried use the Java class to load an image in a functional style:
fun fetch (url: String): ImageBitmap =
URL(url)
.openConnection()
.also { it.setRequestProperty (authorization.header, authorization.basicauth()) }
.getInputStream()
.buffered()
.use { BitmapFactory.decodeStream(it) }
.asImageBitmap()
Now I am wondering how to add the disconnect call?
I want to achieve this:
fun fetch (url: String): ImageBitmap {
var connection: HttpURLConnection? = null
return try {
URL(url)
.openConnection()
.also { it.setRequestProperty(authorization.header, authorization.basicauth()) }
.also { connection = it as HttpURLConnection }
.getInputStream()
.buffered()
.use { BitmapFactory.decodeStream(it) }
.asImageBitmap()
} finally {
connection?.disconnect()
}
}
But in a less ugly manner.
There isn't a stdlib solution for your case, but kotlin has the use() extension method defined on (Auto)Closeable to make this pattern more functional for those interfaces.
You could add a use extension method to HttpUrlConnection yourself that calls its disconnect() method, using the same approach as the source of use.
Of course you would still need to write the try finally once, but it is now hidden when using HttpUrlConnection.
On first sight you'd end up with something like this, you might still need some null handling somewhere.
public fun <T : HttpURLConnection, R> T.use(block: (T) -> R): R {
try {
return block(this)
} finally {
disconnect()
}
}
(URL(url).openConnection() as HttpURLConnection).use {
// do things with connection `it`
}
// connection is now disconnected()
Your existing way of doing this is clear enough as it is. Kotlin is not a purely functional language, so trying to use higher-order functions all the time can sometimes make your code harder to read (See also Principle of least astonishment), not to mention that you are using a Java API, which isn't designed for something like this at all.
One way I've thought of though, is:
fun fetch (url: String) =
(URL(url).openConnection() as HttpURLConnection).apply {
runCatching {
setRequestProperty(authorization.header, authorization.basicauth())
inputStream
.buffered()
.use { BitmapFactory.decodeStream(it) }
.asImageBitmap()
}.also { disconnect() }.getOrThrow()
}
I do the entire operation that could through inside a runCatching to catch any exceptions, disconnect from the connection, then throw the exception back out again.
In terms of the order of execution, this should be the same as try...finally like this:
fun fetch (url: String): ImageBitmap {
val connection = URL(url).openConnection() as HttpURLConnection
return try {
connection
.also { it.setRequestProperty(authorization.header, authorization.basicauth()) }
.getInputStream()
.buffered()
.use { BitmapFactory.decodeStream(it) }
.asImageBitmap()
} finally {
connection.disconnect()
}
}

How do I properly use Kotlin Flow in Ktor streaming responses?

emphasized textI am trying to use Kotlin Flow to process some data asynchronously and in parallel, and stream the responses to the client as they occur, as opposed to waiting until all the jobs are complete.
After unsuccessfully trying to just send the flow itself to the response, like this: call.respond(HttpStatusCode.OK, flow.toList())
... I tinkered for hours trying to figure it out, and came up with the following. Is this correct? It seems there should be a more idiomatic way of sending a Flow<MyData> as a response, like one can with a Flux<MyData> in Spring Boot.
Also, it seems that using the below method does not cancel the Flow when the HTTP request is cancelled, so how would one cancel it in Ktor?
data class MyData(val number: Int)
class MyService {
fun updateAllJobs(): Flow<MyData> =
flow {
buildList { repeat(10) { add(MyData(Random.nextInt())) } }
// Docs recommend using `onEach` to "delay" elements.
// However, if I delay here instead of in `map`, all elements are held
// and emitted at once at the very end of the cumulative delay.
// .onEach { delay(500) }
.map {
// I want to emit elements in a "stream" as each is computed.
delay(500)
emit(it)
}
}
}
fun Route.jobRouter() {
val service: MyService by inject() // injected with Koin
put("/jobs") {
val flow = service.updateAllJobs()
// Just using the default Jackson mapper for this example.
val mapper = jsonMapper { }
// `respondOutputStream` seems to be the only way to send a Flow as a stream.
call.respondOutputStream(ContentType.Application.Json, HttpStatusCode.OK) {
flow.collect {
println(it)
// The data does not stream without the newline and `flush()` call.
write((mapper.writeValueAsString(it) + "\n").toByteArray())
flush()
}
}
}
}
The best solution I was able to find (although I don't like it) is to use respondBytesWriter to write data to a response body channel. In the handler, a new job to collect the flow is launched to be able to cancel it if the channel is closed for writing (HTTP request is canceled):
fun Route.jobRouter(service: MyService) {
put("/jobs") {
val flow = service.updateAllJobs()
val mapper = jsonMapper {}
call.respondBytesWriter(contentType = ContentType.Application.Json) {
val job = launch {
flow.collect {
println(it)
try {
writeStringUtf8(mapper.writeValueAsString(it))
flush()
} catch (_: ChannelWriteException) {
cancel()
}
}
}
job.join()
}
}
}

Any way to detect client disconnet in Ktor server

I'm trying to make a proxy server by ktor, it proxy a http-flv infinite stream, after client close, it should do something to release resource. like
client.get<HttpStatement>(url.value.url).execute { response ->
val channel: ByteReadChannel = response.receive()
while (!channel.isClosedForRead) {
logger.info("status ${call.request.receiveChannel().isClosedForRead}")
val packet = channel.readRemaining(DEFAULT_BUFFER_SIZE.toLong())
logger.info("flushing...")
withContext(Dispatchers.IO) {
write(packet.readBytes())
flush()
}
}
// do something to clean resource
}
try to use router event and flush writer, didn't work

Is there a more idiomatic way to perform a subscribe & async / await operation?

I have a spring boot kotlin app that creates a web socket connection to another spring app, sends multiple "subscribe" messages, and then needs to wait for receipt of one response per subscription on the web socket connection. The number of subscriptions open at a given time could be up to a few thousand.
I've come up with a basic working solution using CompletableFuture and coroutines, as below. Is there a more idiomatic or concise way to do this task, or is this a fine solution? Any suggestions for improvement are appreciated.
// InputObject / ResponseObject are generic placeholders
fun getItems(inputObjects: List<InputObject>): List<ResponseObject> {
val ret: ConcurrentLinkedQueue<ResponseObject> = ConcurrentLinkedQueue()
// create a completable future for each input object
val subscriptions: MutableMap<String, CompletableFuture<ResponseObject>> = mutableMapOf()
inputObjects.forEach {
subscriptions[it.id] = CompletableFuture()
}
// create web socket client configured with a lambda handler to
// fulfill each subscription
// each responseObject.id matches one inputObject.id
val client = createWebSocketClient({
try {
val responseObject = objectMapper.readValue(it, ResponseObject::class.java)
subscriptions[responseObject.id]?.complete(responseObject)
} catch (e: Exception) {
logger.warn("Exception reading data: ${e.message}")
}
})
runBlocking {
coroutineScope {
for (item in inputObjects) {
launch {
// create and send a subscribe request
client.sendMessage(createSubscribe(item.id))
// wait for each future to complete
// uses CompletableFuture extension await() from kotlinx-coroutines-jdk8
val result = subscriptions[item.id]?.await()
if (result != null) {
ret.add(result)
}
}
}
}
}
client.close()
return ret.toList()
}
edit: I found a similar question: How to pass result as it comes using coroutines?
Which options makes the most sense?
fun getItems(inputObjects: List<InputObject>): List<ResponseObject> {
val subscriptions = ids.associateTo(mutableMapOf()) { it.id to CompletableFuture<ResponseObject>() }
val client = createWebSocketClient({
try {
val responseObject = objectMapper.readValue(it, ResponseObject::class.java)
subscriptions[responseObject.id]?.complete(responseObject)
} catch (e: Exception) {
logger.warn("Exception reading data: ${e.message}")
}
})
return runBlocking(Dispatchers.IO) {
inputObjects
.mapNotNull {
client.sendMessage(createSubscribe(item.id))
subscriptions[item.id]?.await()
}
}
}

Kotlin Coroutine - Ktor Server WebSocket

I made a kotlin-ktor application, what i wanted to achieve is that it is modular that anytime any pipelines inside the application maybe removed from the source code without breaking any functionalities. So i decided i want to move a websocket implementation to separate class
but i faced an issue where the coroutine inside the lambda expression terminates immediately.
link-github issue
Can someone enlighten me about the coroutine setup on this, and how I can still keep this as modular without this kind of issue
working ktor websocket
fun Application.socketModule() = runBlocking {
// other declarations
......
routing {
val sessionService = SocketSessionService()
webSocket("/home") {
val chatSession = call.sessions.get<ChatSession>()
println("request session: $chatSession")
if (chatSession == null) {
close(CloseReason(CloseReason.Codes.VIOLATED_POLICY, "empty Session"))
return#webSocket
}
send(Frame.Text("connected to server"))
sessionService.addLiveSocket(chatSession.id, this)
sessionService.checkLiveSocket()
}
thread(start = true, name = "socket-monitor") {
launch {
sessionService.checkLiveSocket()
}
}
}
}
kotlin-ktor auto-close web socket
code below closes the socket automatically
Socket Module
class WebSocketServer {
fun createWebSocket(root: String, routing: Routing) {
println("creating web socket server")
routing.installSocketRoute(root)
}
private fun Routing.installSocketRoute(root: String) {
val base = "/message/so"
val socketsWeb = SocketSessionService()
webSocket("$root$base/{type}") {
call.parameters["type"] ?: throw Exception("missing type")
val session = call.sessions.get<ChatSession>()
if (session == null) {
println( "WEB-SOCKET:: client session is null" )
close(CloseReason(CloseReason.Codes.VIOLATED_POLICY, "No Session"))
return#webSocket
}
socketsWeb.addLiveSocket(session.id, this)
thread(start= true, name = "thread-live-socket") {
launch {
socketsWeb.checkLiveSocket()
}
}
}
}
}
Application Module
fun Application.socketModule() = runBlocking {
// other delcarations
.....
install(Sessions) {
cookie<ChatSession>("SESSION")
}
intercept(ApplicationCallPipeline.Features) {
if (call.sessions.get<ChatSession>() == null) {
val sessionID = generateNonce()
println("generated Session: $sessionID")
call.sessions.set(ChatSession(sessionID))
}
}
routing {
webSocketServer.createWebSocket("/home", this)
}
}
I quite dont understand why the coroutine insdie webSocket lamda is completed.
Can someone show me other/right approach on this one.