From 66c8b9ae4cce7d12b38ef020bc84dfffb5ffd740 Mon Sep 17 00:00:00 2001
From: Alexandre Cesaro <alexandre.cesaro@gmail.com>
Date: Wed, 8 Jul 2015 00:26:19 +0200
Subject: [PATCH] Changed body content type from []byte to func(io.Writer)
 error

It allows streaming the message body when sending an email instead
of buffering it to memory.
---
 message.go      | 34 +++++++++++++++++++---------------
 message_test.go |  6 ++++--
 writeto.go      | 15 +++++++++------
 3 files changed, 32 insertions(+), 23 deletions(-)

diff --git a/message.go b/message.go
index 6612e97..a075186 100644
--- a/message.go
+++ b/message.go
@@ -28,7 +28,7 @@ type header map[string][]string
 
 type part struct {
 	contentType string
-	body        *bytes.Buffer
+	copier      func(io.Writer) error
 }
 
 // NewMessage creates a new message. It uses UTF-8 and quoted-printable encoding
@@ -170,12 +170,15 @@ func (msg *Message) SetBody(contentType, body string) {
 	msg.parts = []part{
 		part{
 			contentType: contentType,
-			body:        bytes.NewBufferString(body),
+			copier: func(w io.Writer) error {
+				_, err := io.WriteString(w, body)
+				return err
+			},
 		},
 	}
 }
 
-// AddAlternative adds an alternative body to the message. Commonly used to
+// AddAlternative adds an alternative part to the message. Commonly used to
 // send HTML emails that default to the plain text version for backward
 // compatibility.
 //
@@ -189,29 +192,30 @@ func (msg *Message) AddAlternative(contentType, body string) {
 	msg.parts = append(msg.parts,
 		part{
 			contentType: contentType,
-			body:        bytes.NewBufferString(body),
+			copier: func(w io.Writer) error {
+				_, err := io.WriteString(w, body)
+				return err
+			},
 		},
 	)
 }
 
-// GetBodyWriter gets a writer that writes to the body. It can be useful with
-// the templates from packages text/template or html/template.
+// AddAlternativeWriter adds an alternative part to the message. It can be
+// useful with the text/template and html/template packages.
 //
 // Example:
 //
-//	w := msg.GetBodyWriter("text/plain")
 //	t := template.Must(template.New("example").Parse("Hello {{.}}!"))
-//	t.Execute(w, "Bob")
-func (msg *Message) GetBodyWriter(contentType string) io.Writer {
-	buf := new(bytes.Buffer)
-	msg.parts = append(msg.parts,
+//	msg.AddAlternativeWriter("text/plain", func(w io.Writer) error {
+//		return t.Execute(w, "Bob")
+//	})
+func (msg *Message) AddAlternativeWriter(contentType string, f func(io.Writer) error) {
+	msg.parts = []part{
 		part{
 			contentType: contentType,
-			body:        buf,
+			copier:      f,
 		},
-	)
-
-	return buf
+	}
 }
 
 // A File represents a file that can be attached or embedded in an email.
diff --git a/message_test.go b/message_test.go
index 5a40e2d..7cb7ea2 100644
--- a/message_test.go
+++ b/message_test.go
@@ -66,8 +66,10 @@ func TestBodyWriter(t *testing.T) {
 	msg := NewMessage()
 	msg.SetHeader("From", "from@example.com")
 	msg.SetHeader("To", "to@example.com")
-	w := msg.GetBodyWriter("text/plain")
-	w.Write([]byte("Test message"))
+	msg.AddAlternativeWriter("text/plain", func(w io.Writer) error {
+		_, err := w.Write([]byte("Test message"))
+		return err
+	})
 
 	want := &message{
 		from: "from@example.com",
diff --git a/writeto.go b/writeto.go
index 1d73319..15c200a 100644
--- a/writeto.go
+++ b/writeto.go
@@ -42,7 +42,7 @@ func (w *messageWriter) writeMessage(msg *Message) {
 			"Content-Type":              []string{contentType},
 			"Content-Transfer-Encoding": []string{string(msg.encoding)},
 		})
-		w.writeBody(part.body.Bytes(), msg.encoding)
+		w.writeBody(part.copier, msg.encoding)
 	}
 	if msg.hasAlternativePart() {
 		w.closeMultipart()
@@ -123,7 +123,10 @@ func (w *messageWriter) addFiles(files []*File, isAttachment bool) {
 			}
 		}
 		w.writeHeaders(h)
-		w.writeBody(f.Content, Base64)
+		w.writeBody(func(w io.Writer) error {
+			_, err := w.Write(f.Content)
+			return err
+		}, Base64)
 	}
 }
 
@@ -175,7 +178,7 @@ func (w *messageWriter) writeHeaders(h map[string][]string) {
 	}
 }
 
-func (w *messageWriter) writeBody(body []byte, enc Encoding) {
+func (w *messageWriter) writeBody(f func(io.Writer) error, enc Encoding) {
 	var subWriter io.Writer
 	if w.depth == 0 {
 		w.writeString("\r\n")
@@ -186,13 +189,13 @@ func (w *messageWriter) writeBody(body []byte, enc Encoding) {
 
 	if enc == Base64 {
 		wc := base64.NewEncoder(base64.StdEncoding, newBase64LineWriter(subWriter))
-		wc.Write(body)
+		w.err = f(wc)
 		wc.Close()
 	} else if enc == Unencoded {
-		subWriter.Write(body)
+		w.err = f(subWriter)
 	} else {
 		wc := quotedprintable.NewWriter(subWriter)
-		wc.Write(body)
+		w.err = f(wc)
 		wc.Close()
 	}
 }