From aafd2a20d9d53f0100f0a8d4b1d52805536fc718 Mon Sep 17 00:00:00 2001
From: Ralph Slooten <axllent@gmail.com>
Date: Sun, 23 Feb 2025 11:14:47 +1300
Subject: [PATCH] Chore: Minor speed & memory improvements when storing
 messages

---
 internal/smtpd/main.go    | 42 ++++++++++-----------------------------
 internal/tools/headers.go | 14 ++++++-------
 server/apiv1/release.go   |  4 ++--
 3 files changed, 20 insertions(+), 40 deletions(-)

diff --git a/internal/smtpd/main.go b/internal/smtpd/main.go
index 2eb2ef6..37f7a68 100644
--- a/internal/smtpd/main.go
+++ b/internal/smtpd/main.go
@@ -34,7 +34,7 @@ func mailHandler(origin net.Addr, from string, to []string, data []byte) (string
 
 // SaveToDatabase will attempt to save a message to the database
 func SaveToDatabase(origin net.Addr, from string, to []string, data []byte) (string, error) {
-	if !config.SMTPStrictRFCHeaders {
+	if !config.SMTPStrictRFCHeaders && bytes.Contains(data, []byte("\r\r\n")) {
 		// replace all <CR><CR><LF> (\r\r\n) with <CR><LF> (\r\n)
 		// @see https://github.com/axllent/mailpit/issues/87 & https://github.com/axllent/mailpit/issues/153
 		data = bytes.ReplaceAll(data, []byte("\r\r\n"), []byte("\r\n"))
@@ -50,21 +50,9 @@ func SaveToDatabase(origin net.Addr, from string, to []string, data []byte) (str
 	// check / set the Return-Path based on SMTP from
 	returnPath := strings.Trim(msg.Header.Get("Return-Path"), "<>")
 	if returnPath != from {
-		if returnPath != "" {
-			// replace Return-Path
-			re := regexp.MustCompile(`(?i)(^|\n)(Return\-Path: .*\n)`)
-			replaced := false
-			data = re.ReplaceAllFunc(data, func(r []byte) []byte {
-				if replaced {
-					return r
-				}
-				replaced = true // only replace first occurrence
-
-				return re.ReplaceAll(r, []byte("${1}Return-Path: <"+from+">\r\n"))
-			})
-		} else {
-			// add Return-Path
-			data = append([]byte("Return-Path: <"+from+">\r\n"), data...)
+		data, err = tools.SetMessageHeader(data, "Return-Path", "<"+from+">")
+		if err != nil {
+			return "", err
 		}
 	}
 
@@ -108,23 +96,15 @@ func SaveToDatabase(origin net.Addr, from string, to []string, data []byte) (str
 
 	// add missing email addresses to Bcc (eg: Laravel doesn't include these in the headers)
 	if len(missingAddresses) > 0 {
+		bccVal := strings.Join(missingAddresses, ", ")
 		if hasBccHeader {
-			// email already has Bcc header, add to existing addresses
-			re := regexp.MustCompile(`(?i)(^|\n)(Bcc: )`)
-			replaced := false
-			data = re.ReplaceAllFunc(data, func(r []byte) []byte {
-				if replaced {
-					return r
-				}
-				replaced = true // only replace first occurrence
+			b := msg.Header.Get("Bcc")
+			bccVal = ", " + b
+		}
 
-				return re.ReplaceAll(r, []byte("${1}Bcc: "+strings.Join(missingAddresses, ", ")+", "))
-			})
-
-		} else {
-			// prepend new Bcc header
-			bcc := []byte(fmt.Sprintf("Bcc: %s\r\n", strings.Join(missingAddresses, ", ")))
-			data = append(bcc, data...)
+		data, err = tools.SetMessageHeader(data, "Bcc", bccVal)
+		if err != nil {
+			return "", err
 		}
 
 		logger.Log().Debugf("[smtpd] added missing addresses to Bcc header: %s", strings.Join(missingAddresses, ", "))
diff --git a/internal/tools/headers.go b/internal/tools/headers.go
index 0a8ad46..589298d 100644
--- a/internal/tools/headers.go
+++ b/internal/tools/headers.go
@@ -58,8 +58,10 @@ func RemoveMessageHeaders(msg []byte, headers []string) ([]byte, error) {
 	return msg, nil
 }
 
-// UpdateMessageHeader scans a message for a header and updates its value if found.
-func UpdateMessageHeader(msg []byte, header, value string) ([]byte, error) {
+// SetMessageHeader scans a message for a header and updates its value if found.
+// It does not consider multiple instances of the same header.
+// If not found it will add the header to the beginning of the message.
+func SetMessageHeader(msg []byte, header, value string) ([]byte, error) {
 	reader := bytes.NewReader(msg)
 	m, err := mail.ReadMessage(reader)
 	if err != nil {
@@ -90,13 +92,11 @@ func UpdateMessageHeader(msg []byte, header, value string) ([]byte, error) {
 			}
 		}
 
-		if len(hdr) > 0 {
-			logger.Log().Debugf("[relay] replaced %s header", hdr)
-			msg = bytes.Replace(msg, hdr, []byte(header+": "+value+"\r\n"), 1)
-		}
+		return bytes.Replace(msg, hdr, []byte(header+": "+value+"\r\n"), 1), nil
 	}
 
-	return msg, nil
+	// no header, so add one to beginning
+	return append([]byte(header+": "+value+"\r\n"), msg...), nil
 }
 
 // OverrideFromHeader scans a message for the From header and replaces it with a different email address.
diff --git a/server/apiv1/release.go b/server/apiv1/release.go
index 3c1c9ce..b394e79 100644
--- a/server/apiv1/release.go
+++ b/server/apiv1/release.go
@@ -170,7 +170,7 @@ func ReleaseMessage(w http.ResponseWriter, r *http.Request) {
 	}
 
 	// update message date
-	msg, err = tools.UpdateMessageHeader(msg, "Date", time.Now().Format(time.RFC1123Z))
+	msg, err = tools.SetMessageHeader(msg, "Date", time.Now().Format(time.RFC1123Z))
 	if err != nil {
 		httpError(w, err.Error())
 		return
@@ -179,7 +179,7 @@ func ReleaseMessage(w http.ResponseWriter, r *http.Request) {
 	// generate unique ID
 	uid := shortuuid.New() + "@mailpit"
 	// update Message-Id with unique ID
-	msg, err = tools.UpdateMessageHeader(msg, "Message-Id", "<"+uid+">")
+	msg, err = tools.SetMessageHeader(msg, "Message-Id", "<"+uid+">")
 	if err != nil {
 		httpError(w, err.Error())
 		return