From e4037549b4f3258c3b2df9cc253e3a409718354a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=AD=B1=E5=82=91?= <jie65535@qq.com>
Date: Thu, 2 Mar 2023 23:42:13 +0800
Subject: [PATCH] Implement WebSocket client

---
 .idea/gradle.xml                              |  1 +
 build.gradle                                  |  2 +
 .../java/top/jie65535/minionebot/Main.java    | 15 ++--
 .../top/jie65535/minionebot/MiniOneBot.java   | 69 ++++++++++++-----
 .../minionebot/MiniOneBotWsClient.java        | 74 +++++++++++++++----
 .../minionebot/MiniOneBotWsServer.java        | 54 ++++++++++----
 .../top/jie65535/minionebot/WsStream.java     | 11 +++
 7 files changed, 174 insertions(+), 52 deletions(-)
 create mode 100644 src/main/java/top/jie65535/minionebot/WsStream.java

diff --git a/.idea/gradle.xml b/.idea/gradle.xml
index ba1ec5c..611e7c8 100644
--- a/.idea/gradle.xml
+++ b/.idea/gradle.xml
@@ -1,5 +1,6 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <project version="4">
+  <component name="GradleMigrationSettings" migrationVersion="1" />
   <component name="GradleSettings">
     <option name="linkedExternalProjectsSettings">
       <GradleProjectSettings>
diff --git a/build.gradle b/build.gradle
index 072457e..45d3066 100644
--- a/build.gradle
+++ b/build.gradle
@@ -15,7 +15,9 @@ dependencies {
     testImplementation 'org.junit.jupiter:junit-jupiter-api:5.8.1'
     testRuntimeOnly 'org.junit.jupiter:junit-jupiter-engine:5.8.1'
 
+    implementation 'org.slf4j:slf4j-simple:2.0.5'
     implementation 'io.javalin:javalin:5.3.2'
+    implementation 'org.java-websocket:Java-WebSocket:1.5.3'
     implementation 'com.google.code.gson:gson:2.10.1'
 }
 
diff --git a/src/main/java/top/jie65535/minionebot/Main.java b/src/main/java/top/jie65535/minionebot/Main.java
index 7eb5da8..e4bb3e3 100644
--- a/src/main/java/top/jie65535/minionebot/Main.java
+++ b/src/main/java/top/jie65535/minionebot/Main.java
@@ -1,11 +1,9 @@
 package top.jie65535.minionebot;
 
-import com.google.gson.Gson;
-import com.google.gson.JsonObject;
 import io.javalin.Javalin;
 
-import java.util.Map;
-import java.util.Objects;
+import java.net.URI;
+import java.net.URISyntaxException;
 import java.util.Scanner;
 
 public class Main {
@@ -28,10 +26,15 @@ public class Main {
 //        System.out.println(Objects.equals(dataMap.get("type"), "text"));
 //        System.out.println((String)dataMap.get("data"));
 
+
         Javalin app = Javalin.create().start(7070);
-        var bot = new MiniOneBot(app);
+        var bot = new MiniOneBot(app, "HelloOneBot");
         bot.startWsServer("/openchat");
-        bot.startWsClient("ws://localhost/onebot");
+        try {
+            bot.startWsClient(new URI("ws://localhost:8080"));
+        } catch (URISyntaxException ignored) {
+            System.out.println("Uri Syntax error!");
+        }
         bot.subscribeGroupMessageEvent(event -> System.out.println(event.message));
         bot.subscribeGuildChannelMessageEvent(event -> System.out.println(event.message));
 
diff --git a/src/main/java/top/jie65535/minionebot/MiniOneBot.java b/src/main/java/top/jie65535/minionebot/MiniOneBot.java
index 79e77ef..5941e6a 100644
--- a/src/main/java/top/jie65535/minionebot/MiniOneBot.java
+++ b/src/main/java/top/jie65535/minionebot/MiniOneBot.java
@@ -6,51 +6,84 @@ import org.slf4j.LoggerFactory;
 import top.jie65535.minionebot.events.GroupMessageHandler;
 import top.jie65535.minionebot.events.GuildChannelMessageHandler;
 
-public class MiniOneBot {
+import java.net.URI;
+
+public class MiniOneBot implements WsStream.WsMessageHandler {
     private final Logger logger = LoggerFactory.getLogger("Console");
     private final Javalin javalin;
+    private final String token;
     private MiniOneBotWsServer server;
     private MiniOneBotWsClient client;
-    public MiniOneBot(Javalin javalin) {
+
+    public MiniOneBot(Javalin javalin, String token) {
         this.javalin = javalin;
+        this.token = token;
         logger.info("Test");
     }
 
+    // region WebSocket
+
     public void startWsServer(String path) {
         if (server == null) {
-            server = new MiniOneBotWsServer(javalin, path, logger);
+            logger.info("Start MiniOneBot WebSocket Server");
+            server = new MiniOneBotWsServer(javalin, path, token, logger);
+            server.subscribe(this);
         }
     }
 
-    public void startWsClient(String wsUrl) {
+    public void startWsClient(URI serverUri) {
         if (client == null) {
-            client = new MiniOneBotWsClient(wsUrl, logger);
+            logger.info("Start MiniOneBot WebSocket Client");
+            client = MiniOneBotWsClient.create(serverUri, token, logger);
+            client.subscribe(this);
         }
     }
 
-    private static String escape(String msg){
+    private void sendMessageToAll(String message) {
+        server.send(message);
+        client.send(message);
+    }
+
+    @Override
+    public void onMessage(String message) {
+        // TODO
+    }
+
+    // endregion
+
+    // region Message API
+
+    public void sendMessage(String sender, String message) {
+        // TODO
+    }
+
+    GroupMessageHandler groupMessageHandler;
+    public void subscribeGroupMessageEvent(GroupMessageHandler handler) {
+        groupMessageHandler = handler;
+    }
+
+    GuildChannelMessageHandler guildChannelMessageHandler;
+    public void subscribeGuildChannelMessageEvent(GuildChannelMessageHandler handler) {
+        guildChannelMessageHandler = handler;
+    }
+
+    // endregion
+
+    // region Utils
+
+    private static String escape(String msg) {
         return msg.replace("&", "&amp;")
                 .replace("[", "&#91;")
                 .replace("]", "&#93;")
                 .replace(",", "&#44;");
     }
 
-    private static String unescape(String msg){
+    private static String unescape(String msg) {
         return msg.replace("&amp;", "&")
                 .replace("&#91;", "[")
                 .replace("&#93;", "]")
                 .replace("&#44;", ",");
     }
 
-    public void sendMessage(String sender, String message) {
-        
-    }
-
-    public void subscribeGroupMessageEvent(GroupMessageHandler handler) {
-
-    }
-
-    public void subscribeGuildChannelMessageEvent(GuildChannelMessageHandler handler) {
-
-    }
+    // endregion
 }
diff --git a/src/main/java/top/jie65535/minionebot/MiniOneBotWsClient.java b/src/main/java/top/jie65535/minionebot/MiniOneBotWsClient.java
index 5155e54..b53879f 100644
--- a/src/main/java/top/jie65535/minionebot/MiniOneBotWsClient.java
+++ b/src/main/java/top/jie65535/minionebot/MiniOneBotWsClient.java
@@ -1,29 +1,73 @@
 package top.jie65535.minionebot;
 
-import org.jetbrains.annotations.NotNull;
+import org.java_websocket.client.WebSocketClient;
+import org.java_websocket.handshake.ServerHandshake;
 import org.slf4j.Logger;
 
+import java.net.URI;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.Timer;
 import java.util.TimerTask;
 
-public class MiniOneBotWsClient {
-
-    private final String wsUrl;
+public class MiniOneBotWsClient extends WebSocketClient implements WsStream {
     private final Logger logger;
 
-    private final Timer wsConnectDaemon;
+    private MiniOneBotWsClient(URI serverUri, Map<String, String> headers, Logger logger) {
+        super(serverUri, headers);
 
-    public MiniOneBotWsClient(@NotNull String wsUrl, @NotNull Logger logger) {
-        this.wsUrl = wsUrl;
         this.logger = logger;
-
-        wsConnectDaemon = new Timer("WsClientDaemon", true);
-        wsConnectDaemon.schedule(new TimerTask() {
-            @Override
-            public void run() {
-                logger.debug("Timer...");
-            }
-        }, 60_000);
     }
 
+    public static MiniOneBotWsClient create(URI serverUri, String token, Logger logger) {
+        var headers = new HashMap<String, String>();
+        headers.put("Authorization", "Bearer " + token);
+        var client = new MiniOneBotWsClient(serverUri, headers, logger);
+        var wsClientDaemon = new Timer("WsClientDaemon", true);
+        wsClientDaemon.schedule(new TimerTask() {
+            @Override
+            public void run() {
+                if (!client.isOpen()) {
+                    logger.debug("Try connect...");
+                    client.connect();
+                }
+            }
+        }, 5_000);
+        return client;
+    }
+
+    private WsMessageHandler callback;
+
+    @Override
+    public void subscribe(WsMessageHandler callback) {
+        this.callback = callback;
+    }
+
+    @Override
+    public void onOpen(ServerHandshake handshakedata) {
+        logger.info("onOpen: statusMessage={}", handshakedata.getHttpStatusMessage());
+    }
+
+    @Override
+    public void onMessage(String message) {
+        logger.info("onMessage: {}", message);
+        callback.onMessage(message);
+    }
+
+    @Override
+    public void onClose(int code, String reason, boolean remote) {
+        logger.info("onClose: code={} reason={} isRemote={}", code, reason, remote);
+    }
+
+    @Override
+    public void onError(Exception ex) {
+        logger.error("onError:", ex);
+    }
+
+    @Override
+    public void send(String message) {
+        if (isOpen()) {
+            super.send(message);
+        }
+    }
 }
diff --git a/src/main/java/top/jie65535/minionebot/MiniOneBotWsServer.java b/src/main/java/top/jie65535/minionebot/MiniOneBotWsServer.java
index d7163b6..4187220 100644
--- a/src/main/java/top/jie65535/minionebot/MiniOneBotWsServer.java
+++ b/src/main/java/top/jie65535/minionebot/MiniOneBotWsServer.java
@@ -2,18 +2,20 @@ package top.jie65535.minionebot;
 
 import io.javalin.Javalin;
 import io.javalin.websocket.*;
-import org.jetbrains.annotations.NotNull;
+import org.eclipse.jetty.websocket.api.CloseStatus;
 import org.slf4j.Logger;
 
 import java.util.Map;
-import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
 
-public class MiniOneBotWsServer {
+public class MiniOneBotWsServer implements WsStream {
 
+    private final String token;
     private final Logger logger;
     private final Map<WsContext, String> connections = new ConcurrentHashMap<>();
-    public MiniOneBotWsServer(@NotNull Javalin javalin, @NotNull String path, @NotNull Logger logger) {
+
+    public MiniOneBotWsServer(Javalin javalin, String path, String token, Logger logger) {
+        this.token = token;
         this.logger = logger;
         javalin.ws(path, ws -> {
             ws.onConnect(this::onConnect);
@@ -25,18 +27,26 @@ public class MiniOneBotWsServer {
         logger.info("WebSocket server started at {}", path);
     }
 
-    public void onConnect(@NotNull WsConnectContext ctx) {
+    public void onConnect(WsConnectContext ctx) {
         logger.debug("onConnect: address={} headers={}", ctx.session.getRemoteAddress(), ctx.headerMap());
-        var selfId = ctx.header("X-Self-ID");
-        if (selfId != null && !selfId.isEmpty()) {
-            logger.info("Bot [{}] WebSocket connected", selfId);
+        var author = ctx.header("Authorization");
+        // Check access token.
+        if (author == null) {
+            ctx.session.close(new CloseStatus(401, "Unauthorized"));
+        } else if (!author.equals("Bearer " + token)) {
+            ctx.session.close(new CloseStatus(403, "Unauthorized"));
         } else {
-            logger.info("[{}] WebSocket connected", ctx.session.getRemoteAddress());
+            var selfId = ctx.header("X-Self-ID");
+            if (selfId != null && !selfId.isEmpty()) {
+                logger.info("Bot [{}] WebSocket connected", selfId);
+            } else {
+                logger.info("[{}] WebSocket connected", ctx.session.getRemoteAddress());
+            }
+            connections.put(ctx, selfId);
         }
-        connections.put(ctx, selfId);
     }
 
-    public void onClose(@NotNull WsCloseContext ctx) {
+    public void onClose(WsCloseContext ctx) {
         logger.debug("onClose: address={} status={} reason={}", ctx.session.getRemoteAddress(), ctx.status(), ctx.reason());
         var selfId = connections.remove(ctx);
         if (selfId != null && !selfId.isEmpty()) {
@@ -46,7 +56,7 @@ public class MiniOneBotWsServer {
         }
     }
 
-    public void onError(@NotNull WsErrorContext ctx) {
+    public void onError(WsErrorContext ctx) {
         logger.warn("onError: address={}", ctx.session.getRemoteAddress(), ctx.error());
         var selfId = connections.remove(ctx);
         if (selfId != null && !selfId.isEmpty()) {
@@ -56,9 +66,10 @@ public class MiniOneBotWsServer {
         }
     }
 
-    public void onMessage(@NotNull WsMessageContext ctx) {
+    public void onMessage(WsMessageContext ctx) {
         logger.debug("onMessage: {}", ctx.message());
 
+        callback.onMessage(ctx.message());
 //        var map = JsonUtils.decode(ctx.message(), Map.class);
 //        // 消息事件上报
 //        if (Objects.equals(map.get("post_type"), "message")) {
@@ -85,4 +96,21 @@ public class MiniOneBotWsServer {
 //            }
 //        }
     }
+
+    private WsMessageHandler callback;
+
+    @Override
+    public void subscribe(WsMessageHandler callback) {
+        this.callback = callback;
+    }
+
+    @Override
+    public void send(String message) {
+        if (connections.isEmpty()) return;
+        for (var ctx : connections.keySet()) {
+            if (ctx.session.isOpen()) {
+                ctx.send(message);
+            }
+        }
+    }
 }
diff --git a/src/main/java/top/jie65535/minionebot/WsStream.java b/src/main/java/top/jie65535/minionebot/WsStream.java
new file mode 100644
index 0000000..242cac5
--- /dev/null
+++ b/src/main/java/top/jie65535/minionebot/WsStream.java
@@ -0,0 +1,11 @@
+package top.jie65535.minionebot;
+
+public interface WsStream {
+    void subscribe(WsMessageHandler callback);
+
+    void send(String message);
+
+    interface WsMessageHandler {
+        void onMessage(String message);
+    }
+}