git » chasquid » commit 4edcd79

diff --git a/internal/smtpsrv/conn.go b/internal/smtpsrv/conn.go
index 16385c1..f051012 100644
--- a/internal/smtpsrv/conn.go
+++ b/internal/smtpsrv/conn.go
@@ -592,6 +592,11 @@ func (c *Conn) DATA(params string) (code int, msg string) {
 	c.data, err = ioutil.ReadAll(dotr)
 	if err != nil {
 		if err == io.ErrUnexpectedEOF {
+			// Message is too big already. But we need to keep reading until we see
+			// the "\r\n.\r\n", otherwise we will treat the remanent data that
+			// the user keeps sending as commands, and that's a security
+			// issue.
+			readUntilDot(c.reader)
 			return 552, fmt.Sprintf("5.3.4 Message too big")
 		}
 		return 554, fmt.Sprintf("5.4.0 Error reading DATA: %v", err)
@@ -868,6 +873,24 @@ func boolToStr(b bool) string {
 	return "0"
 }
 
+func readUntilDot(r *bufio.Reader) {
+	prevMore := false
+	for {
+		// The reader will not read more than the size of the buffer,
+		// so this doesn't cause increased memory consumption.
+		// The reader's data deadline will prevent this from continuing
+		// forever.
+		l, more, err := r.ReadLine()
+		if err != nil {
+			break
+		}
+		if !more && !prevMore && string(l) == "." {
+			break
+		}
+		prevMore = more
+	}
+}
+
 // STARTTLS SMTP command handler.
 func (c *Conn) STARTTLS(params string) (code int, msg string) {
 	if c.onTLS {
diff --git a/internal/smtpsrv/conn_test.go b/internal/smtpsrv/conn_test.go
index 9da4664..412dd73 100644
--- a/internal/smtpsrv/conn_test.go
+++ b/internal/smtpsrv/conn_test.go
@@ -1,7 +1,9 @@
 package smtpsrv
 
 import (
+	"bufio"
 	"net"
+	"strings"
 	"testing"
 
 	"blitiri.com.ar/go/chasquid/internal/domaininfo"
@@ -79,6 +81,57 @@ func TestIsHeader(t *testing.T) {
 	}
 }
 
+func TestReadUntilDot(t *testing.T) {
+	// This must be > than the minimum buffer size for bufio.Reader, which
+	// unfortunately is not available to us. The current value is 16, these
+	// tests will break if it gets increased, and the nonfinal cases will need
+	// to be adjusted.
+	size := 20
+	xs := "12345678901234567890"
+
+	final := []string{
+		"", ".", "..",
+		".\r\n", "\r\n.", "\r\n.\r\n",
+		".\n", "\n.", "\n.\n",
+		".\r", "\r.", "\r.\r",
+		xs + "\r\n.\r\n",
+		xs + "1234\r\n.\r\n",
+		xs + xs + "\r\n.\r\n",
+		xs + xs + xs + "\r\n.\r\n",
+		xs + "." + xs + "\n.",
+		xs + ".\n" + xs + "\n.",
+	}
+	for _, s := range final {
+		t.Logf("testing %q", s)
+		buf := bufio.NewReaderSize(strings.NewReader(s), size)
+		readUntilDot(buf)
+		if r := buf.Buffered(); r != 0 {
+			t.Errorf("%q: there are %d remaining bytes", s, r)
+		}
+	}
+
+	nonfinal := []struct {
+		s string
+		r int
+	}{
+		{".\na", 1},
+		{"\n.\na", 1},
+		{"\n.\nabc", 3},
+		{"\n.\n12345678", 8},
+		{"\n.\n" + xs, size - 3},
+		{"\n.\n" + xs + xs, size - 3},
+		{"\n.\n.\n", 2},
+	}
+	for _, c := range nonfinal {
+		t.Logf("testing %q", c.s)
+		buf := bufio.NewReaderSize(strings.NewReader(c.s), size)
+		readUntilDot(buf)
+		if r := buf.Buffered(); r != c.r {
+			t.Errorf("%q: expected %d remaining bytes, got %d", c.s, c.r, r)
+		}
+	}
+}
+
 func TestAddrLiteral(t *testing.T) {
 	// TCP addresses.
 	casesTCP := []struct {
diff --git a/internal/smtpsrv/server_test.go b/internal/smtpsrv/server_test.go
index 835231d..557b52a 100644
--- a/internal/smtpsrv/server_test.go
+++ b/internal/smtpsrv/server_test.go
@@ -349,6 +349,15 @@ func TestTooMuchData(t *testing.T) {
 	if err == nil || err.Error() != "552 5.3.4 Message too big" {
 		t.Fatalf("Expected message too big, got: %v", err)
 	}
+
+	// Repeat the test once again, the limit should not prevent connection
+	// from continuing.
+	localC.Expect(1)
+	err = sendLargeEmail(t, c, maxDataSizeMiB-1)
+	if err != nil {
+		t.Errorf("Error sending large but ok email: %v", err)
+	}
+	localC.Wait()
 }
 
 func simpleCmd(t *testing.T, c *smtp.Client, cmd string, expected int) string {