TcpSerer/UdpServer支持创建session前拦截

xiongguangjie · Jun 2, 2024 · a3979d7 · a3979d7
1 parent 1e1a990
commit a3979d7
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 7 deletions.
diff --git a/src/Network/TcpServer.h b/src/Network/TcpServer.h
@@ -42,15 +42,18 @@ class TcpServer : public Server {
     * @param host 监听网卡ip
     * @param backlog tcp listen backlog
     */
-    template<typename SessionType>
-    void start(uint16_t port, const std::string &host = "::", uint32_t backlog = 1024) {
+    template <typename SessionType>
+    void start(uint16_t port, const std::string &host = "::", uint32_t backlog = 1024, const std::function<void(std::shared_ptr<SessionType> &)> &cb = nullptr) {
         static std::string cls_name = toolkit::demangle(typeid(SessionType).name());
-        //Session创建器，通过它创建不同类型的服务器
-        _session_alloc = [](const TcpServer::Ptr &server, const Socket::Ptr &sock) {
-            auto session = std::shared_ptr<SessionType>(new SessionType(sock), [](SessionType * ptr) {
+        // Session创建器，通过它创建不同类型的服务器
+        _session_alloc = [cb](const TcpServer::Ptr &server, const Socket::Ptr &sock) {
+            auto session = std::shared_ptr<SessionType>(new SessionType(sock), [](SessionType *ptr) {
                 TraceP(static_cast<Session *>(ptr)) << "~" << cls_name;
                 delete ptr;
             });
+            if (cb) {
+                cb(session);
+            }
             TraceP(static_cast<Session *>(session.get())) << cls_name;
             session->setOnCreateSocket(server->_on_create_socket);
             return std::make_shared<SessionHelper>(server, std::move(session), cls_name);

diff --git a/src/Network/UdpServer.h b/src/Network/UdpServer.h
@@ -29,14 +29,17 @@ class UdpServer : public Server {
      * @brief 开始监听服务器
      */
     template<typename SessionType>
-    void start(uint16_t port, const std::string &host = "::") {
+    void start(uint16_t port, const std::string &host = "::", const std::function<void(std::shared_ptr<SessionType> &)> &cb = nullptr) {
         static std::string cls_name = toolkit::demangle(typeid(SessionType).name());
         // Session 创建器, 通过它创建不同类型的服务器
-        _session_alloc = [](const UdpServer::Ptr &server, const Socket::Ptr &sock) {
+        _session_alloc = [cb](const UdpServer::Ptr &server, const Socket::Ptr &sock) {
             auto session = std::shared_ptr<SessionType>(new SessionType(sock), [](SessionType * ptr) {
                 TraceP(static_cast<Session *>(ptr)) << "~" << cls_name;
                 delete ptr;
             });
+            if (cb) {
+                cb(session);
+            }
             TraceP(static_cast<Session *>(session.get())) << cls_name;
             auto sock_creator = server->_on_create_socket;
             session->setOnCreateSocket([sock_creator](const EventPoller::Ptr &poller) {