websocket

module/Local_Server/server/WebSocket_Manager.cpp

@@ -51,6 +51,24 @@ void WebSocket_Manager::push_data(const char *buf, uint64_t len) {
     conn->send(buf, len, drogon::WebSocketMessageType::Binary);
   }
 }
+void WebSocket_Manager::set_message_handler(Message_Handler handler) {
+  message_handler = std::move(handler);
+}
+void WebSocket_Manager::set_close_handler(Close_Handler handler) {
+  close_handler = std::move(handler);
+}
+void WebSocket_Manager::handle_message(const drogon::WebSocketConnectionPtr &conn,
+                                       std::string &&message,
+                                       const drogon::WebSocketMessageType &type) {
+  if (message_handler) {
+    message_handler(conn, std::move(message), type);
+  }
+}
+void WebSocket_Manager::handle_close(const drogon::WebSocketConnectionPtr &conn) {
+  if (close_handler) {
+    close_handler(conn);
+  }
+}
 size_t WebSocket_Manager::get_websocket_connect_num() {
   std::lock_guard<std::mutex> g(mtx);
   return connections.size();
@@ -77,14 +95,24 @@ std::shared_ptr<WebSocket_Manager> Ws_Mgr::get_ws(std::string_view path) {
   return nullptr;
 }
 
-void DSP_Ws_Mgr::handleNewMessage(const drogon::WebSocketConnectionPtr &,
-                                  std::string &&,
-                                  const drogon::WebSocketMessageType &) {}
+void DSP_Ws_Mgr::handleNewMessage(const drogon::WebSocketConnectionPtr &conn,
+                                  std::string &&message,
+                                  const drogon::WebSocketMessageType &type) {
+  auto &s = conn->getContextRef<WebSocket_Connection>();
+  auto ws = Ws_Mgr::instance()->get_ws(s.type);
+  if (ws) {
+    ws->handle_message(conn, std::move(message), type);
+  }
+}
 void DSP_Ws_Mgr::handleConnectionClosed(
     const drogon::WebSocketConnectionPtr &conn) {
   auto &s = conn->getContextRef<WebSocket_Connection>();
   auto rp = s.type;
   auto ws = Ws_Mgr::instance()->get_ws(rp);
+  if (!ws) {
+    return;
+  }
+  ws->handle_close(conn);
   // LOG_DEBUG
   std::ostringstream oss;
   oss << "已有连接数:" << ws->get_websocket_connect_num() << "  "
@@ -100,6 +128,10 @@ void DSP_Ws_Mgr::handleNewConnection(
   auto rpt = req->path();
   auto rp = rpt.substr(std::string("/ws").size());
   auto ws = Ws_Mgr::instance()->get_ws(rp);
+  if (!ws) {
+    conn->shutdown();
+    return;
+  }
   std::ostringstream oss;
   oss << "已有连接数:" << ws->get_websocket_connect_num() << "  "
       << "新建websocket连接:" << rp << "  "

module/Local_Server/server/WebSocket_Manager.h

@@ -11,17 +11,25 @@
 
 class WebSocket_Manager {
 public:
+  using Message_Handler = std::function<void(const drogon::WebSocketConnectionPtr &conn, std::string &&message, const drogon::WebSocketMessageType &type)>;
+  using Close_Handler = std::function<void(const drogon::WebSocketConnectionPtr &conn)>;
   explicit WebSocket_Manager(std::string_view path);
   ~WebSocket_Manager() = default;
   void remove(const drogon::WebSocketConnectionPtr &conn);
   void add(const drogon::WebSocketConnectionPtr &conn);
   void visit(const std::function<void(const drogon::WebSocketConnectionPtr &conn)> &func);
   void push_data(const char *buf, uint64_t len);
+  void set_message_handler(Message_Handler handler);
+  void set_close_handler(Close_Handler handler);
+  void handle_message(const drogon::WebSocketConnectionPtr &conn, std::string &&message, const drogon::WebSocketMessageType &type);
+  void handle_close(const drogon::WebSocketConnectionPtr &conn);
   size_t get_websocket_connect_num();
   std::string path;
 protected:
   std::mutex                                 mtx;
   std::unordered_set<drogon::WebSocketConnectionPtr> connections;
+  Message_Handler message_handler;
+  Close_Handler close_handler;
 };