From c06b8237657d5fa781bdd52523584c0e0ebecf81 Mon Sep 17 00:00:00 2001
From: cqqqq777 <115757192+cqqqq777@users.noreply.github.com>
Date: Sun, 20 Aug 2023 19:27:29 +0800
Subject: [PATCH] fix: fix the bug that temporarily saved the header (#15)

* fix: fix the bug that temporarily saved the header

* feat: use sync.Pool to improve
---
 reverse_proxy.go | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/reverse_proxy.go b/reverse_proxy.go
index 16df778..1256200 100644
--- a/reverse_proxy.go
+++ b/reverse_proxy.go
@@ -31,6 +31,7 @@ import (
 	"net/textproto"
 	"reflect"
 	"strings"
+	"sync"
 	"unsafe"
 
 	"github.com/cloudwego/hertz/pkg/app"
@@ -213,18 +214,24 @@ func (r *ReverseProxy) defaultErrorHandler(c *app.RequestContext, _ error) {
 	c.Response.Header.SetStatusCode(consts.StatusBadGateway)
 }
 
+var respTmpHeaderPool = sync.Pool{
+	New: func() interface{} {
+		return make(map[string][]string)
+	},
+}
+
 func (r *ReverseProxy) ServeHTTP(c context.Context, ctx *app.RequestContext) {
 	req := &ctx.Request
 	resp := &ctx.Response
 
 	// save tmp resp header
-	respTmpHeader := map[string][]string{}
+	respTmpHeader := respTmpHeaderPool.Get().(map[string][]string)
 	if r.saveOriginResHeader {
 		resp.Header.SetNoDefaultContentType(true)
 		resp.Header.VisitAll(func(key, value []byte) {
 			keyStr := string(key)
 			valueStr := string(value)
-			if _, ok := respTmpHeader[keyStr]; ok {
+			if _, ok := respTmpHeader[keyStr]; !ok {
 				respTmpHeader[keyStr] = []string{valueStr}
 			} else {
 				respTmpHeader[keyStr] = append(respTmpHeader[keyStr], valueStr)
@@ -286,6 +293,12 @@ func (r *ReverseProxy) ServeHTTP(c context.Context, ctx *app.RequestContext) {
 		}
 	}
 
+	// Clear and put respTmpHeader back to respTmpHeaderPool
+	for k := range respTmpHeader {
+		delete(respTmpHeader, k)
+	}
+	respTmpHeaderPool.Put(respTmpHeader)
+
 	removeResponseConnHeaders(ctx)
 
 	for _, h := range hopHeaders {