From 559c2a5146a7360a6b8cd23a2eb6c343053ee708 Mon Sep 17 00:00:00 2001
From: Ian Kent <email@iankent.co.uk>
Date: Sun, 23 Nov 2014 00:42:15 +0000
Subject: [PATCH] Make ident configurable

---
 mailhog/smtp/protocol/protocol.go | 23 ++++++++++++++---------
 mailhog/smtp/server/session.go    |  3 ++-
 2 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/mailhog/smtp/protocol/protocol.go b/mailhog/smtp/protocol/protocol.go
index cf14fbe..fccd15e 100644
--- a/mailhog/smtp/protocol/protocol.go
+++ b/mailhog/smtp/protocol/protocol.go
@@ -19,9 +19,11 @@ type Command struct {
 
 // Protocol is a state machine representing an SMTP session
 type Protocol struct {
-	state    State
-	message  *data.SMTPMessage
-	hostname string
+	state   State
+	message *data.SMTPMessage
+
+	Hostname string
+	Ident    string
 
 	// LogHandler is called for each log message. If nil, log messages will
 	// be output using log.Printf instead.
@@ -46,8 +48,10 @@ type Protocol struct {
 // handler is called when a message is received and should return a message ID
 func NewProtocol() *Protocol {
 	return &Protocol{
-		state:   INVALID,
-		message: &data.SMTPMessage{},
+		state:    INVALID,
+		message:  &data.SMTPMessage{},
+		Hostname: "",
+		Ident:    "ESMTP Go-MailHog",
 	}
 }
 
@@ -64,10 +68,10 @@ func (proto *Protocol) logf(message string, args ...interface{}) {
 
 // Start begins an SMTP conversation with a 220 reply, placing the state
 // machine in ESTABLISH state.
-func (proto *Protocol) Start(hostname string) *Reply {
+func (proto *Protocol) Start() *Reply {
+	proto.logf("Started session, switching to ESTABLISH state")
 	proto.state = ESTABLISH
-	proto.hostname = hostname
-	return ReplyIdent(hostname + " ESMTP Go-MailHog")
+	return ReplyIdent(proto.Hostname + " " + proto.Ident)
 }
 
 // Parse parses a line string and returns any remaining line string
@@ -90,6 +94,7 @@ func (proto *Protocol) Parse(line string) (string, *Reply) {
 		line = ""
 	}
 
+	// TODO collapse AUTH states into separate processing
 	if proto.state == DATA {
 		reply = proto.ProcessData(parts[0])
 	} else {
@@ -111,7 +116,7 @@ func (proto *Protocol) ProcessData(line string) (reply *Reply) {
 		proto.message.Data = strings.TrimSuffix(proto.message.Data, "\r\n.\r\n")
 		proto.state = MAIL
 
-		msg := proto.message.Parse(proto.hostname)
+		msg := proto.message.Parse(proto.Hostname)
 
 		if proto.MessageReceivedHandler == nil {
 			return ReplyStorageFailed("No storage backend")
diff --git a/mailhog/smtp/server/session.go b/mailhog/smtp/server/session.go
index daecfeb..da53589 100644
--- a/mailhog/smtp/server/session.go
+++ b/mailhog/smtp/server/session.go
@@ -24,6 +24,7 @@ type Session struct {
 // Accept starts a new SMTP session using net.TCPConn
 func Accept(conn *net.TCPConn, conf *config.Config) {
 	proto := protocol.NewProtocol()
+	proto.Hostname = conf.Hostname
 	session := &Session{conn, proto, conf, false, ""}
 	proto.LogHandler = session.logf
 	proto.MessageReceivedHandler = session.acceptMessage
@@ -32,7 +33,7 @@ func Accept(conn *net.TCPConn, conf *config.Config) {
 	proto.ValidateAuthenticationHandler = session.validateAuthentication
 
 	session.logf("Starting session")
-	session.Write(proto.Start(conf.Hostname))
+	session.Write(proto.Start())
 	for session.Read() == true {
 	}
 	session.logf("Session ended")