lib/email: make Message Pack works with single text or HTML part

Previouly, the Pack method generate multipart/alternative message only.

Since the Message now can set the body text and HTML, without using
NewMultipart, the Pack method need to be able to accomodate this.

6 files changed, 274 insertions, 34 deletions

diff --git a/lib/email/body.go b/lib/email/body.go
index eda591fc..e31a8fea 100644
--- a/lib/email/body.go
+++ b/lib/email/body.go
@@ -102,6 +102,22 @@ func (body *Body) Add(mime *MIME) {
 }
 
 //
+// getPart get the body part by top and sub content type.
+//
+func (body *Body) getPart(top, sub []byte) (mime *MIME) {
+	for _, mime = range body.Parts {
+		if !bytes.Equal(mime.contentType.Top, top) {
+			continue
+		}
+		if !bytes.Equal(mime.contentType.Sub, sub) {
+			continue
+		}
+		return mime
+	}
+	return nil
+}
+
+//
 // Set replace the MIME content-type with new one, if its exist; otherwise
 // append it.
 //
diff --git a/lib/email/contenttype.go b/lib/email/contenttype.go
index 05c80ca8..b165961e 100644
--- a/lib/email/contenttype.go
+++ b/lib/email/contenttype.go
@@ -13,6 +13,12 @@ import (
 	libio "github.com/shuLhan/share/lib/io"
 )
 
+var (
+	topText  = []byte("text")
+	subPlain = []byte("plain")
+	subHtml  = []byte("html")
+)
+
 //
 // ContentType represent MIME header "Content-Type" field.
 //
@@ -166,7 +172,7 @@ func (ct *ContentType) SetBoundary(boundary []byte) {
 }
 
 //
-// String return text representation of this instance.
+// String return text representation of content type with its parameters.
 //
 func (ct *ContentType) String() string {
 	var sb strings.Builder
diff --git a/lib/email/header.go b/lib/email/header.go
index 39ceda45..d16030be 100644
--- a/lib/email/header.go
+++ b/lib/email/header.go
@@ -7,6 +7,7 @@ package email
 import (
 	"bytes"
 	"fmt"
+	"io"
 	"log"
 
 	"github.com/shuLhan/share/lib/ascii"
@@ -279,3 +280,27 @@ func (hdr *Header) SetMultipart() (err error) {
 
 	return nil
 }
+
+//
+// WriteTo the header into w.
+// The header does not end with an empty line to allow multiple Header
+// written multiple times.
+//
+func (hdr *Header) WriteTo(w io.Writer) (n int, err error) {
+	var (
+		f *Field
+		m int
+	)
+	for _, f = range hdr.fields {
+		if f.Type == FieldTypeContentType {
+			m, err = fmt.Fprintf(w, "%s: %s\r\n", f.Name, f.ContentType.String())
+		} else {
+			m, err = fmt.Fprintf(w, "%s: %s", f.Name, f.Value)
+		}
+		if err != nil {
+			return n, err
+		}
+		n += m
+	}
+	return n, nil
+}
diff --git a/lib/email/message.go b/lib/email/message.go
index dd2e806a..df6baf20 100644
--- a/lib/email/message.go
+++ b/lib/email/message.go
@@ -489,38 +489,121 @@ func (msg *Message) CanonHeader(subHeader *Header, dkimField *Field) []byte {
 //
 // Pack the message for sending.
 //
-func (msg *Message) Pack() (out []byte) {
-	var buf bytes.Buffer
+// The message content type is automatically set based on the Body parts.
+// If the Body only contain text part, the generated content-type will be set
+// to text/plain.
+// If the Body only contain HTML part, the generated content-type will be set
+// to text/html.
+// If both the text and HTML parts exist, the generated content-type will be
+// set to multipart/alternative.
+//
+func (msg *Message) Pack() (out []byte, err error) {
+	var (
+		logp = "Pack"
+	)
 
-	boundary := msg.Header.Boundary()
+	if len(msg.Body.Parts) == 0 {
+		return nil, fmt.Errorf("%s: empty body", logp)
+	}
 
-	for _, f := range msg.Header.fields {
-		if f.Type == FieldTypeContentType {
-			fmt.Fprintf(&buf, "%s: %s\r\n", f.Name, f.ContentType.String())
-		} else {
-			fmt.Fprintf(&buf, "%s: %s", f.Name, f.Value)
+	// TODO: check date, from, to, subject.
+
+	if len(msg.Body.Parts) > 1 {
+		out, err = msg.packMultipartAlternative()
+		if err != nil {
+			return nil, fmt.Errorf("%s: %w", logp, err)
 		}
+		return out, nil
 	}
+
+	out, err = msg.packSingle()
+	if err != nil {
+		return nil, fmt.Errorf("%s: %w", logp, err)
+	}
+
+	return out, nil
+}
+
+func (msg *Message) packMultipartAlternative() (out []byte, err error) {
+	var (
+		mime     *MIME
+		buf      bytes.Buffer
+		boundary []byte
+	)
+
+	boundary = msg.Header.Boundary()
+	if len(boundary) == 0 {
+		// Set the boundary has not been set, so generated one.
+		err = msg.Header.SetMultipart()
+		if err != nil {
+			return nil, err
+		}
+		boundary = msg.Header.Boundary()
+	}
+
+	_, err = msg.Header.WriteTo(&buf)
+	if err != nil {
+		return nil, err
+	}
+
 	buf.WriteString("\r\n")
-	for _, mime := range msg.Body.Parts {
-		if len(boundary) > 0 {
-			fmt.Fprintf(&buf, "--%s\r\n", boundary)
+
+	// Make sure the text part written first.
+	mime = msg.Body.getPart(topText, subPlain)
+	if mime != nil {
+		fmt.Fprintf(&buf, "--%s\r\n", boundary)
+		_, err = mime.WriteTo(&buf)
+		if err != nil {
+			return nil, err
 		}
-		for _, f := range mime.Header.fields {
-			if f.Type == FieldTypeContentType {
-				fmt.Fprintf(&buf, "%s: %s\r\n", f.Name,
-					f.ContentType.String())
-			} else {
-				fmt.Fprintf(&buf, "%s: %s", f.Name, f.Value)
-			}
+	}
+
+	mime = msg.Body.getPart(topText, subHtml)
+	if mime != nil {
+		fmt.Fprintf(&buf, "--%s\r\n", boundary)
+		_, err = mime.WriteTo(&buf)
+		if err != nil {
+			return nil, err
 		}
-		buf.WriteString("\r\n")
-		buf.Write(mime.Content)
 	}
-	if len(boundary) > 0 {
-		fmt.Fprintf(&buf, "--%s--\r\n", boundary)
+
+	// Write the rest of the parts, but skip the plain and HTML parts.
+	for _, mime = range msg.Body.Parts {
+		if mime.isContentType(topText, subPlain) {
+			continue
+		}
+		if mime.isContentType(topText, subHtml) {
+			continue
+		}
+
+		fmt.Fprintf(&buf, "--%s\r\n", boundary)
+		_, err = mime.WriteTo(&buf)
+		if err != nil {
+			return nil, err
+		}
 	}
-	return buf.Bytes()
+	fmt.Fprintf(&buf, "--%s--\r\n", boundary)
+	return buf.Bytes(), nil
+}
+
+func (msg *Message) packSingle() (out []byte, err error) {
+	var (
+		mime = msg.Body.Parts[0]
+
+		buf bytes.Buffer
+	)
+
+	_, err = msg.Header.WriteTo(&buf)
+	if err != nil {
+		return nil, err
+	}
+
+	_, err = mime.WriteTo(&buf)
+	if err != nil {
+		return nil, err
+	}
+
+	return buf.Bytes(), nil
 }
 
 //
diff --git a/lib/email/message_test.go b/lib/email/message_test.go
index abc0b137..faf9827a 100644
--- a/lib/email/message_test.go
+++ b/lib/email/message_test.go
@@ -14,6 +14,12 @@ import (
 )
 
 func TestNewMultipart(t *testing.T) {
+	var (
+		gotMsg *Message
+		msgb   []byte
+		err    error
+	)
+
 	dateInUtc = true
 	Epoch = func() int64 {
 		return 1645811431
@@ -57,12 +63,17 @@ func TestNewMultipart(t *testing.T) {
 	}}
 
 	for _, c := range cases {
-		gotMsg, err := NewMultipart(c.from, c.to, c.subject, c.bodyText, c.bodyHTML)
+		gotMsg, err = NewMultipart(c.from, c.to, c.subject, c.bodyText, c.bodyHTML)
 		if err != nil {
 			t.Fatal(err)
 		}
 
-		test.Assert(t, "NewMultipart", c.expMsg, string(gotMsg.Pack()))
+		msgb, err = gotMsg.Pack()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		test.Assert(t, "NewMultipart", c.expMsg, string(msgb))
 	}
 }
 
@@ -306,6 +317,64 @@ func TestMessageDKIMSign(t *testing.T) {
 	}
 }
 
+func TestMessage_packSingle(t *testing.T) {
+	type testCase struct {
+		desc     string
+		exp      string
+		bodyText []byte
+		bodyHtml []byte
+	}
+
+	var (
+		msg   Message
+		err   error
+		cases []testCase
+		got   []byte
+	)
+
+	cases = []testCase{{
+		desc:     "With body text",
+		bodyText: []byte(`this is a body text`),
+		exp: "" +
+			"content-type: text/plain; charset=\"utf-8\"\r\n" +
+			"mime-version: 1.0\r\n" +
+			"content-transfer-encoding: quoted-printable\r\n" +
+			"\r\n" +
+			"this is a body text\r\n",
+	}, {
+		desc:     "With body HTML",
+		bodyHtml: []byte(`<p>this is an HTML body</p>`),
+		exp: "" +
+			"content-type: text/html; charset=\"utf-8\"\r\n" +
+			"mime-version: 1.0\r\n" +
+			"content-transfer-encoding: quoted-printable\r\n" +
+			"\r\n" +
+			"<p>this is an HTML body</p>\r\n",
+	}}
+
+	for _, c := range cases {
+		msg.Body.Parts = nil
+
+		if len(c.bodyText) > 0 {
+			err = msg.SetBodyText(c.bodyText)
+			if err != nil {
+				t.Fatal(err)
+			}
+		} else {
+			err = msg.SetBodyHtml(c.bodyHtml)
+			if err != nil {
+				t.Fatal(err)
+			}
+		}
+		got, err = msg.Pack()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		test.Assert(t, c.desc, c.exp, string(got))
+	}
+}
+
 func TestMessage_SetBodyText(t *testing.T) {
 	var (
 		msg Message
@@ -313,12 +382,12 @@ func TestMessage_SetBodyText(t *testing.T) {
 	)
 	cases := []struct {
 		desc    string
-		content []byte
 		expMsg  string
+		content []byte
 	}{{
 		desc:    "With empty Body",
 		content: []byte("text body"),
-		expMsg: "\r\n" +
+		expMsg: "" +
 			"content-type: text/plain; charset=\"utf-8\"\r\n" +
 			"mime-version: 1.0\r\n" +
 			"content-transfer-encoding: quoted-printable\r\n" +
@@ -327,7 +396,7 @@ func TestMessage_SetBodyText(t *testing.T) {
 	}, {
 		desc:    "With new text",
 		content: []byte("new text body"),
-		expMsg: "\r\n" +
+		expMsg: "" +
 			"content-type: text/plain; charset=\"utf-8\"\r\n" +
 			"mime-version: 1.0\r\n" +
 			"content-transfer-encoding: quoted-printable\r\n" +
@@ -341,7 +410,12 @@ func TestMessage_SetBodyText(t *testing.T) {
 			t.Fatal(err)
 		}
 
-		test.Assert(t, c.desc, string(c.expMsg), string(msg.Pack()))
+		msgb, err := msg.Pack()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		test.Assert(t, c.desc, string(c.expMsg), string(msgb))
 	}
 }
 
@@ -434,18 +508,18 @@ func TestMessage_SetSubject(t *testing.T) {
 		expMsg  string
 	}{{
 		subject: "a subject",
-		expMsg:  "subject: a subject\r\n\r\n",
+		expMsg:  "subject:a subject\r\n\r\n",
 	}, {
-		expMsg: "subject: a subject\r\n\r\n",
+		expMsg: "subject:a subject\r\n\r\n",
 	}, {
 		subject: "new subject",
-		expMsg:  "subject: new subject\r\n\r\n",
+		expMsg:  "subject:new subject\r\n\r\n",
 	}}
 
 	for _, c := range cases {
 		msg.SetSubject(c.subject)
 
-		test.Assert(t, "SetSubject", c.expMsg, string(msg.Pack()))
+		test.Assert(t, "SetSubject", c.expMsg, msg.String())
 	}
 }
 
diff --git a/lib/email/mime.go b/lib/email/mime.go
index 12a77ce0..c6da311b 100644
--- a/lib/email/mime.go
+++ b/lib/email/mime.go
@@ -7,6 +7,7 @@ package email
 import (
 	"bytes"
 	"errors"
+	"io"
 	"mime/quotedprintable"
 	"strings"
 
@@ -145,6 +146,13 @@ func ParseBodyPart(raw, boundary []byte) (mime *MIME, rest []byte, err error) {
 	return mime, rest, err
 }
 
+func (mime *MIME) isContentType(top, sub []byte) bool {
+	if bytes.Equal(mime.contentType.Top, top) {
+		return bytes.Equal(mime.contentType.Sub, sub)
+	}
+	return false
+}
+
 //
 // String return string representation of MIME object.
 //
@@ -160,3 +168,31 @@ func (mime *MIME) String() string {
 
 	return sb.String()
 }
+
+//
+// WriteTo write the MIME header and content into Writer w.
+//
+func (mime *MIME) WriteTo(w io.Writer) (n int, err error) {
+	var (
+		m int
+	)
+	m, err = mime.Header.WriteTo(w)
+	if err != nil {
+		return n, err
+	}
+	n += m
+
+	m, err = w.Write([]byte("\r\n"))
+	if err != nil {
+		return n, err
+	}
+	n += m
+
+	m, err = w.Write(mime.Content)
+	if err != nil {
+		return n, err
+	}
+	n += m
+
+	return n, nil
+}