internal/escape: make HexEscape injective across escaped-rune boundaries

internal/escape/escape.go

@@ -56,7 +56,9 @@ func IsASCIIAlphanumeric(r rune) bool {
 func HexEscape(s string, shouldEscape func(s []rune, i int) bool) string {
 	// Do a first pass to see which runes (if any) need escaping.
 	runes := []rune(s)
-	var toEscape []int
+	escapeAt := make([]bool, len(runes))
+	var pending []int
+	n := 0
 	for i := range runes {
 		// In addition to whatever the caller wants escaped, always escape
 		// the start of any substring that HexUnescape would itself parse as
@@ -71,26 +73,43 @@ func HexEscape(s string, shouldEscape func(s []rune, i int) bool) string {
 		// escaped string as a storage key or path, since it lets one
 		// caller-chosen key collide with another's.
 		if shouldEscape(runes, i) || isEscapeSequence(runes, i) {
-			toEscape = append(toEscape, i)
+			escapeAt[i] = true
+			pending = append(pending, i)
+			n++
+		}
+	}
+	// Escaping a rune emits "__0x<hex>__", which starts with "__". If the
+	// input right before an escaped rune ends with "__0x<hex>" or
+	// "__0x<hex>_" (not an escape sequence on its own, since the input does
+	// not continue with "__"), the emitted "__" completes it into one, and
+	// HexUnescape parses the output as a different string: "__0x2f/" would
+	// escape to "__0x2f__0x2f__", which is also what "/0x2f__" escapes to.
+	// Escape the leading "_" of any such dangling prefix. That emits another
+	// "__", so check the newly escaped position the same way.
+	for len(pending) > 0 {
+		i := pending[len(pending)-1]
+		pending = pending[:len(pending)-1]
+		if j, ok := danglingEscapePrefix(runes, i); ok && !escapeAt[j] {
+			escapeAt[j] = true
+			pending = append(pending, j)
+			n++
 		}
 	}
-	if len(toEscape) == 0 {
+	if n == 0 {
 		return s
 	}
 	// Each escaped rune turns into at most 14 runes ("__0x7fffffff__"),
 	// so allocate an extra 13 for each. We'll reslice at the end
 	// if we didn't end up using them.
-	escaped := make([]rune, len(runes)+13*len(toEscape))
-	n := 0 // current index into toEscape
+	escaped := make([]rune, len(runes)+13*n)
 	j := 0 // current index into escaped
 	for i, r := range runes {
-		if n < len(toEscape) && i == toEscape[n] {
+		if escapeAt[i] {
 			// We were asked to escape this rune.
 			for _, x := range fmt.Sprintf("__%#x__", r) {
 				escaped[j] = x
 				j++
 			}
-			n++
 		} else {
 			escaped[j] = r
 			j++
@@ -99,6 +118,32 @@ func HexEscape(s string, shouldEscape func(s []rune, i int) bool) string {
 	return string(escaped[0:j])
 }
 
+// danglingEscapePrefix reports whether r[:end] ends with "__0x<hex>" or
+// "__0x<hex>_": the start of an escape sequence that only lacks the
+// trailing "__" (or its second "_"). When the rune at r[end] is escaped,
+// the "__" that begins its escape sequence supplies what is missing, so
+// HexUnescape would read the output as an escape sequence that was never
+// in the input. It returns the index of the leading "_" of that prefix.
+func danglingEscapePrefix(r []rune, end int) (int, bool) {
+	// The input may already supply the first "_" of the trailing "__".
+	if end > 0 && r[end-1] == '_' {
+		end--
+	}
+	// HexUnescape reads hex digits up to the next "_", so walk back over
+	// the run of non-"_" runes: it must be "0x<hex>", preceded by "__".
+	start := end
+	for start > 0 && r[start-1] != '_' {
+		start--
+	}
+	if start < 2 || end-start < 2 || r[start] != '0' || r[start+1] != 'x' || r[start-2] != '_' {
+		return 0, false
+	}
+	if _, err := strconv.ParseInt(string(r[start+2:end]), 16, 32); err != nil {
+		return 0, false
+	}
+	return start - 2, true
+}
+
 // isEscapeSequence reports whether the rune sequence starting at r[i] would
 // be parsed by HexUnescape as an escape sequence, regardless of whether it
 // was actually produced by a previous call to HexEscape.

internal/escape/escape_test.go

@@ -16,6 +16,7 @@ package escape
 
 import (
 	"testing"
+	"unicode/utf8"
 )
 
 func TestHexEscape(t *testing.T) {
@@ -83,24 +84,97 @@ func TestHexEscapeInjective(t *testing.T) {
 		return i > 1 && r[i] == '/' && r[i-1] == '.' && r[i-2] == '.'
 	}
 
-	pairs := [][2]string{
-		{"a/../b", "a/..__0x2f__b"},
-		{"x/../y", "x/..__0x2f__y"},
-		{"../z", "..__0x2f__z"},
+	// Mirrors the "escape every '/'" style of rule; the dangling-prefix
+	// pairs below need the escaped rune to be reachable without a "..".
+	slash := func(r []rune, i int) bool {
+		return r[i] == '/'
+	}
+
+	pairs := []struct {
+		a, b         string
+		shouldEscape func(r []rune, i int) bool
+	}{
+		{"a/../b", "a/..__0x2f__b", dotDotSlash},
+		{"x/../y", "x/..__0x2f__y", dotDotSlash},
+		{"../z", "..__0x2f__z", dotDotSlash},
+		// A dangling "__0x<hex>" right before an escaped rune used to be
+		// completed into an escape sequence by the "__" that the escaped
+		// rune emits: both sides escaped to "__0x2f__0x2f__".
+		{"__0x2f/", "/0x2f__", slash},
+		{"a__0x2f/b", "a/0x2f__b", slash},
+		{"___0x2f/", "_/0x2f__", slash},
+		// Same with the input supplying the first "_" of the trailing "__":
+		// both sides escaped to "__0x2f___0x2f__".
+		{"__0x2f_/", "/_0x2f__", slash},
+		{"__0x2f..__0x2f/", "__0x2f../0x2f__", slash},
 	}
 	for _, p := range pairs {
-		e0 := HexEscape(p[0], dotDotSlash)
-		e1 := HexEscape(p[1], dotDotSlash)
+		e0 := HexEscape(p.a, p.shouldEscape)
+		e1 := HexEscape(p.b, p.shouldEscape)
 		if e0 == e1 {
-			t.Errorf("collision: HexEscape(%q) == HexEscape(%q) == %q", p[0], p[1], e0)
+			t.Errorf("collision: HexEscape(%q) == HexEscape(%q) == %q", p.a, p.b, e0)
+		}
+		if got := HexUnescape(e0); got != p.a {
+			t.Errorf("HexUnescape(HexEscape(%q)) = %q, want %q", p.a, got, p.a)
+		}
+		if got := HexUnescape(e1); got != p.b {
+			t.Errorf("HexUnescape(HexEscape(%q)) = %q, want %q", p.b, got, p.b)
+		}
+	}
+}
+
+// TestHexEscapeRoundTripExhaustive checks HexUnescape(HexEscape(s)) == s for
+// every string up to a fixed length over an alphabet built from the pieces
+// of an escape sequence. Round-tripping every input also proves HexEscape is
+// injective on that set, since two inputs with the same escaped form would
+// have to unescape to the same string.
+func TestHexEscapeRoundTripExhaustive(t *testing.T) {
+	alphabet := []rune{'_', '0', 'x', '2', 'f', '/'}
+	rules := map[string]func(r []rune, i int) bool{
+		"slash":     func(r []rune, i int) bool { return r[i] == '/' },
+		"non-alnum": func(r []rune, i int) bool { return !IsASCIIAlphanumeric(r[i]) },
+		"none":      func(r []rune, i int) bool { return false },
+	}
+	// Long enough for both forms of dangling prefix, "__0x2/" and "__0x2_/".
+	const maxLen = 7
+	buf := make([]rune, 0, maxLen)
+	var walk func(name string, rule func(r []rune, i int) bool)
+	walk = func(name string, rule func(r []rune, i int) bool) {
+		s := string(buf)
+		if got := HexUnescape(HexEscape(s, rule)); got != s {
+			t.Errorf("%s: HexUnescape(HexEscape(%q)) = %q", name, s, got)
 		}
-		if got := HexUnescape(e0); got != p[0] {
-			t.Errorf("HexUnescape(HexEscape(%q)) = %q, want %q", p[0], got, p[0])
+		if len(buf) == maxLen {
+			return
 		}
-		if got := HexUnescape(e1); got != p[1] {
-			t.Errorf("HexUnescape(HexEscape(%q)) = %q, want %q", p[1], got, p[1])
+		for _, r := range alphabet {
+			buf = append(buf, r)
+			walk(name, rule)
+			buf = buf[:len(buf)-1]
 		}
 	}
+	for name, rule := range rules {
+		buf = buf[:0]
+		walk(name, rule)
+	}
+}
+
+// FuzzHexEscapeRoundTrip is the fuzzing counterpart of
+// TestHexEscapeRoundTripExhaustive, for inputs of arbitrary shape.
+func FuzzHexEscapeRoundTrip(f *testing.F) {
+	for _, s := range []string{"", "a/../b", "__0x2f__", "__0x2f/", "__0x2f_/", "___0x2f/", "__0x/", "__0xzz/", "__0x8fffffff/"} {
+		f.Add(s)
+	}
+	slash := func(r []rune, i int) bool { return r[i] == '/' }
+	f.Fuzz(func(t *testing.T, s string) {
+		if !utf8.ValidString(s) {
+			// Non-UTF-8 input is documented as lossy.
+			return
+		}
+		if got := HexUnescape(HexEscape(s, slash)); got != s {
+			t.Errorf("HexUnescape(HexEscape(%q)) = %q", s, got)
+		}
+	})
 }
 
 func TestHexEscapeUnescapeWeirdStrings(t *testing.T) {