Merge branch 'develop' into fix/unexpected_generated_package_path

cloudwego · Jul 18, 2024 · f0c0741 · f0c0741
2 parents 0a0c434 + 764df4d
commit f0c0741
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 0 deletions.
diff --git a/pkg/route/engine.go b/pkg/route/engine.go
@@ -75,6 +75,7 @@ import (
 	"github.com/cloudwego/hertz/pkg/protocol/http1"
 	"github.com/cloudwego/hertz/pkg/protocol/http1/factory"
 	"github.com/cloudwego/hertz/pkg/protocol/suite"
+	"github.com/cloudwego/hertz/pkg/route/param"
 )
 
 const unknownTransporterName = "unknown"
@@ -749,6 +750,12 @@ func (engine *Engine) ServeHTTP(c context.Context, ctx *app.RequestContext) {
 		return
 	}
 
+	// if Params is re-assigned in HandlerFunc and the capacity is not enough we need to realloc
+	maxParams := int(engine.maxParams)
+	if cap(ctx.Params) < maxParams {
+		ctx.Params = make(param.Params, 0, maxParams)
+	}
+
 	// Find root of the tree for the given HTTP method
 	t := engine.trees
 	paramsPointer := &ctx.Params

diff --git a/pkg/route/engine_test.go b/pkg/route/engine_test.go
@@ -1029,3 +1029,33 @@ func TestAcquireHijackConn(t *testing.T) {
 	assert.DeepEqual(t, engine, hijackConn.e)
 	assert.DeepEqual(t, conn, hijackConn.Conn)
 }
+
+func TestHandleParamsReassignInHandleFunc(t *testing.T) {
+	e := NewEngine(config.NewOptions(nil))
+	routes := []string{
+		"/:a/:b/:c",
+	}
+	for _, r := range routes {
+		e.GET(r, func(c context.Context, ctx *app.RequestContext) {
+			ctx.Params = make([]param.Param, 1)
+			ctx.String(consts.StatusOK, "")
+		})
+	}
+	testRoutes := []string{
+		"/aaa/bbb/ccc",
+		"/asd/alskja/alkdjad",
+		"/asd/alskja/alkdjad",
+		"/asd/alskja/alkdjad",
+		"/asd/alskja/alkdjad",
+		"/alksjdlakjd/ooo/askda",
+		"/alksjdlakjd/ooo/askda",
+		"/alksjdlakjd/ooo/askda",
+	}
+	ctx := e.ctxPool.Get().(*app.RequestContext)
+	for _, tr := range testRoutes {
+		r := protocol.NewRequest(http.MethodGet, tr, nil)
+		r.CopyTo(&ctx.Request)
+		e.ServeHTTP(context.Background(), ctx)
+		ctx.ResetWithoutConn()
+	}
+}