From c20b44175667d6d03d86caaa132ad405f021c3fb Mon Sep 17 00:00:00 2001
From: "fox.cpp" <fox.cpp@disroot.org>
Date: Fri, 17 Jul 2020 15:21:15 +0300
Subject: [PATCH] client: Implement command timeouts

---
 client.go      | 30 +++++++++++++++++++++++++++++-
 client_test.go |  6 +++---
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/client.go b/client.go
index 0b88a08..b189017 100644
--- a/client.go
+++ b/client.go
@@ -14,6 +14,7 @@ import (
 	"net/textproto"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/emersion/go-sasl"
 )
@@ -38,6 +39,12 @@ type Client struct {
 	didHello   bool     // whether we've said HELO/EHLO/LHLO
 	helloError error    // the error from the hello
 	rcpts      []string // recipients accumulated for the current session
+
+	// Time to wait for command responses (this includes 3xx reply to DATA).
+	CommandTimeout time.Duration
+
+	// Time to wait for responses after final dot.
+	SubmissionTimeout time.Duration
 }
 
 // Dial returns a new Client connected to an SMTP server at addr.
@@ -89,7 +96,20 @@ func NewClient(conn net.Conn, host string) (*Client, error) {
 		return nil, err
 	}
 	_, isTLS := conn.(*tls.Conn)
-	c := &Client{Text: text, conn: conn, serverName: host, localName: "localhost", tls: isTLS}
+	c := &Client{
+		Text:       text,
+		conn:       conn,
+		serverName: host,
+		localName:  "localhost",
+		tls:        isTLS,
+		// As recommended by RFC 5321. For DATA command reply (3xx one) RFC
+		// recommends a slightly shorter timeout but we do not bother
+		// differentiating these.
+		CommandTimeout: 5 * time.Minute,
+		// 10 minutes + 2 minute buffer in case the server is doing transparent
+		// forwarding and also follows recommended timeouts.
+		SubmissionTimeout: 12 * time.Minute,
+	}
 	return c, nil
 }
 
@@ -142,6 +162,9 @@ func (c *Client) Hello(localName string) error {
 // cmd is a convenience function that sends a command and returns the response
 // textproto.Error returned by c.Text.ReadResponse is converted into SMTPError.
 func (c *Client) cmd(expectCode int, format string, args ...interface{}) (int, string, error) {
+	c.conn.SetDeadline(time.Now().Add(c.CommandTimeout))
+	defer c.conn.SetDeadline(time.Time{})
+
 	id, err := c.Text.Cmd(format, args...)
 	if err != nil {
 		return 0, "", err
@@ -174,6 +197,7 @@ func (c *Client) ehlo() error {
 	if c.lmtp {
 		cmd = "LHLO"
 	}
+
 	_, msg, err := c.cmd(250, "%s %s", cmd, c.localName)
 	if err != nil {
 		return err
@@ -375,6 +399,10 @@ type dataCloser struct {
 
 func (d *dataCloser) Close() error {
 	d.WriteCloser.Close()
+
+	d.c.conn.SetDeadline(time.Now().Add(d.c.SubmissionTimeout))
+	defer d.c.conn.SetDeadline(time.Time{})
+
 	expectedResponses := len(d.c.rcpts)
 	if d.c.lmtp {
 		for expectedResponses > 0 {
diff --git a/client_test.go b/client_test.go
index 4a7793e..cb12b7d 100644
--- a/client_test.go
+++ b/client_test.go
@@ -80,7 +80,7 @@ func TestBasic(t *testing.T) {
 	bcmdbuf := bufio.NewWriter(&cmdbuf)
 	var fake faker
 	fake.ReadWriter = bufio.NewReadWriter(bufio.NewReader(strings.NewReader(server)), bcmdbuf)
-	c := &Client{Text: textproto.NewConn(fake), localName: "localhost"}
+	c := &Client{Text: textproto.NewConn(fake), conn: fake, localName: "localhost"}
 
 	if err := c.helo(); err != nil {
 		t.Fatalf("HELO failed: %s", err)
@@ -912,7 +912,7 @@ func TestLMTP(t *testing.T) {
 	bcmdbuf := bufio.NewWriter(&cmdbuf)
 	var fake faker
 	fake.ReadWriter = bufio.NewReadWriter(bufio.NewReader(strings.NewReader(server)), bcmdbuf)
-	c := &Client{Text: textproto.NewConn(fake), lmtp: true}
+	c := &Client{Text: textproto.NewConn(fake), conn: fake, lmtp: true}
 
 	if err := c.Hello("localhost"); err != nil {
 		t.Fatalf("LHLO failed: %s", err)
@@ -997,7 +997,7 @@ func TestLMTPData(t *testing.T) {
 	bcmdbuf := bufio.NewWriter(&cmdbuf)
 	var fake faker
 	fake.ReadWriter = bufio.NewReadWriter(bufio.NewReader(strings.NewReader(server)), bcmdbuf)
-	c := &Client{Text: textproto.NewConn(fake), lmtp: true}
+	c := &Client{Text: textproto.NewConn(fake), conn: fake, lmtp: true}
 
 	if err := c.Hello("localhost"); err != nil {
 		t.Fatalf("LHLO failed: %s", err)