diff --git a/build.gradle b/build.gradle index 14fc335..ee87cf1 100644 --- a/build.gradle +++ b/build.gradle @@ -10,12 +10,15 @@ sourceCompatibility = 1.8 repositories { mavenCentral() + google() + jcenter() } dependencies { compile "org.jetbrains.kotlin:kotlin-stdlib-jdk8" compile 'io.javalin:javalin:2.6.0' compile 'org.slf4j:slf4j-simple:1.7.25' + implementation 'com.beust:klaxon:5.0.1' testCompile group: 'junit', name: 'junit', version: '4.12' } diff --git a/src/main/kotlin/com/zecqtwallet/wormhole/Service.kt b/src/main/kotlin/com/zecqtwallet/wormhole/Service.kt index fa43fa8..e1d0033 100644 --- a/src/main/kotlin/com/zecqtwallet/wormhole/Service.kt +++ b/src/main/kotlin/com/zecqtwallet/wormhole/Service.kt @@ -1,6 +1,18 @@ package com.zecqtwallet.wormhole +import com.beust.klaxon.JsonObject +import com.beust.klaxon.Parser +import com.beust.klaxon.json import io.javalin.Javalin +import io.javalin.websocket.WsSession +import java.util.concurrent.ConcurrentHashMap + + +private val usermap = ConcurrentHashMap() + +// Allow maps to be bidirectional +fun Map.getKeys(value: V) : List = + entries.filter { it.value == value } .map { it.key } fun main(args : Array) { @@ -10,14 +22,62 @@ fun main(args : Array) { ws.onConnect { session -> println("Connected Session") } + ws.onClose { session, status, message -> println("Closed session") + usermap.remove(session) } + ws.onMessage { session, message -> - println("message") - session.send("Reply to $message") + // Limit message size to 50kb of hex encoded text + if (message.length > 2 * 50 * 1024) { + sendError(session, "Message too big") + } + + // Parse the message as json + try { + val j = Parser.default().parse(StringBuilder(message)) as JsonObject + + if (j.contains("register")) { + doRegister(session, j["register"].toString()) + return@onMessage + } + + if (j.contains("to")) { + val s = usermap.getKeys(j["to"].toString()).filter { it != session } + if (s.isEmpty()) { + // Not connected + sendError(session, "Peer is not connected") + return@onMessage + } + + s[0].send(message) + return@onMessage + } + + } catch (e: Throwable) { + session.close(1000, "Invalid json") + } } } }.start(7070) } + +fun doRegister(session: WsSession, id: String) { + if (usermap.contains(session)) { + // TODO: Make JSON + sendError(session, "Already registered a session") + + usermap.remove(session) + session.close() + } + + usermap[session] = id +} + +fun sendError(session: WsSession, err: String) { + if (session.isOpen) { + session.send(json { obj("error" to err) }.toJsonString()) + } +} \ No newline at end of file