From 256849f2c67997f20d730ecbaedd3af549d6dc81 Mon Sep 17 00:00:00 2001
From: Ian Kent <email@iankent.co.uk>
Date: Sat, 22 Nov 2014 19:20:47 +0000
Subject: [PATCH] Add sender/recipient validate hooks

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

diff --git a/mailhog/smtp/protocol/protocol.go b/mailhog/smtp/protocol/protocol.go
index bcdf2c5..46005ee 100644
--- a/mailhog/smtp/protocol/protocol.go
+++ b/mailhog/smtp/protocol/protocol.go
@@ -23,8 +23,10 @@ type Protocol struct {
 	message  *data.SMTPMessage
 	hostname string
 
-	LogHandler             func(message string, args ...interface{})
-	MessageReceivedHandler func(*data.Message) (string, error)
+	LogHandler               func(message string, args ...interface{})
+	MessageReceivedHandler   func(*data.Message) (string, error)
+	ValidateSenderHandler    func(from string) bool
+	ValidateRecipientHandler func(to string) bool
 }
 
 // NewProtocol returns a new SMTP state machine in INVALID state
@@ -85,7 +87,6 @@ func (proto *Protocol) Parse(line string) (string, *Reply) {
 // ProcessData handles content received (with newlines stripped) while
 // in the SMTP DATA state
 func (proto *Protocol) ProcessData(line string) (reply *Reply) {
-
 	proto.message.Data += line + "\n"
 
 	if strings.HasSuffix(proto.message.Data, "\r\n.\r\n") {
@@ -191,6 +192,12 @@ func (proto *Protocol) Command(command *Command) (reply *Reply) {
 			if err != nil {
 				return ReplyError(err)
 			}
+			if proto.ValidateSenderHandler != nil {
+				if !proto.ValidateSenderHandler(from) {
+					// TODO correct sender error response
+					return ReplyError(errors.New("Invalid sender " + from))
+				}
+			}
 			proto.message.From = from
 			proto.state = RCPT
 			return ReplySenderOk(from)
@@ -210,6 +217,12 @@ func (proto *Protocol) Command(command *Command) (reply *Reply) {
 			if err != nil {
 				return ReplyError(err)
 			}
+			if proto.ValidateRecipientHandler != nil {
+				if !proto.ValidateRecipientHandler(to) {
+					// TODO correct send error response
+					return ReplyError(errors.New("Invalid recipient " + to))
+				}
+			}
 			proto.message.To = append(proto.message.To, to)
 			proto.state = RCPT
 			return ReplyRecipientOk(to)
diff --git a/mailhog/smtp/server/session.go b/mailhog/smtp/server/session.go
index 59f84b6..202ae24 100644
--- a/mailhog/smtp/server/session.go
+++ b/mailhog/smtp/server/session.go
@@ -29,6 +29,8 @@ func Accept(conn *net.TCPConn, conf *config.Config) {
 	session := &Session{conn, proto, conf, false, ""}
 	proto.LogHandler = session.logf
 	proto.MessageReceivedHandler = session.acceptMessageHandler
+	proto.ValidateSenderHandler = session.validateSender
+	proto.ValidateRecipientHandler = session.validateRecipient
 
 	session.logf("Starting session")
 	session.Write(proto.Start(conf.Hostname))
@@ -37,6 +39,14 @@ func Accept(conn *net.TCPConn, conf *config.Config) {
 	session.logf("Session ended")
 }
 
+func (c *Session) validateRecipient(to string) bool {
+	return true
+}
+
+func (c *Session) validateSender(from string) bool {
+	return true
+}
+
 func (c *Session) acceptMessageHandler(msg *data.Message) (id string, err error) {
 	switch c.conf.Storage.(type) {
 	case *storage.MongoDB: