Add regex support for websocket controller (drogonframework#1779)

hwc0919 · Jun 4, 2024 · 9a96a20 · 9a96a20
1 parent f37a1d0
commit 9a96a20
Show file tree

Hide file tree

Showing 7 changed files with 105 additions and 1 deletion.
diff --git a/examples/websocket_server/WebSocketServer.cc b/examples/websocket_server/WebSocketServer.cc
@@ -14,6 +14,7 @@ class WebSocketChat : public drogon::WebSocketController<WebSocketChat>
                              const WebSocketConnectionPtr &) override;
     WS_PATH_LIST_BEGIN
     WS_PATH_ADD("/chat", Get);
+    WS_ADD_PATH_VIA_REGEX("/[^/]*", Get);
     WS_PATH_LIST_END
   private:
     PubSubService<std::string> chatRooms_;

diff --git a/lib/inc/drogon/HttpAppFramework.h b/lib/inc/drogon/HttpAppFramework.h
@@ -618,6 +618,18 @@ class DROGON_EXPORT HttpAppFramework : public trantor::NonCopyable
         const std::string &ctrlName,
         const std::vector<internal::HttpConstraint> &constraints = {}) = 0;
 
+    /// Register a WebSocketController into the framework.
+    /**
+     * The parameters of this method are the same as those in the
+     * registerHttpSimpleController() method but using regular
+     * expression string for path.
+     */
+    virtual HttpAppFramework &registerWebSocketControllerRegex(
+        const std::string &regExp,
+        const std::string &ctrlName,
+        const std::vector<internal::HttpConstraint> &constraints =
+            std::vector<internal::HttpConstraint>{}) = 0;
+
     /// Register controller objects created and initialized by the user
     /**
      * @details Drogon can only automatically create controllers using the

diff --git a/lib/inc/drogon/WebSocketController.h b/lib/inc/drogon/WebSocketController.h
@@ -28,6 +28,8 @@
     static void initPathRouting() \
     {
 #define WS_PATH_ADD(path, ...) registerSelf__(path, {__VA_ARGS__})
+#define WS_ADD_PATH_VIA_REGEX(regExp, ...) \
+    registerSelfRegex__(regExp, {__VA_ARGS__})
 #define WS_PATH_LIST_END }
 
 namespace drogon
@@ -94,6 +96,19 @@ class WebSocketController : public DrObject<T>, public WebSocketControllerBase
             constraints);
     }
 
+    static void registerSelfRegex__(
+        const std::string &regExp,
+        const std::vector<internal::HttpConstraint> &constraints)
+    {
+        LOG_TRACE << "register websocket controller("
+                  << WebSocketController<T, AutoCreation>::classTypeName()
+                  << ") on regExp:" << regExp;
+        app().registerWebSocketControllerRegex(
+            regExp,
+            WebSocketController<T, AutoCreation>::classTypeName(),
+            constraints);
+    }
+
   private:
     class pathRegistrator
     {

diff --git a/lib/src/HttpAppFrameworkImpl.cc b/lib/src/HttpAppFrameworkImpl.cc
@@ -294,6 +294,17 @@ HttpAppFramework &HttpAppFrameworkImpl::registerWebSocketController(
     return *this;
 }
 
+HttpAppFramework &HttpAppFrameworkImpl::registerWebSocketControllerRegex(
+    const std::string &regExp,
+    const std::string &ctrlName,
+    const std::vector<internal::HttpConstraint> &constraints)
+{
+    assert(!routersInit_);
+    HttpControllersRouter::instance().registerWebSocketControllerRegex(
+        regExp, ctrlName, constraints);
+    return *this;
+}
+
 HttpAppFramework &HttpAppFrameworkImpl::registerHttpSimpleController(
     const std::string &pathName,
     const std::string &ctrlName,

diff --git a/lib/src/HttpAppFrameworkImpl.h b/lib/src/HttpAppFrameworkImpl.h
@@ -89,6 +89,10 @@ class HttpAppFrameworkImpl final : public HttpAppFramework
         const std::string &pathName,
         const std::string &ctrlName,
         const std::vector<internal::HttpConstraint> &constraints) override;
+    HttpAppFramework &registerWebSocketControllerRegex(
+        const std::string &regExp,
+        const std::string &ctrlName,
+        const std::vector<internal::HttpConstraint> &constraints) override;
     HttpAppFramework &registerHttpSimpleController(
         const std::string &pathName,
         const std::string &ctrlName,

diff --git a/lib/src/HttpControllersRouter.cc b/lib/src/HttpControllersRouter.cc
@@ -105,7 +105,10 @@ std::vector<HttpHandlerInfo> HttpControllersRouter::getHandlersInfo() const
                 }
                 else if constexpr (std::is_same_v<
                                        std::decay_t<decltype(item)>,
-                                       WebSocketControllerRouterItem>)
+                                       WebSocketControllerRouterItem> ||
+                                   std::is_same_v<
+                                       std::decay_t<decltype(item)>,
+                                       RegExWebSocketControllerRouterItem>)
                 {
                     description = std::string("WebsocketController: ") +
                                   item.binders_[i]->handlerName_;
@@ -140,6 +143,10 @@ std::vector<HttpHandlerInfo> HttpControllersRouter::getHandlersInfo() const
     {
         gatherInfo(path, item);
     }
+    for (auto &item : wsCtrlVector_)
+    {
+        gatherInfo(item.pathPattern_, item);
+    }
     return ret;
 }
 
@@ -276,6 +283,30 @@ void HttpControllersRouter::registerWebSocketController(
     addCtrlBinderToRouterItem(binder, item, result.validMethods);
 }
 
+void HttpControllersRouter::registerWebSocketControllerRegex(
+    const std::string &regExp,
+    const std::string &ctrlName,
+    const std::vector<internal::HttpConstraint> &constraints)
+{
+    assert(!regExp.empty());
+    assert(!ctrlName.empty());
+    auto result = processSimpleControllerParams(regExp, constraints);
+    auto binder = std::make_shared<WebsocketControllerBinder>();
+    binder->handlerName_ = ctrlName;
+    binder->middlewareNames_ = result.middlewares;
+    drogon::app().getLoop()->queueInLoop([binder, ctrlName]() {
+        auto &object_ = DrClassMap::getSingleInstance(ctrlName);
+        auto controller =
+            std::dynamic_pointer_cast<WebSocketControllerBase>(object_);
+        binder->controller_ = controller;
+    });
+    struct RegExWebSocketControllerRouterItem router;
+    router.pathPattern_ = regExp;
+    router.regex_ = regExp;
+    addCtrlBinderToRouterItem(binder, router, result.validMethods);
+    wsCtrlVector_.push_back(std::move(router));
+}
+
 void HttpControllersRouter::addHttpRegex(
     const std::string &regExp,
     const internal::HttpBinderBasePtr &binder,
@@ -681,6 +712,24 @@ RouteResult HttpControllersRouter::routeWs(const HttpRequestImplPtr &req)
             }
             return {RouteResult::Success, binder};
         }
+        else
+        {
+            for (auto &ctrlInfo : wsCtrlVector_)
+            {
+                auto const &wsCtrlRegex = ctrlInfo.regex_;
+                std::smatch result;
+                if (std::regex_match(req->path(), result, wsCtrlRegex))
+                {
+                    req->setMatchedPathPattern(ctrlInfo.pathPattern_);
+                    auto &binder = ctrlInfo.binders_[req->method()];
+                    if (!binder)
+                    {
+                        return {RouteResult::MethodNotAllowed, nullptr};
+                    }
+                    return {RouteResult::Success, binder};
+                }
+            }
+        }
     }
     return {RouteResult::NotFound, nullptr};
 }

diff --git a/lib/src/HttpControllersRouter.h b/lib/src/HttpControllersRouter.h
@@ -50,6 +50,10 @@ class HttpControllersRouter : public trantor::NonCopyable
         const std::string &pathName,
         const std::string &ctrlName,
         const std::vector<internal::HttpConstraint> &constraints);
+    void registerWebSocketControllerRegex(
+        const std::string &regExp,
+        const std::string &ctrlName,
+        const std::vector<internal::HttpConstraint> &constraints);
     void addHttpPath(const std::string &path,
                      const internal::HttpBinderBasePtr &binder,
                      const std::vector<HttpMethod> &validMethods,
@@ -89,9 +93,17 @@ class HttpControllersRouter : public trantor::NonCopyable
         std::shared_ptr<WebsocketControllerBinder> binders_[Invalid]{nullptr};
     };
 
+    struct RegExWebSocketControllerRouterItem
+    {
+        std::string pathPattern_;
+        std::regex regex_;
+        std::shared_ptr<WebsocketControllerBinder> binders_[Invalid]{nullptr};
+    };
+
     std::unordered_map<std::string, SimpleControllerRouterItem> simpleCtrlMap_;
     std::unordered_map<std::string, HttpControllerRouterItem> ctrlMap_;
     std::vector<HttpControllerRouterItem> ctrlVector_;  // for regexp path
     std::unordered_map<std::string, WebSocketControllerRouterItem> wsCtrlMap_;
+    std::vector<RegExWebSocketControllerRouterItem> wsCtrlVector_;
 };
 }  // namespace drogon