Enable 'prepare' for server task too.

src/factory/WFTask.h

@@ -200,6 +200,12 @@ class WFNetworkTask : public CommRequest
 		}
 	}
 
+public:
+	void set_prepare(std::function<void (WFNetworkTask<REQ, RESP> *)> prep)
+	{
+		this->prepare = std::move(prep);
+	}
+
 public:
 	void set_callback(std::function<void (WFNetworkTask<REQ, RESP> *)> cb)
 	{
@@ -219,6 +225,7 @@ class WFNetworkTask : public CommRequest
 	int watch_timeo;
 	REQ req;
 	RESP resp;
+	std::function<void (WFNetworkTask<REQ, RESP> *)> prepare;
 	std::function<void (WFNetworkTask<REQ, RESP> *)> callback;
 
 protected:
@@ -742,6 +749,7 @@ class WFRepeaterTask : public WFGenericTask
 		this->create = std::move(create);
 	}
 
+public:
 	void set_callback(std::function<void (WFRepeaterTask *)> cb)
 	{
 		this->callback = std::move(cb);

src/factory/WFTask.inl

@@ -47,7 +47,7 @@ class WFClientTask : public WFNetworkTask<REQ, RESP>
 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<void (WFNetworkTask<REQ, RESP> *)> prep)
-	{
-		this->prepare = std::move(prep);
-	}
-
-protected:
-	std::function<void (WFNetworkTask<REQ, RESP> *)> prepare;
-
 public:
 	WFClientTask(CommSchedObject *object, CommScheduler *scheduler,
 				 std::function<void (WFNetworkTask<REQ, RESP> *)>&& cb) :
@@ -115,7 +106,16 @@ template<class REQ, class RESP>
 class WFServerTask : public WFNetworkTask<REQ, RESP>
 {
 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);