added support for multiple values for same header key (#80)

* added support for multiple values for same header key * style: direct err return Fixes a lint failure by returning the err directly, rather than via a conditional. --------- Co-authored-by: Dom Dwyer <[email protected]>
domodwyer · Mar 26, 2023 · 87b7a70 · 87b7a70
1 parent 8cabd8d
commit 87b7a70
Show file tree

Hide file tree

Showing 5 changed files with 53 additions and 28 deletions.
diff --git a/mailyak.go b/mailyak.go
@@ -24,7 +24,7 @@ type MailYak struct {
 	fromAddr       string
 	fromName       string
 	replyTo        string
-	headers        map[string]string // arbitrary headers
+	headers        map[string][]string // arbitrary headers
 	attachments    []attachment
 	trimRegex      *regexp.Regexp
 	auth           smtp.Auth
@@ -54,7 +54,7 @@ const mailDateFormat = time.RFC1123Z
 // connection, or to provide a custom tls.Config, use NewWithTLS() instead.
 func New(host string, auth smtp.Auth) *MailYak {
 	return &MailYak{
-		headers:        map[string]string{},
+		headers:        map[string][]string{},
 		host:           host,
 		auth:           auth,
 		sender:         newSenderWithStartTLS(host),

diff --git a/mailyak_test.go b/mailyak_test.go
@@ -28,7 +28,7 @@ func TestMailYakStringer(t *testing.T) {
 
 	mail.date = "a date"
 
-	want := "&MailYak{date: \"a date\", from: \"[email protected]\", fromName: \"From Example\", html: 31 bytes, plain: 42 bytes, toAddrs: [[email protected]], bccAddrs: [[email protected] [email protected]], subject: \"Test subject\", Precedence: \"bulk\", host: \"mail.host.com:25\", attachments (2): [{filename: test.html} {filename: test2.html}], auth set: true, explicit tls: false}"
+	want := "&MailYak{date: \"a date\", from: \"[email protected]\", fromName: \"From Example\", html: 31 bytes, plain: 42 bytes, toAddrs: [[email protected]], bccAddrs: [[email protected] [email protected]], subject: \"Test subject\", Precedence: [\"bulk\"], host: \"mail.host.com:25\", attachments (2): [{filename: test.html} {filename: test2.html}], auth set: true, explicit tls: false}"
 	got := fmt.Sprintf("%+v", mail)
 	if got != want {
 		t.Errorf("MailYak.String() = %v, want %v", got, want)

diff --git a/mime.go b/mime.go
@@ -76,11 +76,7 @@ func (m *MailYak) buildMimeWithBoundaries(w io.Writer, mb, ab string) error {
 		return err
 	}
 
-	if err := mixed.Close(); err != nil {
-		return err
-	}
-
-	return nil
+	return mixed.Close()
 }
 
 // writeHeaders writes the Mime-Version, Date, Reply-To, From, To and Subject headers,
@@ -118,8 +114,10 @@ func (m *MailYak) writeHeaders(w io.Writer) error {
 		fmt.Fprintf(w, "BCC: %s\r\n", commaSeparatedBCCAddrs)
 	}
 
-	for k, v := range m.headers {
-		fmt.Fprintf(w, "%s: %s\r\n", k, v)
+	for k, values := range m.headers {
+		for _, v := range values {
+			fmt.Fprintf(w, "%s: %s\r\n", k, v)
+		}
 	}
 
 	return nil

diff --git a/setters.go b/setters.go
@@ -1,6 +1,8 @@
 package mailyak
 
-import "mime"
+import (
+	"mime"
+)
 
 // To sets a list of recipient addresses.
 //
@@ -136,13 +138,27 @@ func (m *MailYak) Subject(sub string) {
 }
 
 // AddHeader adds an arbitrary email header.
+// It appends to any existing values associated with key.
 //
 // If value contains non-ASCII characters, it is Q-encoded according to RFC1342.
 // As always, validate any user input before adding it to a message, as this
 // method may enable an attacker to override the standard headers and, for
 // example, BCC themselves in a password reset email to a different user.
 func (m *MailYak) AddHeader(name, value string) {
-	m.headers[m.trimRegex.ReplaceAllString(name, "")] = mime.QEncoding.Encode("UTF-8", m.trimRegex.ReplaceAllString(value, ""))
+	key := m.trimRegex.ReplaceAllString(name, "")
+	m.headers[key] = append(m.headers[key], mime.QEncoding.Encode("UTF-8", m.trimRegex.ReplaceAllString(value, "")))
+}
+
+// SetHeader sets the header entries associated with key to
+// the single element value. It replaces any existing
+// values associated with key.
+//
+// If value contains non-ASCII characters, it is Q-encoded according to RFC1342.
+// As always, validate any user input before adding it to a message, as this
+// method may enable an attacker to override the standard headers and, for
+// example, BCC themselves in a password reset email to a different user.
+func (m *MailYak) SetHeader(name, value string) {
+	m.headers[m.trimRegex.ReplaceAllString(name, "")] = []string{mime.QEncoding.Encode("UTF-8", m.trimRegex.ReplaceAllString(value, ""))}
 }
 
 // LocalName sets the sender domain name.

diff --git a/setters_test.go b/setters_test.go
@@ -256,28 +256,37 @@ func TestMailYakAddHeader(t *testing.T) {
 		// Test description.
 		name string
 		// Parameters.
-		from map[string]string
+		from map[string][]string
 		// Want
-		want map[string]string
+		want map[string][]string
 	}{
 		{
 			"ASCII",
-			map[string]string{
-				"List-Unsubscribe": "http://example.com",
-				"X-NASTY":          "true\r\nBcc: [email protected]",
+			map[string][]string{
+				"List-Unsubscribe": {"http://example.com"},
+				"X-NASTY":          {"true\r\nBcc: [email protected]"},
 			},
-			map[string]string{
-				"List-Unsubscribe": "http://example.com",
-				"X-NASTY":          "trueBcc: [email protected]",
+			map[string][]string{
+				"List-Unsubscribe": {"http://example.com"},
+				"X-NASTY":          {"trueBcc: [email protected]"},
 			},
 		},
 		{
 			"Q-encoded",
-			map[string]string{
-				"X-BEETHOVEN": "für Elise",
+			map[string][]string{
+				"X-BEETHOVEN": {"für Elise"},
 			},
-			map[string]string{
-				"X-BEETHOVEN": "=?UTF-8?q?f=C3=BCr_Elise?=",
+			map[string][]string{
+				"X-BEETHOVEN": {"=?UTF-8?q?f=C3=BCr_Elise?="},
+			},
+		},
+		{
+			"Multi",
+			map[string][]string{
+				"X-MailGun-Tag": {"marketing", "transactional"},
+			},
+			map[string][]string{
+				"X-MailGun-Tag": {"marketing", "transactional"},
 			},
 		},
 	}
@@ -287,12 +296,14 @@ func TestMailYakAddHeader(t *testing.T) {
 			t.Parallel()
 
 			m := &MailYak{
-				headers:   map[string]string{},
+				headers:   map[string][]string{},
 				trimRegex: regexp.MustCompile("\r?\n"),
 			}
 
-			for k, v := range tt.from {
-				m.AddHeader(k, v)
+			for k, values := range tt.from {
+				for _, v := range values {
+					m.AddHeader(k, v)
+				}
 			}
 
 			if !reflect.DeepEqual(m.headers, tt.want) {
@@ -330,7 +341,7 @@ func TestMailYakLocalName(t *testing.T) {
 			t.Parallel()
 
 			m := &MailYak{
-				headers:   map[string]string{},
+				headers:   map[string][]string{},
 				trimRegex: regexp.MustCompile("\r?\n"),
 			}