From e06f3891c8b7cdf67473c85f05e5af39452d1808 Mon Sep 17 00:00:00 2001
From: Maxwell Jeffress <max@maxwellj.xyz>
Date: Mon, 28 Oct 2024 20:37:20 +1100
Subject: [PATCH] Anti timeout ping, user list

---
 server/build.gradle.kts        |  1 +
 server/src/main/kotlin/Main.kt | 86 +++++++++++++++++++++++++++++-----
 2 files changed, 74 insertions(+), 13 deletions(-)

diff --git a/server/build.gradle.kts b/server/build.gradle.kts
index d78af3f..a2c18b0 100644
--- a/server/build.gradle.kts
+++ b/server/build.gradle.kts
@@ -24,6 +24,7 @@ tasks.withType<Jar> {
 
 dependencies {
     testImplementation(kotlin("test"))
+    implementation("org.jetbrains.kotlinx:kotlinx-coroutines-core:1.7.3")
     implementation("io.javalin:javalin:6.3.0")
     implementation("org.slf4j:slf4j-simple:2.0.16")
 }
diff --git a/server/src/main/kotlin/Main.kt b/server/src/main/kotlin/Main.kt
index f7d4b9b..56fc657 100644
--- a/server/src/main/kotlin/Main.kt
+++ b/server/src/main/kotlin/Main.kt
@@ -5,6 +5,8 @@ import io.javalin.websocket.WsContext
 import java.util.concurrent.ConcurrentHashMap
 import java.util.UUID
 
+import kotlin.concurrent.fixedRateTimer
+
 import java.io.File
 import java.io.BufferedReader
 
@@ -17,12 +19,50 @@ fun md5(input:String): String {
 }
 
 object WsSessionManager {
+    var peopleOnline = mutableListOf("")
+    var sessionsList = mutableListOf("")
+
     private val sessions = ConcurrentHashMap<WsContext, String>()
     private val sessionIds = ConcurrentHashMap<String, WsContext>()
 
+    init {
+        fixedRateTimer("websocket-ping", period = 5000) {
+            sendPing()
+        }
+    }
+    private fun sendPing() {
+        val deadSessions = mutableListOf<WsContext>()
+
+        sessions.keys.forEach { ctx ->
+            try {
+                if (ctx.session.isOpen) {
+                    ctx.send("ping")
+                } else {
+                    deadSessions.add(ctx)
+                }
+            } catch (e: Exception) {
+                println("Error sending ping: ${e.message}")
+                deadSessions.add(ctx)
+            }
+        }
+
+        // Clean up any dead sessions
+        deadSessions.forEach { removeSession(it) }
+    }
+
+    fun broadcastOnlineUsers() {
+        broadcast("!users:{${peopleOnline.joinToString(",")}}")
+    }
+
+    fun handleUserLogin(username: String) {
+        peopleOnline += username
+        broadcastOnlineUsers()
+    }
+
     fun addSession(ctx: WsContext) {
         try {
             val sessionId = UUID.randomUUID().toString()
+            sessionsList += sessionId
             sessions[ctx] = sessionId
             sessionIds[sessionId] = ctx
         } catch (e: Exception) {
@@ -34,8 +74,11 @@ object WsSessionManager {
         try {
             val sessionId = sessions[ctx]
             if (sessionId != null) {
+                peopleOnline.removeAt(sessionsList.indexOf(sessionId))
+                sessionsList.removeAt(sessionsList.indexOf(sessionId))
                 sessions.remove(ctx)
                 sessionIds.remove(sessionId)
+                broadcastOnlineUsers()
             }
         } catch (e: Exception) {
             println("Error removing session: ${e.message}")
@@ -101,6 +144,7 @@ fun handleSentMessage(inputData: String): String {
     var message = ""
     var dataType = ""
     var command = ""
+    var commandArg = ""
     var isParsingData = 0
     for (char in inputData) {
         val character = char
@@ -119,6 +163,8 @@ fun handleSentMessage(inputData: String): String {
                     message += character
                 } else if (dataType == "command") {
                     command += character
+                } else if (dataType == "commandArg") {
+                    commandArg += character
                 }
             }
         } else {
@@ -177,9 +223,17 @@ fun handleSentMessage(inputData: String): String {
         chatHistory.appendText("$username: $message ${System.lineSeparator()}")
         message = ""
         return("Success")
+    } else if (command != "") {
+        if (command == "sync") {
+            return(chatHistoryView.readText())
+        } else if (command == "login") {
+            WsSessionManager.handleUserLogin(commandArg)
+            return("Login successful")
+        }
     } else {
         return("No data provided")
     }
+    return("System: Welcome to Chookpen, $username!")
 }
 
 fun syncMessages(inputData: String): String {
@@ -391,7 +445,8 @@ fun authKey(inputData: String): String {
 }
 
 fun main(args: Array<String>) {
-    var users = listOf("")
+    WsSessionManager.peopleOnline.removeAt(0)
+    WsSessionManager.sessionsList.removeAt(0)
     val app = Javalin.create()
         .get("/") { ctx -> ctx.result("dingus") }
         .get("/api/send/{content}") { ctx -> 
@@ -413,20 +468,25 @@ fun main(args: Array<String>) {
             ws.onClose { ctx ->
                 WsSessionManager.removeSession(ctx)
             }
-            ws.onMessage { ctx -> 
-                println(ctx.message()) 
-                val successState = handleSentMessage(ctx.message())
-                if (successState != "Success") {
-                    try {
-                        ctx.send(successState)
-                    } catch (e: Exception) {
-                        println("Error sending error message: ${e.message}")
+            ws.onMessage { ctx ->
+                when (ctx.message()) {
+                    "pong" -> {}
+                else -> {
+                    println(ctx.message()) 
+                    val successState = handleSentMessage(ctx.message())
+                    if (successState != "Success") {
+                        try {
+                            ctx.send(successState)
+                        } catch (e: Exception) {
+                            println("Error sending error message: ${e.message}")
+                        }
+                    } else {
+                        val messageContent = extractMessageContent(ctx.message())
+                        WsSessionManager.broadcast(messageContent)
                     }
-                } else {
-                    val messageContent = extractMessageContent(ctx.message())
-                    WsSessionManager.broadcast(messageContent)
                 }
-            } 
+            }
+            }
         }
     .start(7070)