update unmarshaler

Author: IncSW

unmarshaler.go

@@ -1,103 +1,101 @@
 package bencode
 
 import (
+	"bytes"
 	"errors"
 	"strconv"
 )
 
-func readUntil(data []byte, symbol byte) ([]byte, int, bool) {
-	for i, b := range data {
-		if b == symbol {
-			return data[:i], i, true
-		}
-	}
-	return nil, 0, false
+func Unmarshal(data []byte) (interface{}, error) {
+	return (&unmarshaler{
+		data:   data,
+		length: len(data),
+	}).unmarshal()
 }
 
-func Unmarshal(data []byte) (interface{}, error) {
-	result, _, err := unmarshal(data)
-	if err != nil {
-		return nil, err
-	}
-	return result, nil
+type unmarshaler struct {
+	data   []byte
+	length int
+	cursor int
 }
 
-func unmarshal(data []byte) (interface{}, int, error) {
-	switch data[0] {
+func (u *unmarshaler) unmarshal() (interface{}, error) {
+	switch u.data[u.cursor] {
 	case 'i':
-		integerBuffer, length, ok := readUntil(data[1:], 'e')
-		if !ok {
-			return nil, 0, errors.New("bencode: invalid integer field")
+		u.cursor += 1
+		index := bytes.IndexByte(u.data[u.cursor:], 'e')
+		if index == -1 {
+			return nil, errors.New("bencode: invalid integer field")
 		}
-		integer, err := strconv.ParseInt(b2s(integerBuffer), 10, 64)
+		index += u.cursor
+		integer, err := strconv.ParseInt(b2s(u.data[u.cursor:index]), 10, 64)
 		if err != nil {
-			return nil, 0, err
+			return nil, err
 		}
-		return integer, length + 2, nil
+		u.cursor = index + 1
+		return integer, nil
 
 	case 'l':
+		u.cursor += 1
 		list := []interface{}{}
-		totalLength := 2
-		data = data[1:]
 		for {
-			if len(data) == 0 {
-				return nil, 0, errors.New("bencode: invalid list field")
+			if u.cursor == u.length {
+				return nil, errors.New("bencode: invalid list field")
 			}
-			if data[0] == 'e' {
-				return list, totalLength, nil
+			if u.data[u.cursor] == 'e' {
+				u.cursor += 1
+				return list, nil
 			}
-			value, length, err := unmarshal(data)
+			value, err := u.unmarshal()
 			if err != nil {
-				return nil, 0, err
+				return nil, err
 			}
 			list = append(list, value)
-			data = data[length:]
-			totalLength += length
 		}
 
 	case 'd':
+		u.cursor += 1
 		dictionary := map[string]interface{}{}
-		totalLength := 2
-		data = data[1:]
 		for {
-			if len(data) == 0 {
-				return nil, 0, errors.New("bencode: invalid dictionary field")
+			if u.cursor == u.length {
+				return nil, errors.New("bencode: invalid dictionary field")
 			}
-			if data[0] == 'e' {
-				return dictionary, totalLength, nil
+			if u.data[u.cursor] == 'e' {
+				u.cursor += 1
+				return dictionary, nil
 			}
-			value, length, err := unmarshal(data)
+			value, err := u.unmarshal()
 			if err != nil {
-				return nil, 0, err
+				return nil, err
 			}
 			key, ok := value.([]byte)
 			if !ok {
-				return nil, 0, errors.New("bencode: non-string dictionary key")
+				return nil, errors.New("bencode: non-string dictionary key")
 			}
-			data = data[length:]
-			totalLength += length
-			value, length, err = unmarshal(data)
+			value, err = u.unmarshal()
 			if err != nil {
-				return nil, 0, err
+				return nil, err
 			}
 			dictionary[b2s(key)] = value
-			data = data[length:]
-			totalLength += length
 		}
 
 	default:
-		stringLengthBuffer, length, ok := readUntil(data, ':')
-		if !ok {
-			return nil, 0, errors.New("bencode: invalid string field")
+		index := bytes.IndexByte(u.data[u.cursor:], ':')
+		if index == -1 {
+			return nil, errors.New("bencode: invalid string field")
 		}
-		stringLength, err := strconv.ParseInt(b2s(stringLengthBuffer), 10, 64)
+		index += u.cursor
+		stringLength, err := strconv.ParseInt(b2s(u.data[u.cursor:index]), 10, 64)
 		if err != nil {
-			return nil, 0, err
+			return nil, err
 		}
-		endPosition := length + 1 + int(stringLength)
-		if endPosition > len(data) {
-			return nil, 0, errors.New("bencode: not a valid bencoded string")
+		index += 1
+		endIndex := index + int(stringLength)
+		if endIndex > u.length {
+			return nil, errors.New("bencode: not a valid bencoded string")
 		}
-		return data[length+1 : endPosition], endPosition, nil
+		value := u.data[index:endIndex]
+		u.cursor = endIndex
+		return value, nil
 	}
 }

unmarshaler_test.go

@@ -14,28 +14,13 @@ var (
 func TestUnmarshal(t *testing.T) {
 	assert := assert.New(t)
 
-	_, err := Unmarshal(unmarshalTestData)
-	if !assert.NoError(err) {
-		return
-	}
-
-	_, err = Unmarshal(unmarshalWrongTestData)
-	if !assert.Error(err) || !assert.Equal("bencode: not a valid bencoded string", err.Error()) {
-		return
-	}
-
-	data, length, ok := readUntil([]byte("38:udp://tracker.publicbt.com:80/announce"), ':')
-	if !assert.True(ok) || !assert.Equal(2, length) || !assert.Equal([]byte("38"), data) {
-		return
-	}
-
 	result, err := Unmarshal([]byte("i38e"))
 	if !assert.NoError(err) || !assert.Equal(int64(38), result) {
 		return
 	}
 
 	result, err = Unmarshal([]byte("38:udp://tracker.publicbt.com:80/announce"))
-	if !assert.NoError(err) || !assert.Equal([]byte("udp://tracker.publicbt.com:80/announce"), result) {
+	if !assert.NoError(err) || !assert.Equal("udp://tracker.publicbt.com:80/announce", string(result.([]byte))) {
 		return
 	}
 
@@ -58,6 +43,16 @@ func TestUnmarshal(t *testing.T) {
 		return
 	}
 
+	_, err = Unmarshal(unmarshalTestData)
+	if !assert.NoError(err) {
+		return
+	}
+
+	_, err = Unmarshal(unmarshalWrongTestData)
+	if !assert.Error(err) || !assert.Equal("bencode: not a valid bencoded string", err.Error()) {
+		return
+	}
+
 	result, err = Unmarshal([]byte("i38"))
 	if !assert.Error(err) || !assert.Nil(result) || !assert.Equal("bencode: invalid integer field", err.Error()) {
 		return
@@ -87,6 +82,11 @@ func TestUnmarshal(t *testing.T) {
 	if !assert.Error(err) || !assert.Nil(result) || !assert.Equal("bencode: invalid string field", err.Error()) {
 		return
 	}
+
+	result, err = Unmarshal([]byte("10:wasd"))
+	if !assert.Error(err) || !assert.Nil(result) || !assert.Equal("bencode: not a valid bencoded string", err.Error()) {
+		return
+	}
 }
 
 func BenchmarkUnmarshal(b *testing.B) {