Added AddHeader() to SMTP state

smtp/protocol.go

@@ -402,80 +402,6 @@ func (id *Id) String() string {
 	return strconv.FormatInt(id.Timestamp, 16) + strconv.FormatInt(int64(id.Counter), 16)
 }
 
-// State contains all the state for a single client
-type State struct {
-	From          *MailAddress
-	To            []*MailAddress
-	Data          []byte
-	EightBitMIME  bool
-	Secure        bool
-	SessionId     Id
-	Ip            net.IP
-	Hostname      string
-	Authenticated bool
-	User          User
-}
-
-// User denotes an authenticated SMTP user.
-type User interface {
-	// Username returns the username / email address of the user.
-	Username() string
-}
-
-// reset the state
-func (s *State) Reset() {
-	s.From = nil
-	s.To = []*MailAddress{}
-	s.Data = []byte{}
-	s.EightBitMIME = false
-}
-
-// Checks the state if the client can send a MAIL command.
-func (s *State) CanReceiveMail() (bool, string) {
-	if s.From != nil {
-		return false, "Sender already specified"
-	}
-
-	return true, ""
-}
-
-// Checks the state if the client can send a RCPT command.
-func (s *State) CanReceiveRcpt() (bool, string) {
-	if s.From == nil {
-		return false, "Need mail before RCPT"
-	}
-
-	return true, ""
-}
-
-// Checks the state if the client can send a DATA command.
-func (s *State) CanReceiveData() (bool, string) {
-	if s.From == nil {
-		return false, "Need mail before DATA"
-	}
-
-	if len(s.To) == 0 {
-		return false, "Need RCPT before DATA"
-	}
-
-	return true, ""
-}
-
-// Check whether the auth user is allowed to send from the MAIL FROM email address and to the RCPT TO address.
-func (s *State) AuthMatchesRcptAndMail() (bool, string) {
-
-	// TODO: what if one of those variables is nil?
-
-	// TODO: handle if user can send from multiple email addresses
-	if s.From.Address != s.User.Username() {
-		return false, fmt.Sprintf("5.7.1 Sender address rejected: not owned by user %s", s.User.Username())
-	}
-
-	// TODO: check for recipient?
-
-	return true, ""
-}
-
 // Protocol Used as communication layer so we can easily switch between a real socket
 // and a test implementation.
 type Protocol interface {

smtp/state.go

@@ -0,0 +1,86 @@
+package smtp
+
+import (
+	"fmt"
+	"net"
+)
+
+// State contains all the state for a single client
+type State struct {
+	From          *MailAddress
+	To            []*MailAddress
+	Data          []byte
+	EightBitMIME  bool
+	Secure        bool
+	SessionId     Id
+	Ip            net.IP
+	Hostname      string
+	Authenticated bool
+	User          User
+}
+
+// User denotes an authenticated SMTP user.
+type User interface {
+	// Username returns the username / email address of the user.
+	Username() string
+}
+
+// reset the state
+func (s *State) Reset() {
+	s.From = nil
+	s.To = []*MailAddress{}
+	s.Data = []byte{}
+	s.EightBitMIME = false
+}
+
+// Checks the state if the client can send a MAIL command.
+func (s *State) CanReceiveMail() (bool, string) {
+	if s.From != nil {
+		return false, "Sender already specified"
+	}
+
+	return true, ""
+}
+
+// Checks the state if the client can send a RCPT command.
+func (s *State) CanReceiveRcpt() (bool, string) {
+	if s.From == nil {
+		return false, "Need mail before RCPT"
+	}
+
+	return true, ""
+}
+
+// Checks the state if the client can send a DATA command.
+func (s *State) CanReceiveData() (bool, string) {
+	if s.From == nil {
+		return false, "Need mail before DATA"
+	}
+
+	if len(s.To) == 0 {
+		return false, "Need RCPT before DATA"
+	}
+
+	return true, ""
+}
+
+// Check whether the auth user is allowed to send from the MAIL FROM email address and to the RCPT TO address.
+func (s *State) AuthMatchesRcptAndMail() (bool, string) {
+
+	// TODO: what if one of those variables is nil?
+
+	// TODO: handle if user can send from multiple email addresses
+	if s.From.Address != s.User.Username() {
+		return false, fmt.Sprintf("5.7.1 Sender address rejected: not owned by user %s", s.User.Username())
+	}
+
+	// TODO: check for recipient?
+
+	return true, ""
+}
+
+// AddHeader prepends the given header to the state.
+func (s *State) AddHeader(headerKey string, headerValue string) {
+	header := fmt.Sprintf("%s: %s\r\n", headerKey, headerValue)
+	s.Data = append([]byte(header), s.Data...)
+}

smtp/state_test.go

@@ -0,0 +1,38 @@
+package smtp
+
+import (
+	"net/mail"
+	"strings"
+	"testing"
+
+	. "github.com/smartystreets/goconvey/convey"
+)
+
+func TestState(t *testing.T) {
+
+	Convey("AddHeader()", t, func() {
+
+		// Create a state object with a message
+		message := `From: [email protected]
+To: [email protected]
+Subject: Test Subject
+
+This is the body of the email.`
+
+		state := &State{
+			Data: []byte(message),
+		}
+
+		// Add the header
+		state.AddHeader("MessageId", "some-value@localhost")
+
+		// If we now parse the state data again, it should be valid and the header should be present
+		parsedMessage, err := mail.ReadMessage(strings.NewReader(string(state.Data)))
+		So(err, ShouldBeNil)
+		So(parsedMessage.Header.Get("MessageId"), ShouldEqual, "some-value@localhost")
+
+		// and make sure the rest is also still there...
+		So(parsedMessage.Header.Get("From"), ShouldEqual, "[email protected]")
+
+	})
+}