From 1262266ac2a72762166be4c9c0425e36c0cb3251 Mon Sep 17 00:00:00 2001
From: TaoYu <71547744+zhouyusd@users.noreply.github.com>
Date: Thu, 15 Jun 2023 23:31:15 +0800
Subject: [PATCH] feat: httpx add common handler (#3269)

---
 rest/httpx/responses.go | 35 ++++++++++++++++++++++++++++++++---
 1 file changed, 32 insertions(+), 3 deletions(-)

diff --git a/rest/httpx/responses.go b/rest/httpx/responses.go
index 99b17fd8..29f54a5a 100644
--- a/rest/httpx/responses.go
+++ b/rest/httpx/responses.go
@@ -13,9 +13,12 @@ import (
 )
 
 var (
-	errorHandler    func(error) (int, any)
-	errorHandlerCtx func(context.Context, error) (int, any)
-	lock            sync.RWMutex
+	errorHandler     func(error) (int, any)
+	errorHandlerCtx  func(context.Context, error) (int, any)
+	commonHandler    func(any) any
+	commonHandlerCtx func(context.Context, any) any
+	lock             sync.RWMutex
+	cLock            sync.RWMutex
 )
 
 // Error writes err into w.
@@ -53,11 +56,23 @@ func Ok(w http.ResponseWriter) {
 
 // OkJson writes v into w with 200 OK.
 func OkJson(w http.ResponseWriter, v any) {
+	cLock.RLock()
+	handler := commonHandler
+	cLock.RUnlock()
+	if handler != nil {
+		v = handler(v)
+	}
 	WriteJson(w, http.StatusOK, v)
 }
 
 // OkJsonCtx writes v into w with 200 OK.
 func OkJsonCtx(ctx context.Context, w http.ResponseWriter, v any) {
+	cLock.RLock()
+	handlerCtx := commonHandlerCtx
+	cLock.RUnlock()
+	if handlerCtx != nil {
+		v = handlerCtx(ctx, v)
+	}
 	WriteJsonCtx(ctx, w, http.StatusOK, v)
 }
 
@@ -75,6 +90,20 @@ func SetErrorHandlerCtx(handlerCtx func(context.Context, error) (int, any)) {
 	errorHandlerCtx = handlerCtx
 }
 
+// SetCommonHandler sets the common handler, which is called on calling OkJson.
+func SetCommonHandler(handler func(any) any) {
+	cLock.Lock()
+	defer cLock.Unlock()
+	commonHandler = handler
+}
+
+// SetCommonHandlerCtx sets the common handler, which is called on calling OkJson.
+func SetCommonHandlerCtx(handlerCtx func(context.Context, any) any) {
+	cLock.Lock()
+	defer cLock.Unlock()
+	commonHandlerCtx = handlerCtx
+}
+
 // WriteJson writes v as json string into w with code.
 func WriteJson(w http.ResponseWriter, code int, v any) {
 	if err := doWriteJson(w, code, v); err != nil {