diff --git a/src/factory/WFTask.h b/src/factory/WFTask.h index 78c8c318..94bdf12d 100644 --- a/src/factory/WFTask.h +++ b/src/factory/WFTask.h @@ -200,6 +200,12 @@ public: } } +public: + void set_prepare(std::function *)> prep) + { + this->prepare = std::move(prep); + } + public: void set_callback(std::function *)> cb) { @@ -219,6 +225,7 @@ protected: int watch_timeo; REQ req; RESP resp; + std::function *)> prepare; std::function *)> callback; protected: @@ -742,6 +749,7 @@ public: this->create = std::move(create); } +public: void set_callback(std::function cb) { this->callback = std::move(cb); diff --git a/src/factory/WFTask.inl b/src/factory/WFTask.inl index ace1eaaa..d85bf613 100644 --- a/src/factory/WFTask.inl +++ b/src/factory/WFTask.inl @@ -47,7 +47,7 @@ class WFClientTask : public WFNetworkTask protected: virtual CommMessageOut *message_out() { - /* By using prepare function, users can modify request after + /* By using prepare function, users can modify the request after * the connection is established. */ if (this->prepare) this->prepare(this); @@ -91,15 +91,6 @@ protected: return series->pop(); } -public: - void set_prepare(std::function *)> prep) - { - this->prepare = std::move(prep); - } - -protected: - std::function *)> prepare; - public: WFClientTask(CommSchedObject *object, CommScheduler *scheduler, std::function *)>&& cb) : @@ -115,7 +106,16 @@ template class WFServerTask : public WFNetworkTask { protected: - virtual CommMessageOut *message_out() { return &this->resp; } + virtual CommMessageOut *message_out() + { + /* By using prepare function, users can modify the response before + * replying to the client. */ + if (this->prepare) + this->prepare(this); + + return &this->resp; + } + virtual CommMessageIn *message_in() { return &this->req; } virtual void handle(int state, int error);