From 32097526fc64017577fe6036199f8acf1a582c1a Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Fri, 7 Jan 2022 19:05:03 +1100
Subject: [PATCH] Add bencode.Decoder.MaxStrLen

---
 bencode/decode.go      | 24 +++++++++++++++++++++++-
 bencode/decode_test.go | 41 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/bencode/decode.go b/bencode/decode.go
index 53ce6efd..67c25a33 100644
--- 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 461b3690..8cb972f8 100644
--- 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)
+}
-- 
2.50.0