Move header size restriction from textproto to message

It's not really useful to individually check the line length or the
number of fields, what we really care about is the total header size.
This is also what net/http checks for.

Remove size checks from textproto, since callers can implement them. Use
a variant of io.LimitedReader in message instead.

entity.go

@@ -2,7 +2,9 @@ package message
 
 import (
 	"bufio"
+	"errors"
 	"io"
+	"math"
 	"strings"
 
 	"github.com/emersion/go-message/textproto"
@@ -77,6 +79,28 @@ func NewMultipart(header Header, parts []*Entity) (*Entity, error) {
 	return New(header, r)
 }
 
+const maxHeaderBytes = 1 << 20 // 1 MB
+
+var errHeaderTooBig = errors.New("message: header exceeds maximum size")
+
+// limitedReader is the same as io.LimitedReader, but returns a custom error.
+type limitedReader struct {
+	R io.Reader
+	N int64
+}
+
+func (lr *limitedReader) Read(p []byte) (int, error) {
+	if lr.N <= 0 {
+		return 0, errHeaderTooBig
+	}
+	if int64(len(p)) > lr.N {
+		p = p[0:lr.N]
+	}
+	n, err := lr.R.Read(p)
+	lr.N -= int64(n)
+	return n, err
+}
+
 // Read reads a message from r. The message's encoding and charset are
 // automatically decoded to raw UTF-8. Note that this function only reads the
 // message header.
@@ -85,12 +109,16 @@ func NewMultipart(header Header, parts []*Entity) (*Entity, error) {
 // error that verifies IsUnknownCharset or IsUnknownEncoding, but also returns
 // an Entity that can be read.
 func Read(r io.Reader) (*Entity, error) {
-	br := bufio.NewReader(r)
+	lr := &limitedReader{R: r, N: maxHeaderBytes}
+	br := bufio.NewReader(lr)
+
 	h, err := textproto.ReadHeader(br)
 	if err != nil {
 		return nil, err
 	}
 
+	lr.N = math.MaxInt64
+
 	return New(Header{h}, br)
 }
 

entity_test.go

@@ -152,6 +152,16 @@ func TestRead_single(t *testing.T) {
 	}
 }
 
+func TestRead_tooBig(t *testing.T) {
+	raw := "Subject: " + strings.Repeat("A", 4096 * 1024) + "\r\n" +
+		"\r\n" +
+		"This header is too big.\r\n"
+	_, err := Read(strings.NewReader(raw))
+	if err != errHeaderTooBig {
+		t.Fatalf("Read() = %q, want %q", err, errHeaderTooBig)
+	}
+}
+
 func TestEntity_WriteTo_decode(t *testing.T) {
 	e := testMakeEntity()
 

textproto/header.go

@@ -367,16 +367,6 @@ func (h *Header) FieldsByKey(k string) HeaderFields {
 	return &headerFieldsByKey{h, textproto.CanonicalMIMEHeaderKey(k), -1}
 }
 
-// TooBigError is returned by ReadHeader if one of header components are larger
-// than allowed.
-type TooBigError struct {
-	desc string
-}
-
-func (err TooBigError) Error() string {
-	return "textproto: length limit exceeded: " + err.desc
-}
-
 func readLineSlice(r *bufio.Reader, line []byte) ([]byte, error) {
 	for {
 		l, more, err := r.ReadLine()
@@ -385,10 +375,6 @@ func readLineSlice(r *bufio.Reader, line []byte) ([]byte, error) {
 			return line, err
 		}
 
-		if len(line) > maxLineOctets {
-			return line, TooBigError{"line"}
-		}
-
 		if !more {
 			break
 		}
@@ -429,24 +415,19 @@ func hasContinuationLine(r *bufio.Reader) bool {
 	return isSpace(c)
 }
 
-func readContinuedLineSlice(r *bufio.Reader, maxLines int) (int, []byte, error) {
+func readContinuedLineSlice(r *bufio.Reader) ([]byte, error) {
 	// Read the first line. We preallocate slice that it enough
 	// for most fields.
 	line, err := readLineSlice(r, make([]byte, 0, 256))
 	if err == io.EOF && len(line) == 0 {
 		// Header without a body
-		return 0, nil, nil
+		return nil, nil
 	} else if err != nil {
-		return 0, nil, err
-	}
-
-	maxLines--
-	if maxLines <= 0 {
-		return 0, nil, TooBigError{"lines"}
+		return nil, err
 	}
 
 	if len(line) == 0 { // blank line - no continuation
-		return maxLines, line, nil
+		return line, nil
 	}
 
 	line = append(line, '\r', '\n')
@@ -458,15 +439,10 @@ func readContinuedLineSlice(r *bufio.Reader, maxLines int) (int, []byte, error)
 			break // bufio will keep err until next read.
 		}
 
-		maxLines--
-		if maxLines <= 0 {
-			return 0, nil, TooBigError{"lines"}
-		}
-
 		line = append(line, '\r', '\n')
 	}
 
-	return maxLines, line, nil
+	return line, nil
 }
 
 func writeContinued(b *strings.Builder, l []byte) {
@@ -501,13 +477,12 @@ func trimAroundNewlines(v []byte) string {
 	return b.String()
 }
 
-const (
-	maxHeaderLines = 1000
-	maxLineOctets  = 4000
-)
-
 // ReadHeader reads a MIME header from r. The header is a sequence of possibly
 // continued Key: Value lines ending in a blank line.
+//
+// To avoid denial of service attacks, the provided bufio.Reader should be
+// reading from an io.LimitedReader or a similar Reader to bound the size of
+// headers.
 func ReadHeader(r *bufio.Reader) (Header, error) {
 	fs := make([]*headerField, 0, 32)
 
@@ -521,14 +496,12 @@ func ReadHeader(r *bufio.Reader) (Header, error) {
 		return newHeader(fs), fmt.Errorf("message: malformed MIME header initial line: %v", string(line))
 	}
 
-	maxLines := maxHeaderLines
-
 	for {
 		var (
 			kv  []byte
 			err error
 		)
-		maxLines, kv, err = readContinuedLineSlice(r, maxLines)
+		kv, err = readContinuedLineSlice(r)
 		if len(kv) == 0 {
 			return newHeader(fs), err
 		}

textproto/header_test.go

@@ -230,32 +230,6 @@ func TestInvalidHeader(t *testing.T) {
 	}
 }
 
-func TestReadHeader_TooBig(t *testing.T) {
-	testHeader := "Received: from example.com by example.org\r\n" +
-		"Received: from localhost by example.com\r\n" +
-		"To: Taki Tachibana <taki.tachibana@example.org> " + strings.Repeat("A", 4000) + "\r\n" +
-		"From: Mitsuha Miyamizu <mitsuha.miyamizu@example.com>\r\n\r\n"
-	_, err := ReadHeader(bufio.NewReader(strings.NewReader(testHeader)))
-	if err == nil {
-		t.Fatalf("ReadHeader() succeeded")
-	}
-	if _, ok := err.(TooBigError); !ok {
-		t.Fatalf("Not TooBigError returned: %T", err)
-	}
-
-	testHeader = "Received: from example.com by example.org\r\n" +
-		"Received: from localhost by example.com\r\n" +
-		"To: Taki Tachibana <taki.tachibana@example.org>\r\n" +
-		strings.Repeat("From: Mitsuha Miyamizu <mitsuha.miyamizu@example.com>\r\n", 1001) + "\r\n"
-	_, err = ReadHeader(bufio.NewReader(strings.NewReader(testHeader)))
-	if err == nil {
-		t.Fatalf("ReadHeader() succeeded")
-	}
-	if _, ok := err.(TooBigError); !ok {
-		t.Fatalf("Not TooBigError returned: %T", err)
-	}
-}
-
 const testHeaderWithoutBody = "Received: from example.com by example.org\r\n" +
 	"Received: from localhost by example.com\r\n" +
 	"To: Taki Tachibana <taki.tachibana@example.org>\r\n" +