Merge pull request #33 from soheilhy/fix-b32

Fix index out of range in patricia tree

.travis.yml

@@ -24,4 +24,5 @@ before_script:
   - if [[ $TRAVIS_GO_VERSION == 1.6* ]]; then go tool vet --shadow .; fi
 
 script:
-  - go test -v ./...
+  - go test -bench . -v ./...
+  - go test -race -bench . -v ./...

bench_test.go

@@ -62,12 +62,14 @@ func BenchmarkCMuxConnHTTP1(b *testing.B) {
 	wg.Add(b.N)
 
 	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
-		c := &mockConn{
-			r: bytes.NewReader(benchHTTP1Payload),
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			wg.Add(1)
+			m.serve(&mockConn{
+				r: bytes.NewReader(benchHTTP1Payload),
+			}, donec, &wg)
 		}
-		m.serve(c, donec, &wg)
-	}
+	})
 }
 
 func BenchmarkCMuxConnHTTP2(b *testing.B) {
@@ -80,12 +82,14 @@ func BenchmarkCMuxConnHTTP2(b *testing.B) {
 	wg.Add(b.N)
 
 	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
-		c := &mockConn{
-			r: bytes.NewReader(benchHTTP2Payload),
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			wg.Add(1)
+			m.serve(&mockConn{
+				r: bytes.NewReader(benchHTTP2Payload),
+			}, donec, &wg)
 		}
-		m.serve(c, donec, &wg)
-	}
+	})
 }
 
 func BenchmarkCMuxConnHTTP1n2(b *testing.B) {
@@ -98,15 +102,16 @@ func BenchmarkCMuxConnHTTP1n2(b *testing.B) {
 
 	donec := make(chan struct{})
 	var wg sync.WaitGroup
-	wg.Add(b.N)
 
 	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
-		c := &mockConn{
-			r: bytes.NewReader(benchHTTP2Payload),
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			wg.Add(1)
+			m.serve(&mockConn{
+				r: bytes.NewReader(benchHTTP2Payload),
+			}, donec, &wg)
 		}
-		m.serve(c, donec, &wg)
-	}
+	})
 }
 
 func BenchmarkCMuxConnHTTP2n1(b *testing.B) {
@@ -119,13 +124,14 @@ func BenchmarkCMuxConnHTTP2n1(b *testing.B) {
 
 	donec := make(chan struct{})
 	var wg sync.WaitGroup
-	wg.Add(b.N)
 
 	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
-		c := &mockConn{
-			r: bytes.NewReader(benchHTTP1Payload),
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			wg.Add(1)
+			m.serve(&mockConn{
+				r: bytes.NewReader(benchHTTP1Payload),
+			}, donec, &wg)
 		}
-		m.serve(c, donec, &wg)
-	}
+	})
 }

patricia.go

@@ -22,8 +22,8 @@ import (
 // patriciaTree is a simple patricia tree that handles []byte instead of string
 // and cannot be changed after instantiation.
 type patriciaTree struct {
-	root *ptNode
-	buf  []byte // preallocated buffer to read data while matching
+	root     *ptNode
+	maxDepth int // max depth of the tree.
 }
 
 func newPatriciaTree(bs ...[]byte) *patriciaTree {
@@ -34,8 +34,8 @@ func newPatriciaTree(bs ...[]byte) *patriciaTree {
 		}
 	}
 	return &patriciaTree{
-		root: newNode(bs),
-		buf:  make([]byte, max+1),
+		root:     newNode(bs),
+		maxDepth: max + 1,
 	}
 }
 
@@ -48,13 +48,15 @@ func newPatriciaTreeString(strs ...string) *patriciaTree {
 }
 
 func (t *patriciaTree) matchPrefix(r io.Reader) bool {
-	n, _ := io.ReadFull(r, t.buf)
-	return t.root.match(t.buf[:n], true)
+	buf := make([]byte, t.maxDepth)
+	n, _ := io.ReadFull(r, buf)
+	return t.root.match(buf[:n], true)
 }
 
 func (t *patriciaTree) match(r io.Reader) bool {
-	n, _ := io.ReadFull(r, t.buf)
-	return t.root.match(t.buf[:n], false)
+	buf := make([]byte, t.maxDepth)
+	n, _ := io.ReadFull(r, buf)
+	return t.root.match(buf[:n], false)
 }
 
 type ptNode struct {
@@ -159,6 +161,10 @@ func (n *ptNode) match(b []byte, prefix bool) bool {
 		return true
 	}
 
+	if l >= len(b) {
+		return false
+	}
+
 	nextN, ok := n.next[b[l]]
 	if !ok {
 		return false

patricia_test.go

@@ -33,6 +33,13 @@ func testPTree(t *testing.T, strs ...string) {
 		if pt.match(strings.NewReader(s + s)) {
 			t.Errorf("%s matches %s", s+s, s)
 		}
+
+		// The following tests are just to catch index out of
+		// range and off-by-one errors and not the functionality.
+		pt.matchPrefix(strings.NewReader(s[:len(s)-1]))
+		pt.match(strings.NewReader(s[:len(s)-1]))
+		pt.matchPrefix(strings.NewReader(s + "$"))
+		pt.match(strings.NewReader(s + "$"))
 	}
 }
 
@@ -45,5 +52,5 @@ func TestPatriciaNonOverlapping(t *testing.T) {
 }
 
 func TestPatriciaOverlapping(t *testing.T) {
-	testPTree(t, "foo", "far", "farther", "boo", "bar")
+	testPTree(t, "foo", "far", "farther", "boo", "ba", "bar")
 }