pocoproject / poco

The POCO C++ Libraries are powerful cross-platform C++ libraries for building network- and internet-based applications that run on desktop, server, mobile, IoT, and embedded systems.
https://pocoproject.org
Other
8.04k stars 2.11k forks source link

Causing websocket to disconnect and the client not receiving any messages #4560

Closed KanLai closed 1 month ago

KanLai commented 1 month ago
#pragma once

#include <Poco/Net/WebSocket.h>
#include <Poco/Net/HTTPRequestHandler.h>
#include <Poco/Net/HTTPRequestHandlerFactory.h>
#include <Poco/Util/ServerApplication.h>
#include <Poco/Net/HTTPServerRequest.h>
#include <Poco/Net/HTTPServerResponse.h>
#include <Poco/Net/HTTPServer.h>
#include <Poco/Net/HTTPResponse.h>
#include <Poco/Net/HTTPServerParams.h>
#include <Poco/FunctionDelegate.h>
#include <Poco/ThreadPool.h>
#include <Poco/Net/ServerSocket.h>
#include <Poco/RunnableAdapter.h>
#include <Poco/TaskManager.h>
#include <Poco/UUID.h>
#include <Poco/UUIDGenerator.h>
#include <set>
#include <mutex>
#include <iostream>
#include <utility>
#include <unordered_map>
#include <functional>

namespace WebSocket {

    class WebSocketManager {
    public:
        static void receive(Poco::Net::WebSocket &ws, const std::string &message) {
            try {
                ws.sendFrame("123",3, Poco::Net::WebSocket::FRAME_OP_TEXT);
            } catch (const Poco::Exception &exc) {
                std::cerr << "WebSocketConnection send error: " << exc.displayText() << std::endl;
            }
        }
    };

    class AsyncRequestTask {
    public:
        explicit AsyncRequestTask(Poco::Net::WebSocket& ws)
                : ws(ws) {}

        void runTask() {
            try {
                char buffer[1024];
                int flags;
                int n;
                do {
                    n = ws.receiveFrame(buffer, sizeof(buffer), flags);
                    if (n > 0) {
                        std::string message(buffer, static_cast<unsigned long>(n));
                        if (flags & Poco::Net::WebSocket::FRAME_OP_BITMASK) {
                            WebSocketManager::receive(ws, message);// Just send the frame directly, there's no problem
                        }
                    }
                } while (n > 0 && (flags & Poco::Net::WebSocket::FRAME_OP_BITMASK) != Poco::Net::WebSocket::FRAME_OP_CLOSE);
            } catch (const Poco::Exception &exc) {
                std::cerr << "WebSocketTask exception: " << exc.displayText() << std::endl;
            }
        }

    private:
        Poco::Net::WebSocket& ws;
    };

    class WebSocketRequestHandler : public Poco::Net::HTTPRequestHandler {
    public:
        explicit WebSocketRequestHandler(Poco::TaskManager &taskManager)
                : _taskManager(taskManager) {}

        void handleRequest(Poco::Net::HTTPServerRequest &request, Poco::Net::HTTPServerResponse &response) override {
            if (request.find("Upgrade") != request.end() && Poco::icompare(request["Upgrade"], "websocket") == 0) {
                Poco::Net::WebSocket ws(request, response);
                auto task = new AsyncRequestTask(ws);
                task->runTask();
                delete task;
                task = nullptr;
            } else {
                response.setStatusAndReason(Poco::Net::HTTPResponse::HTTP_OK);
                response.sendBuffer("Hello, world!", 13);
                response.end();
            }
        }

    private:
        Poco::TaskManager &_taskManager;
    };

    class RequestHandlerFactory : public Poco::Net::HTTPRequestHandlerFactory {
    public:
        explicit RequestHandlerFactory(Poco::TaskManager &taskManager) : _taskManager(taskManager) {}

        Poco::Net::HTTPRequestHandler *createRequestHandler(const Poco::Net::HTTPServerRequest &request) override {
            return new WebSocketRequestHandler(_taskManager);
        }

    private:
        Poco::TaskManager &_taskManager;
    };

    class WebSocketServer : public Poco::Util::ServerApplication {
    public:
        static WebSocketServer *instance() {
            if (_instance == nullptr) {
                _instance = new WebSocketServer();
            }
            return _instance;
        }

        void StartServer(const Poco::UInt16 &port) {
            if (_isRunning.exchange(true)) {
                return;
            }
            _serverThread = std::thread([this, port] {
                try {
                    Poco::TaskManager taskManager(Poco::ThreadPool::defaultPool());
                    Poco::Net::ServerSocket svs(port);
                    Poco::Net::HTTPServer server(new RequestHandlerFactory(taskManager), svs,
                                                 new Poco::Net::HTTPServerParams);
                    server.start();
                    {
                        std::unique_lock<std::mutex> lock(_mutex);
                        _condVar.wait(lock, [this] { return !_isRunning.load(); });
                    }
                    server.stop();
                } catch (const Poco::Exception &exc) {
                    std::cerr << exc.displayText() << std::endl;
                }
            });
            _serverThread.detach();
        }

        void Stop() {
            {
                std::lock_guard<std::mutex> lock(_mutex);
                _isRunning = false;
            }
            _condVar.notify_all();
        }

    private:
        std::thread _serverThread;
        std::atomic<bool> _isRunning = false;
        std::condition_variable _condVar;
        static WebSocketServer *_instance;
        std::mutex _mutex;
    };

    WebSocketServer *WebSocketServer::_instance = nullptr;

}

WebSocketManager::receive(ws, message);

I don't know if it's a problem with my code or the library, which caused the websocket to disconnect and the client won't receive any messages