From 6abc5c53a24c3cc1776e5093c054b5f4cdc2f245 Mon Sep 17 00:00:00 2001
From: wssccc <wssccc@qq.com>
Date: Fri, 3 Sep 2021 22:36:19 +0800
Subject: [PATCH] Add literal colon support (#1432)

---
 gin.go       | 23 +++++++++++++++++++++++
 tree.go      | 17 ++++++++++++++++-
 tree_test.go | 24 ++++++++++++++++++++++++
 3 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/gin.go b/gin.go
index 6ab2be66d0..04d340ec41 100644
--- a/gin.go
+++ b/gin.go
@@ -345,6 +345,7 @@ func (engine *Engine) Run(addr ...string) (err error) {
 	if err != nil {
 		return err
 	}
+	engine.updateRouteTrees()
 
 	address := resolveAddress(addr)
 	debugPrint("Listening and serving HTTP on %s\n", address)
@@ -394,6 +395,28 @@ func (engine *Engine) parseTrustedProxies() error {
 	return err
 }
 
+// updateRouteTree do update to the route tree recursively
+func updateRouteTree(n *node) {
+	const escapedColon = "\\:"
+	const colon = ":"
+	const backslash = "\\"
+	n.path = strings.ReplaceAll(n.path, escapedColon, colon)
+	n.fullPath = strings.ReplaceAll(n.fullPath, escapedColon, colon)
+	n.indices = strings.ReplaceAll(n.indices, backslash, colon)
+	if n.children != nil {
+		for _, child := range n.children {
+			updateRouteTree(child)
+		}
+	}
+}
+
+// updateRouteTrees do update to the route trees
+func (engine *Engine) updateRouteTrees() {
+	for _, tree := range engine.trees {
+		updateRouteTree(tree.root)
+	}
+}
+
 // parseIP parse a string representation of an IP and returns a net.IP with the
 // minimum byte representation or nil if input is invalid.
 func parseIP(ip string) net.IP {
diff --git a/tree.go b/tree.go
index fb0a5935c2..f505689f91 100644
--- a/tree.go
+++ b/tree.go
@@ -261,7 +261,20 @@ walk:
 // Returns -1 as index, if no wildcard was found.
 func findWildcard(path string) (wildcard string, i int, valid bool) {
 	// Find start
+	escapeColon := false
 	for start, c := range []byte(path) {
+		if escapeColon {
+			if c == ':' {
+				escapeColon = false
+				continue
+			} else {
+				panic("invalid escaped char in " + path)
+			}
+		}
+		if c == '\\' {
+			escapeColon = true
+			continue
+		}
 		// A wildcard starts with ':' (param) or '*' (catch-all)
 		if c != ':' && c != '*' {
 			continue
@@ -402,6 +415,7 @@ func (n *node) getValue(path string, params *Params, unescape bool) (value nodeV
 	var (
 		skippedPath string
 		latestNode  = n // Caching the latest node
+		latestPath  = path
 	)
 
 walk: // Outer loop for walking the tree
@@ -427,6 +441,7 @@ walk: // Outer loop for walking the tree
 								handlers:  n.handlers,
 								fullPath:  n.fullPath,
 							}
+							latestPath = path
 						}
 
 						n = n.children[i]
@@ -457,7 +472,7 @@ walk: // Outer loop for walking the tree
 					// fix truncate the parameter
 					// tree_test.go  line: 204
 					if matched {
-						path = prefix + path
+						path = latestPath
 						// The saved path is used after the prefix route is intercepted by matching
 						if n.indices == "/" {
 							path = skippedPath[1:]
diff --git a/tree_test.go b/tree_test.go
index cbb37340ef..9f37e64c6b 100644
--- a/tree_test.go
+++ b/tree_test.go
@@ -185,6 +185,7 @@ func TestTreeWildcard(t *testing.T) {
 		"/get/abc/123abg/:param",
 		"/get/abc/123abf/:param",
 		"/get/abc/123abfff/:param",
+		"/get/abc/escaped_colon/test\\:param",
 	}
 	for _, route := range routes {
 		tree.addRoute(route, fakeHandler(route))
@@ -305,6 +306,7 @@ func TestTreeWildcard(t *testing.T) {
 		{"/get/abc/123abg/test", false, "/get/abc/123abg/:param", Params{Param{Key: "param", Value: "test"}}},
 		{"/get/abc/123abf/testss", false, "/get/abc/123abf/:param", Params{Param{Key: "param", Value: "testss"}}},
 		{"/get/abc/123abfff/te", false, "/get/abc/123abfff/:param", Params{Param{Key: "param", Value: "te"}}},
+		{"/get/abc/escaped_colon/test\\:param", false, "/get/abc/escaped_colon/test\\:param", nil},
 	})
 
 	checkPriorities(t, tree)
@@ -407,6 +409,9 @@ func TestTreeWildcardConflict(t *testing.T) {
 		{"/user_:name", false},
 		{"/id:id", false},
 		{"/id/:id", false},
+		{"/escape/test\\:d1", false},
+		{"/escape/test\\:d2", false},
+		{"/escape/test:param", false},
 	}
 	testRoutes(t, routes)
 }
@@ -886,3 +891,22 @@ func TestTreeWildcardConflictEx(t *testing.T) {
 		}
 	}
 }
+
+func TestTreeInvalidEscape(t *testing.T) {
+	routes := map[string]bool{
+		"/r/r":      true,
+		"/r/:r":     true,
+		"/r/\\:r":   true,
+		"/r/\\\\:r": false,
+		"/r/\\~:r":  false,
+	}
+	tree := &node{}
+	for route, valid := range routes {
+		recv := catchPanic(func() {
+			tree.addRoute(route, fakeHandler(route))
+		})
+		if recv == nil != valid {
+			t.Fatalf("%s should be %t but got %v", route, valid, recv)
+		}
+	}
+}
\ No newline at end of file