SMTP state: GetHeader()

smtp/state.go

@@ -1,8 +1,11 @@
 package smtp
 
 import (
+	"bufio"
+	"bytes"
 	"fmt"
 	"net"
+	"strings"
 )
 
 // State contains all the state for a single client
@@ -84,3 +87,27 @@ func (s *State) AddHeader(headerKey string, headerValue string) {
 	header := fmt.Sprintf("%s: %s\r\n", headerKey, headerValue)
 	s.Data = append([]byte(header), s.Data...)
 }
+
+// GetHeader gets a header from the state.
+func (s *State) GetHeader(headerKey string) (string, error) {
+	reader := bufio.NewReader(bytes.NewReader(s.Data))
+	for {
+		line, err := reader.ReadString('\n')
+		if err != nil {
+			break
+		}
+
+		// Headers end with an empty line
+		if len(strings.TrimSpace(line)) == 0 {
+			break
+		}
+
+		if strings.HasPrefix(strings.ToLower(line), strings.ToLower(headerKey)+":") {
+			// Found the header, extract the value
+			headerValue := strings.TrimSpace(line[len(headerKey)+1:])
+			return headerValue, nil
+		}
+	}
+
+	return "", fmt.Errorf("header not found")
+}

smtp/state_test.go

@@ -35,4 +35,38 @@ This is the body of the email.`
 		So(parsedMessage.Header.Get("From"), ShouldEqual, "[email protected]")
 
 	})
+
+	Convey("GetHeader()", t, func() {
+		// Create a state object with a message
+		message := `From: [email protected]
+To: [email protected]
+X-Spam-Score: -5.1
+Subject: Test Subject
+
+This is the body of the email.`
+
+		state := &State{
+			Data: []byte(message),
+		}
+
+		headerValue, err := state.GetHeader("To")
+		So(err, ShouldBeNil)
+		So(headerValue, ShouldEqual, "[email protected]")
+
+		headerValue, err = state.GetHeader("to")
+		So(err, ShouldBeNil)
+		So(headerValue, ShouldEqual, "[email protected]")
+
+		headerValue, err = state.GetHeader("TO")
+		So(err, ShouldBeNil)
+		So(headerValue, ShouldEqual, "[email protected]")
+
+		headerValue, err = state.GetHeader("From")
+		So(err, ShouldBeNil)
+		So(headerValue, ShouldEqual, "[email protected]")
+
+		headerValue, err = state.GetHeader("X-Spam-Score")
+		So(err, ShouldBeNil)
+		So(headerValue, ShouldEqual, "-5.1")
+	})
 }