client: introduce DataCommand

Closes: #189

Author: emersion

client.go

@@ -12,6 +12,7 @@ import (
 	"io"
 	"net"
 	"net/textproto"
+	"sort"
 	"strconv"
 	"strings"
 	"time"
@@ -540,85 +541,152 @@ func (c *Client) Rcpt(to string, opts *RcptOptions) error {
 	return nil
 }
 
-type dataCloser struct {
-	c *Client
-	io.WriteCloser
-	statusCb func(rcpt string, status *SMTPError)
-	closed   bool
+// DataCommand is a pending DATA command. DataCommand is an io.WriteCloser.
+// See Client.Data.
+type DataCommand struct {
+	client *Client
+	wc     io.WriteCloser
+
+	closeErr error
+}
+
+var _ io.WriteCloser = (*DataCommand)(nil)
+
+// Write implements io.Writer.
+func (cmd *DataCommand) Write(b []byte) (int, error) {
+	return cmd.wc.Write(b)
+}
+
+// Close implements io.Closer.
+func (cmd *DataCommand) Close() error {
+	var err error
+	if cmd.client.lmtp {
+		_, err = cmd.CloseWithLMTPResponse()
+	} else {
+		_, err = cmd.CloseWithResponse()
+	}
+	return err
 }
 
-func (d *dataCloser) Close() error {
-	if d.closed {
-		return fmt.Errorf("smtp: data writer closed twice")
+// CloseWithResponse is equivalent to Close, but also returns the server
+// response. It cannot be called when the LMTP protocol is used.
+//
+// If server returns an error, it will be of type *SMTPError.
+func (cmd *DataCommand) CloseWithResponse() (*DataResponse, error) {
+	if cmd.client.lmtp {
+		return nil, errors.New("smtp: CloseWithResponse used with an LMTP client")
 	}
 
-	if err := d.WriteCloser.Close(); err != nil {
-		return err
+	if err := cmd.close(); err != nil {
+		return nil, err
 	}
 
-	d.c.conn.SetDeadline(time.Now().Add(d.c.SubmissionTimeout))
-	defer d.c.conn.SetDeadline(time.Time{})
+	cmd.client.conn.SetDeadline(time.Now().Add(cmd.client.SubmissionTimeout))
+	defer cmd.client.conn.SetDeadline(time.Time{})
 
-	expectedResponses := len(d.c.rcpts)
-	if d.c.lmtp {
-		for expectedResponses > 0 {
-			rcpt := d.c.rcpts[len(d.c.rcpts)-expectedResponses]
-			if _, _, err := d.c.readResponse(250); err != nil {
-				if smtpErr, ok := err.(*SMTPError); ok {
-					if d.statusCb != nil {
-						d.statusCb(rcpt, smtpErr)
-					}
-				} else {
-					return err
+	_, msg, err := cmd.client.readResponse(250)
+	if err != nil {
+		cmd.closeErr = err
+		return nil, err
+	}
+
+	return &DataResponse{StatusText: msg}, nil
+}
+
+// CloseWithLMTPResponse is equivalent to Close, but also returns per-recipient
+// server responses. It can only be called when the LMTP protocol is used.
+//
+// If server returns an error, it will be of type LMTPDataError.
+func (cmd *DataCommand) CloseWithLMTPResponse() (map[string]*DataResponse, error) {
+	if !cmd.client.lmtp {
+		return nil, errors.New("smtp: CloseWithLMTPResponse used without an LMTP client")
+	}
+
+	if err := cmd.close(); err != nil {
+		return nil, err
+	}
+
+	cmd.client.conn.SetDeadline(time.Now().Add(cmd.client.SubmissionTimeout))
+	defer cmd.client.conn.SetDeadline(time.Time{})
+
+	resp := make(map[string]*DataResponse, len(cmd.client.rcpts))
+	lmtpErr := make(LMTPDataError, len(cmd.client.rcpts))
+	for i := 0; i < len(cmd.client.rcpts); i++ {
+		rcpt := cmd.client.rcpts[i]
+		_, msg, err := cmd.client.readResponse(250)
+		if err != nil {
+			if smtpErr, ok := err.(*SMTPError); ok {
+				lmtpErr[rcpt] = smtpErr
+			} else {
+				if len(lmtpErr) > 0 {
+					return resp, errors.Join(err, lmtpErr)
 				}
-			} else if d.statusCb != nil {
-				d.statusCb(rcpt, nil)
+				return resp, err
 			}
-			expectedResponses--
-		}
-	} else {
-		_, _, err := d.c.readResponse(250)
-		if err != nil {
-			return err
+		} else {
+			resp[rcpt] = &DataResponse{StatusText: msg}
 		}
 	}
 
-	d.closed = true
-	return nil
+	if len(lmtpErr) > 0 {
+		return resp, lmtpErr
+	}
+	return resp, nil
 }
 
-// Data issues a DATA command to the server and returns a writer that
-// can be used to write the mail headers and body. The caller should
-// close the writer before calling any more methods on c. A call to
-// Data must be preceded by one or more calls to Rcpt.
-//
-// If server returns an error, it will be of type *SMTPError.
-func (c *Client) Data() (io.WriteCloser, error) {
-	_, _, err := c.cmd(354, "DATA")
-	if err != nil {
-		return nil, err
+func (cmd *DataCommand) close() error {
+	if cmd.closeErr != nil {
+		return cmd.closeErr
 	}
-	return &dataCloser{c: c, WriteCloser: c.text.DotWriter()}, nil
+
+	if err := cmd.wc.Close(); err != nil {
+		cmd.closeErr = err
+		return err
+	}
+
+	cmd.closeErr = errors.New("smtp: data writer closed twice")
+	return nil
 }
 
-// LMTPData is the LMTP-specific version of the Data method. It accepts a callback
-// that will be called for each status response received from the server.
-//
-// Status callback will receive a SMTPError argument for each negative server
-// reply and nil for each positive reply. I/O errors will not be reported using
-// callback and instead will be returned by the Close method of io.WriteCloser.
-// Callback will be called for each successfull Rcpt call done before in the
-// same order.
-func (c *Client) LMTPData(statusCb func(rcpt string, status *SMTPError)) (io.WriteCloser, error) {
-	if !c.lmtp {
-		return nil, errors.New("smtp: not a LMTP client")
+// DataResponse is the response returned by a DATA command. See
+// DataCommand.CloseWithResponse.
+type DataResponse struct {
+	// StatusText is the status text returned by the server. It may contain
+	// tracking information.
+	StatusText string
+}
+
+// LMTPDataError is a collection of errors returned by an LMTP server for a
+// DATA command. It holds per-recipient errors.
+type LMTPDataError map[string]*SMTPError
+
+// Error implements error.
+func (lmtpErr LMTPDataError) Error() string {
+	return errors.Join(lmtpErr.Unwrap()...).Error()
+}
+
+// Unwrap returns all per-recipient errors returned by the server.
+func (lmtpErr LMTPDataError) Unwrap() []error {
+	l := make([]error, 0, len(lmtpErr))
+	for rcpt, smtpErr := range lmtpErr {
+		l = append(l, fmt.Errorf("<%v>: %w", rcpt, smtpErr))
 	}
+	sort.Slice(l, func(i, j int) bool {
+		return l[i].Error() < l[j].Error()
+	})
+	return l
+}
 
+// Data issues a DATA command to the server and returns a writer that
+// can be used to write the mail headers and body. The caller should
+// close the writer before calling any more methods on c. A call to
+// Data must be preceded by one or more calls to Rcpt.
+func (c *Client) Data() (*DataCommand, error) {
 	_, _, err := c.cmd(354, "DATA")
 	if err != nil {
 		return nil, err
 	}
-	return &dataCloser{c: c, WriteCloser: c.text.DotWriter(), statusCb: statusCb}, nil
+	return &DataCommand{client: c, wc: c.text.DotWriter()}, nil
 }
 
 // SendMail will use an existing connection to send an email from

client_test.go

@@ -150,8 +150,10 @@ Goodbye.`
 	if _, err := w.Write([]byte(msg)); err != nil {
 		t.Fatalf("Data write failed: %s", err)
 	}
-	if err := w.Close(); err != nil {
+	if resp, err := w.CloseWithResponse(); err != nil {
 		t.Fatalf("Bad data response: %s", err)
+	} else if want := "Data OK"; resp.StatusText != want {
+		t.Errorf("Bad data status text: got %q, want %q", resp.StatusText, want)
 	}
 
 	if err := c.Quit(); err != nil {
@@ -916,35 +918,33 @@ Line 1
 .Leading dot line .
 Goodbye.`
 
-	rcpts := []string{}
-	errors := []*SMTPError{}
-
-	w, err := c.LMTPData(func(rcpt string, status *SMTPError) {
-		rcpts = append(rcpts, rcpt)
-		errors = append(errors, status)
-	})
+	w, err := c.Data()
 	if err != nil {
 		t.Fatalf("DATA failed: %s", err)
 	}
 	if _, err := w.Write([]byte(msg)); err != nil {
 		t.Fatalf("Data write failed: %s", err)
 	}
-	if err := w.Close(); err != nil {
-		t.Fatalf("Bad data response: %s", err)
+	resp, err := w.CloseWithLMTPResponse()
+
+	var lmtpErr LMTPDataError
+	if !errors.As(err, &lmtpErr) {
+		t.Fatalf("Want error of type LMTPDataError")
 	}
 
-	if !reflect.DeepEqual(rcpts, []string{"golang-nuts@googlegroups.com", "golang-not-nuts@googlegroups.com"}) {
-		t.Fatal("Status callbacks called for wrong recipients:", rcpts)
+	wantResp := map[string]*DataResponse{
+		"golang-nuts@googlegroups.com": {StatusText: "This recipient is fine"},
 	}
 
-	if len(errors) != 2 {
-		t.Fatalf("Wrong amount of status callback calls: %v", len(errors))
+	if !reflect.DeepEqual(resp, wantResp) {
+		t.Fatalf("resp = %v, want %v", resp, wantResp)
 	}
-	if errors[0] != nil {
-		t.Fatalf("Unexpected error status for the first recipient: %v", errors[0])
+
+	if len(lmtpErr) != 1 {
+		t.Fatalf("len(lmtpErr) = %v, want 1", len(lmtpErr))
 	}
-	if errors[1] == nil {
-		t.Fatalf("Unexpected success status for the second recipient")
+	if lmtpErr["golang-not-nuts@googlegroups.com"] == nil {
+		t.Fatalf("Want error for second recipient")
 	}
 
 	if err := c.Quit(); err != nil {