git » chasquid » commit d7006d0

diff --git a/internal/smtp/smtp.go b/internal/smtp/smtp.go
index bb91b85..c011f5e 100644
--- a/internal/smtp/smtp.go
+++ b/internal/smtp/smtp.go
@@ -7,6 +7,8 @@
 package smtp
 
 import (
+	"bufio"
+	"io"
 	"net"
 	"net/smtp"
 	"net/textproto"
@@ -28,6 +30,14 @@ func NewClient(conn net.Conn, host string) (*Client, error) {
 	if err != nil {
 		return nil, err
 	}
+
+	// Wrap the textproto.Conn reader so we are not exposed to a memory
+	// exhaustion DoS on very long replies from the server.
+	// Limit to 2 MiB total (all replies through the lifetime of the client),
+	// which should be plenty for our uses of SMTP.
+	lr := &io.LimitedReader{R: c.Text.Reader.R, N: 2 * 1024 * 1024}
+	c.Text.Reader.R = bufio.NewReader(lr)
+
 	return &Client{c}, nil
 }
 
diff --git a/internal/smtp/smtp_test.go b/internal/smtp/smtp_test.go
index 5cc1f7d..75e5cfb 100644
--- a/internal/smtp/smtp_test.go
+++ b/internal/smtp/smtp_test.go
@@ -4,8 +4,8 @@ import (
 	"bufio"
 	"bytes"
 	"fmt"
+	"io"
 	"net"
-	"net/smtp"
 	"net/textproto"
 	"strings"
 	"testing"
@@ -48,8 +48,19 @@ func TestIsASCII(t *testing.T) {
 	}
 }
 
+func mustNewClient(t *testing.T, nc net.Conn) *Client {
+	t.Helper()
+
+	c, err := NewClient(nc, "")
+	if err != nil {
+		t.Fatalf("failed to create client: %v", err)
+	}
+	return c
+}
+
 func TestBasic(t *testing.T) {
-	fake, client := fakeDialog(`> EHLO a_test
+	fake, client := fakeDialog(`< 220 welcome
+> EHLO a_test
 < 250-server replies your hello
 < 250-SIZE 35651584
 < 250-SMTPUTF8
@@ -61,8 +72,7 @@ func TestBasic(t *testing.T) {
 < 250 RCPT TO is fine
 `)
 
-	c := &Client{
-		Client: &smtp.Client{Text: textproto.NewConn(fake)}}
+	c := mustNewClient(t, fake)
 	if err := c.Hello("a_test"); err != nil {
 		t.Fatalf("Hello failed: %v", err)
 	}
@@ -78,7 +88,8 @@ func TestBasic(t *testing.T) {
 }
 
 func TestSMTPUTF8(t *testing.T) {
-	fake, client := fakeDialog(`> EHLO araña
+	fake, client := fakeDialog(`< 220 welcome
+> EHLO araña
 < 250-chasquid replies your hello
 < 250-SIZE 35651584
 < 250-SMTPUTF8
@@ -90,8 +101,7 @@ func TestSMTPUTF8(t *testing.T) {
 < 250 RCPT TO is fine
 `)
 
-	c := &Client{
-		Client: &smtp.Client{Text: textproto.NewConn(fake)}}
+	c := mustNewClient(t, fake)
 	if err := c.Hello("araña"); err != nil {
 		t.Fatalf("Hello failed: %v", err)
 	}
@@ -107,15 +117,15 @@ func TestSMTPUTF8(t *testing.T) {
 }
 
 func TestSMTPUTF8NotSupported(t *testing.T) {
-	fake, client := fakeDialog(`> EHLO araña
+	fake, client := fakeDialog(`< 220 welcome
+> EHLO araña
 < 250-chasquid replies your hello
 < 250-SIZE 35651584
 < 250-8BITMIME
 < 250 HELP
 `)
 
-	c := &Client{
-		Client: &smtp.Client{Text: textproto.NewConn(fake)}}
+	c := mustNewClient(t, fake)
 	if err := c.Hello("araña"); err != nil {
 		t.Fatalf("Hello failed: %v", err)
 	}
@@ -135,7 +145,8 @@ func TestSMTPUTF8NotSupported(t *testing.T) {
 }
 
 func TestFallbackToIDNA(t *testing.T) {
-	fake, client := fakeDialog(`> EHLO araña
+	fake, client := fakeDialog(`< 220 welcome
+> EHLO araña
 < 250-chasquid replies your hello
 < 250-SIZE 35651584
 < 250-8BITMIME
@@ -146,8 +157,7 @@ func TestFallbackToIDNA(t *testing.T) {
 < 250 RCPT TO is fine
 `)
 
-	c := &Client{
-		Client: &smtp.Client{Text: textproto.NewConn(fake)}}
+	c := mustNewClient(t, fake)
 	if err := c.Hello("araña"); err != nil {
 		t.Fatalf("Hello failed: %v", err)
 	}
@@ -166,6 +176,38 @@ func TestFallbackToIDNA(t *testing.T) {
 	}
 }
 
+func TestLineTooLong(t *testing.T) {
+	// Fake the server sending a >2MiB reply.
+	dialog := `< 220 welcome
+> EHLO araña
+< 250 HELP
+> NOOP
+< 250 longreply:` + fmt.Sprintf("%2097152s", "x") + `:
+> NOOP
+< 250 ok
+`
+
+	fake, client := fakeDialog(dialog)
+
+	c := mustNewClient(t, fake)
+	if err := c.Hello("araña"); err != nil {
+		t.Fatalf("Hello failed: %v", err)
+	}
+
+	if err := c.Noop(); err != nil {
+		t.Errorf("Noop failed: %v", err)
+	}
+
+	if err := c.Noop(); err != io.EOF {
+		t.Errorf("Expected EOF, got: %v", err)
+	}
+
+	cmds := fake.Client()
+	if client != cmds {
+		t.Errorf("Got:\n%s\nExpected:\n%s", cmds, client)
+	}
+}
+
 type faker struct {
 	buf *bytes.Buffer
 	*bufio.ReadWriter
@@ -182,6 +224,8 @@ func (f faker) Client() string {
 	return f.buf.String()
 }
 
+var _ net.Conn = faker{}
+
 // Takes a dialog, returns the corresponding faker and expected client
 // messages.  Ideally we would check this interactively, and it's not that
 // difficult, but this is good enough for now.