From b7f245fe64f0976c3c3fe157487f392639b754c5 Mon Sep 17 00:00:00 2001
From: Simon Ser <contact@emersion.fr>
Date: Wed, 24 Apr 2024 16:32:31 +0200
Subject: [PATCH] client: introduce DataCommand

Closes: https://github.com/emersion/go-smtp/issues/189
---
 client.go | 44 ++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 42 insertions(+), 2 deletions(-)

diff --git a/client.go b/client.go
index e17fdd4..12782c4 100644
--- a/client.go
+++ b/client.go
@@ -530,6 +530,46 @@ func (c *Client) Rcpt(to string, opts *RcptOptions) error {
 	return nil
 }
 
+type DataCommand struct {
+	c  *Client
+	wc io.WriteCloser
+
+	closeErr   error
+	statusText string
+}
+
+func (cmd *DataCommand) Write(b []byte) (int, error) {
+	return cmd.wc.Write(b)
+}
+
+func (cmd *DataCommand) Close() error {
+	if cmd.closeErr != nil {
+		return cmd.closeErr
+	}
+
+	if err := cmd.wc.Close(); err != nil {
+		cmd.closeErr = err
+		return err
+	}
+
+	cmd.c.conn.SetDeadline(time.Now().Add(cmd.c.SubmissionTimeout))
+	defer cmd.c.conn.SetDeadline(time.Time{})
+
+	_, msg, err := cmd.c.readResponse(250)
+	if err != nil {
+		cmd.closeErr = err
+		return err
+	}
+
+	cmd.statusText = msg
+	cmd.closeErr = errors.New("smtp: data writer closed twice")
+	return nil
+}
+
+func (cmd *DataCommand) StatusText() string {
+	return cmd.statusText
+}
+
 type dataCloser struct {
 	c *Client
 	io.WriteCloser
@@ -583,12 +623,12 @@ func (d *dataCloser) Close() error {
 // 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) {
+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()}, nil
+	return &DataCommand{c: c, wc: c.text.DotWriter()}, nil
 }
 
 // LMTPData is the LMTP-specific version of the Data method. It accepts a callback