From e39446cd4743aa4455b2c19dea0e96d6a7dfc378 Mon Sep 17 00:00:00 2001
From: yenon <bhogm4@gmail.com>
Date: Thu, 7 Dec 2023 19:09:58 +0100
Subject: [PATCH] New connection logic, a little more intuitive.

---
 src/main/kotlin/yenon/squadcompanion/Main.kt  |  9 +++-
 .../yenon/squadcompanion/RconConnection.kt    | 42 ++++++++++++-------
 2 files changed, 33 insertions(+), 18 deletions(-)

diff --git a/src/main/kotlin/yenon/squadcompanion/Main.kt b/src/main/kotlin/yenon/squadcompanion/Main.kt
index b1638a1..e469b40 100644
--- a/src/main/kotlin/yenon/squadcompanion/Main.kt
+++ b/src/main/kotlin/yenon/squadcompanion/Main.kt
@@ -9,6 +9,7 @@ import io.ktor.server.plugins.statuspages.*
 import io.ktor.server.request.*
 import io.ktor.server.response.*
 import io.ktor.server.routing.*
+import kotlinx.coroutines.runBlocking
 import yenon.squadcompanion.logreader.LogReaderTail
 import java.io.PrintWriter
 import java.io.StringWriter
@@ -37,8 +38,12 @@ object Main {
             }
         }
 
-        val connection = RconConnection("195.201.62.238", 21114, "ThisIsJustATestServer")
-        val abstraction = RconAbstraction(connection)
+        val abstraction = runBlocking {
+            val connection = RconConnection("195.201.62.238", 21114, "ThisIsJustATestServer")
+            connection.connect()
+            RconAbstraction(connection)
+        }
+
 
         embeddedServer(Netty, 8080) {
             install(StatusPages) {
diff --git a/src/main/kotlin/yenon/squadcompanion/RconConnection.kt b/src/main/kotlin/yenon/squadcompanion/RconConnection.kt
index db2497a..1856c55 100644
--- a/src/main/kotlin/yenon/squadcompanion/RconConnection.kt
+++ b/src/main/kotlin/yenon/squadcompanion/RconConnection.kt
@@ -1,6 +1,7 @@
 package yenon.squadcompanion
 
 import io.ktor.utils.io.core.*
+import io.ktor.utils.io.errors.*
 import io.ktor.utils.io.nio.*
 import kotlinx.coroutines.*
 import kotlinx.coroutines.flow.*
@@ -60,7 +61,12 @@ class RconConnection(private val host: String, private val port: Int, password:
     private val sameIdPacketList = LinkedList<RconPacket>()
     private val packetsFlow = MutableSharedFlow<RconPacket>(extraBufferCapacity = 8)
 
-    init {
+    suspend fun connect(): ConnectionStatus {
+        val status = reconnect()
+        if (status != ConnectionStatus.CONNECTED) {
+            return status
+        }
+
         coroutineScope.launch {
             while (isActive) {
                 delay(5000)
@@ -108,17 +114,30 @@ class RconConnection(private val host: String, private val port: Int, password:
                 sameIdPacketList.clear()
             }
         }
+        return ConnectionStatus.CONNECTED
     }
 
-    private suspend fun reconnect(): Boolean = withContext(Dispatchers.IO) {
+    enum class ConnectionStatus {
+        CONNECTED, WRONG_PASSWORD, TIMEOUT
+    }
+
+    private suspend fun reconnect(): ConnectionStatus = withContext(Dispatchers.IO) {
         socketChannel.close()
-        socketChannel = SocketChannel.open().apply {
-            connect(InetSocketAddress(host, port))
-            finishConnect()
-            println("Socket is open.")
+        try {
+            socketChannel = SocketChannel.open().apply {
+                connect(InetSocketAddress(host, port))
+                finishConnect()
+                println("Socket is open.")
+            }
+        } catch (ex: IOException) {
+            return@withContext ConnectionStatus.TIMEOUT
         }
 
-        authenticate()
+        return@withContext if (authenticate()) {
+            ConnectionStatus.CONNECTED
+        } else {
+            ConnectionStatus.WRONG_PASSWORD
+        }
     }
 
     private suspend fun authenticate(): Boolean {
@@ -163,15 +182,6 @@ class RconConnection(private val host: String, private val port: Int, password:
         }
     }
 
-    private fun readPacket(): RconPacket {
-        var packet = readRawPacket()
-        while (packet.id == 0 && packet.type == 1) {
-            println("Console ${packet.data.size}: " + String(packet.data))
-            packet = readRawPacket()
-        }
-        return packet
-    }
-
     private fun readRawPacket(): RconPacket {
         print("Reading new packet")
         val size = socketChannel.readPacketExact(4L).readIntLittleEndian()