add max_header_size, check header size in smtp session

internal/endpoint/smtp/session.go

@@ -40,6 +40,31 @@ import (
 	"github.com/foxcpp/maddy/framework/module"
 )
 
+func limitReader(r io.Reader, n int64, err error) *limitedReader {
+	return &limitedReader{R: r, N: n, E: err, Enabled: true}
+}
+
+type limitedReader struct {
+	R       io.Reader
+	N       int64
+	E       error
+	Enabled bool
+}
+
+// same as io.LimitedReader.Read except returning the custom error and the option
+// to be disabled
+func (l *limitedReader) Read(p []byte) (n int, err error) {
+	if l.Enabled && l.N <= 0 {
+		return 0, l.E
+	}
+	if int64(len(p)) > l.N {
+		p = p[0:l.N]
+	}
+	n, err = l.R.Read(p)
+	l.N -= int64(n)
+	return
+}
+
 type Session struct {
 	endp *Endpoint
 
@@ -340,7 +365,13 @@ func (s *Session) Logout() error {
 }
 
 func (s *Session) prepareBody(ctx context.Context, r io.Reader) (textproto.Header, buffer.Buffer, error) {
-	bufr := bufio.NewReader(r)
+	limitr := limitReader(r, int64(s.endp.maxHeaderBytes), &exterrors.SMTPError{
+		Code:         552,
+		EnhancedCode: exterrors.EnhancedCode{5, 3, 4},
+		Message:      "Message header size exceeds limit",
+	})
+
+	bufr := bufio.NewReader(limitr)
 	header, err := textproto.ReadHeader(bufr)
 	if err != nil {
 		return textproto.Header{}, nil, fmt.Errorf("I/O error while parsing header: %w", err)
@@ -353,6 +384,9 @@ func (s *Session) prepareBody(ctx context.Context, r io.Reader) (textproto.Heade
 		}
 	}
 
+	// the header size check is done. The message size will be checked by go-smtp
+	limitr.Enabled = false
+
 	buf, err := s.endp.buffer(bufr)
 	if err != nil {
 		return textproto.Header{}, nil, fmt.Errorf("I/O error while writing buffer: %w", err)

internal/endpoint/smtp/smtp.go

@@ -67,6 +67,7 @@ type Endpoint struct {
 	deferServerReject   bool
 	maxLoggedRcptErrors int
 	maxReceived         int
+	maxHeaderBytes      int
 
 	listenersWg sync.WaitGroup
 
@@ -245,6 +246,7 @@ func (endp *Endpoint) setConfig(cfg *config.Map) error {
 	cfg.Duration("write_timeout", false, false, 1*time.Minute, &endp.serv.WriteTimeout)
 	cfg.Duration("read_timeout", false, false, 10*time.Minute, &endp.serv.ReadTimeout)
 	cfg.DataSize("max_message_size", false, false, 32*1024*1024, &endp.serv.MaxMessageBytes)
+	cfg.DataSize("max_header_size", false, false, 1*1024*1024, &endp.maxHeaderBytes)
 	cfg.Int("max_recipients", false, false, 20000, &endp.serv.MaxRecipients)
 	cfg.Int("max_received", false, false, 50, &endp.maxReceived)
 	cfg.Custom("buffer", false, false, func() (interface{}, error) {

tests/smtp_test.go

@@ -24,6 +24,7 @@ import (
 	"errors"
 	"io/ioutil"
 	"path/filepath"
+	"strings"
 	"testing"
 
 	"github.com/foxcpp/go-mockdns"
@@ -456,3 +457,41 @@ func TestCheckCommand(tt *testing.T) {
 	conn.Writeln("QUIT")
 	conn.ExpectPattern("221 *")
 }
+
+func TestHeaderSizeConstraint(tt *testing.T) {
+	tt.Parallel()
+	t := tests.NewT(tt)
+	t.DNS(nil)
+	t.Port("smtp")
+	t.Config(`
+		smtp tcp://127.0.0.1:{env:TEST_PORT_smtp} {
+			hostname mx.maddy.test
+			tls off
+			deliver_to dummy
+			max_header_size 1K
+		}
+	`)
+	t.Run(1)
+	defer t.Close()
+
+	conn := t.Conn("smtp")
+	defer conn.Close()
+	conn.SMTPNegotation("localhost", nil, nil)
+	conn.Writeln("MAIL FROM:<testsender@maddy.test>")
+	conn.ExpectPattern("250 *")
+	conn.Writeln("RCPT TO:<testing@maddy.test>")
+	conn.ExpectPattern("250 *")
+	conn.Writeln("DATA")
+	conn.ExpectPattern("354 *")
+	conn.Writeln("From: <testing@sender.test>")
+	conn.Writeln("To: <testing@maddy.test>")
+	conn.Writeln("Subject: " + strings.Repeat("A", 2*1024))
+	conn.Writeln("")
+	conn.Writeln("Hi")
+	conn.Writeln(".")
+
+	conn.ExpectPattern("552 5.3.4 Message header size exceeds limit *")
+
+	conn.Writeln("QUIT")
+	conn.ExpectPattern("221 *")
+}