Add bencode.Decoder.MaxStrLen

diff --git a/bencode/decode.go b/bencode/decode.go
index 53ce6efd9035f9bb1034f3652b7c59e45bfd54be..67c25a33ab92af61f8cbd48ee6719b5cc99bfd1a 100644 (file)
--- a/bencode/decode.go
+++ b/bencode/decode.go
@@ -12,7 +12,16 @@ import (
        "sync"
 )
 
+// The default bencode string length limit. This is a poor attempt to prevent excessive memory
+// allocation when parsing, but also leaves the window open to implement a better solution.
+const DefaultDecodeMaxStrLen = 1<<27 - 1 // ~128MiB
+
+type MaxStrLen = int64
+
 type Decoder struct {
+       // Maximum parsed bencode string length. Defaults to DefaultMaxStrLen if zero.
+       MaxStrLen MaxStrLen
+
        r interface {
                io.ByteScanner
                io.Reader
@@ -182,8 +191,14 @@ func (d *Decoder) parseStringLength() (uint64, error) {
        if err := d.checkBufferedInt(); err != nil {
                return 0, err
        }
-       length, err := strconv.ParseUint(bytesAsString(d.buf.Bytes()), 10, 32)
+       // Really the limit should be the uint size for the platform. But we can't pass in an allocator,
+       // or limit total memory use in Go, the best we might hope to do is limit the size of a single
+       // decoded value (by reading it in in-place and then operating on a view).
+       length, err := strconv.ParseUint(bytesAsString(d.buf.Bytes()), 10, 0)
        checkForIntParseError(err, start)
+       if int64(length) > d.getMaxStrLen() {
+               err = fmt.Errorf("parsed string length %v exceeds limit (%v)", length, DefaultDecodeMaxStrLen)
+       }
        d.buf.Reset()
        return length, err
 }
@@ -707,3 +722,10 @@ func (d *Decoder) parseListInterface() (list []interface{}) {
        }
        return
 }
+
+func (d *Decoder) getMaxStrLen() int64 {
+       if d.MaxStrLen == 0 {
+               return DefaultDecodeMaxStrLen
+       }
+       return d.MaxStrLen
+}

diff --git a/bencode/decode_test.go b/bencode/decode_test.go
index 461b3690862743214c32c3bfd2d0254ac4cb5fde..8cb972f81e6cde9868fe7af544e1d90261b3b4a3 100644 (file)
--- a/bencode/decode_test.go
+++ b/bencode/decode_test.go
@@ -2,6 +2,7 @@ package bencode
 
 import (
        "bytes"
+       "fmt"
        "io"
        "math/big"
        "reflect"
@@ -193,3 +194,43 @@ func TestUnmarshalDictKeyNotString(t *testing.T) {
        t.Log(err)
        c.Check(err, qt.Not(qt.IsNil))
 }
+
+type arbitraryReader struct{}
+
+func (arbitraryReader) Read(b []byte) (int, error) {
+       return len(b), nil
+}
+
+func decodeHugeString(t *testing.T, strLen int64, header, tail string, v interface{}, maxStrLen MaxStrLen) error {
+       r, w := io.Pipe()
+       go func() {
+               fmt.Fprintf(w, header, strLen)
+               io.CopyN(w, arbitraryReader{}, strLen)
+               w.Write([]byte(tail))
+               w.Close()
+       }()
+       d := NewDecoder(r)
+       d.MaxStrLen = maxStrLen
+       return d.Decode(v)
+
+}
+
+// Ensure that bencode strings in various places obey the Decoder.MaxStrLen field.
+func TestDecodeMaxStrLen(t *testing.T) {
+       t.Parallel()
+       c := qt.New(t)
+       test := func(header, tail string, v interface{}, maxStrLen MaxStrLen) {
+               strLen := maxStrLen
+               if strLen == 0 {
+                       strLen = DefaultDecodeMaxStrLen
+               }
+               c.Assert(decodeHugeString(t, strLen, header, tail, v, maxStrLen), qt.IsNil)
+               c.Assert(decodeHugeString(t, strLen+1, header, tail, v, maxStrLen), qt.IsNotNil)
+       }
+       test("d%d:", "i0ee", new(interface{}), 0)
+       test("%d:", "", new(interface{}), DefaultDecodeMaxStrLen)
+       test("%d:", "", new([]byte), 1)
+       test("d3:420%d:", "e", new(struct {
+               Hi []byte `bencode:"420"`
+       }), 69)
+}