From aaf8c81c91487e441226ab0483a8e41a82e4fd46 Mon Sep 17 00:00:00 2001
From: suiwenfeng <suiwenfeng@live.cn>
Date: Tue, 14 Jul 2020 22:37:04 +0800
Subject: [PATCH] fix wrong param value updated while --http-set-param

---
 proto/proto.go      | 26 ++++++++++++++------------
 proto/proto_test.go | 11 ++++++++++-
 2 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/proto/proto.go b/proto/proto.go
index 9bc13ac..c949309 100644
--- a/proto/proto.go
+++ b/proto/proto.go
@@ -323,21 +323,23 @@ func SetPath(payload, path []byte) []byte {
 func PathParam(payload, name []byte) (value []byte, valueStart, valueEnd int) {
 	path := Path(payload)
 
-	if paramStart := bytes.Index(path, append(name, '=')); paramStart != -1 {
-		valueStart := paramStart + len(name) + 1
-		paramEnd := bytes.IndexByte(path[valueStart:], '&')
-
-		// Param can end with '&' (another param), or end of line
-		if paramEnd == -1 { // It is final param
-			paramEnd = len(path)
-		} else {
-			paramEnd += valueStart
+	paramStart := -1
+	if paramStart = bytes.Index(path, append([]byte{'&'}, append(name, '=')...)); paramStart == -1 {
+		if paramStart = bytes.Index(path, append([]byte{'?'}, append(name, '=')...)); paramStart == -1 {
+			return []byte(""), -1, -1
 		}
-
-		return path[valueStart:paramEnd], valueStart, paramEnd
 	}
 
-	return []byte(""), -1, -1
+	valueStart = paramStart + len(name) + 2
+	paramEnd := bytes.IndexByte(path[valueStart:], '&')
+
+	// Param can end with '&' (another param), or end of line
+	if paramEnd == -1 { // It is final param
+		paramEnd = len(path)
+	} else {
+		paramEnd += valueStart
+	}
+	return path[valueStart:paramEnd], valueStart, paramEnd
 }
 
 // SetPathParam takes payload and updates path Query attribute
diff --git a/proto/proto_test.go b/proto/proto_test.go
index 210a9c1..fc46e52 100644
--- a/proto/proto_test.go
+++ b/proto/proto_test.go
@@ -271,7 +271,7 @@ func TestSetPath(t *testing.T) {
 func TestPathParam(t *testing.T) {
 	var payload []byte
 
-	payload = []byte("POST /post?param=test&user_id=1 HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2")
+	payload = []byte("POST /post?param=test&user_id=1&d_type=1&type=2&d_type=3 HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2")
 
 	if val, _, _ := PathParam(payload, []byte("param")); !bytes.Equal(val, []byte("test")) {
 		t.Error("Should detect attribute", string(val))
@@ -280,6 +280,15 @@ func TestPathParam(t *testing.T) {
 	if val, _, _ := PathParam(payload, []byte("user_id")); !bytes.Equal(val, []byte("1")) {
 		t.Error("Should detect attribute", string(val))
 	}
+
+	if val, _, _ := PathParam(payload, []byte("type")); !bytes.Equal(val, []byte("2")) {
+		t.Error("Should detect attribute", string(val))
+	}
+
+	if val, _, _ := PathParam(payload, []byte("d_type")); !bytes.Equal(val, []byte("1")) {
+		// this function is not designed for cases with duplicate param keys
+		t.Error("Should detect attribute", string(val))
+	}
 }
 
 func TestSetPathParam(t *testing.T) {