From 3a1c757d044d78b60b3c667f2fb701124e3f2eb5 Mon Sep 17 00:00:00 2001
From: James Hillyerd <james@hillyerd.com>
Date: Sat, 17 Aug 2019 15:46:28 -0700
Subject: [PATCH] smtp: Handle late EHLO, fixes #141

---
 pkg/server/smtp/handler.go      | 11 +++++++++++
 pkg/server/smtp/handler_test.go | 13 +++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/pkg/server/smtp/handler.go b/pkg/server/smtp/handler.go
index b90333e..bcedfe6 100644
--- a/pkg/server/smtp/handler.go
+++ b/pkg/server/smtp/handler.go
@@ -346,6 +346,11 @@ func (s *Session) readyHandler(cmd string, arg string) {
 		s.logger.Info().Msgf("Mail from: %v", from)
 		s.send(fmt.Sprintf("250 Roger, accepting mail from <%v>", from))
 		s.enterState(MAIL)
+	} else if cmd == "EHLO" {
+		// Reset session
+		s.logger.Debug().Msgf("Resetting session state on EHLO request")
+		s.reset()
+		s.send("250 Session reset")
 	} else {
 		s.ooSeq(cmd)
 	}
@@ -394,6 +399,12 @@ func (s *Session) mailHandler(cmd string, arg string) {
 		}
 		s.enterState(DATA)
 		return
+	case "EHLO":
+		// Reset session
+		s.logger.Debug().Msgf("Resetting session state on EHLO request")
+		s.reset()
+		s.send("250 Session reset")
+		return
 	}
 	s.ooSeq(cmd)
 }
diff --git a/pkg/server/smtp/handler_test.go b/pkg/server/smtp/handler_test.go
index bc63a58..5aa3e35 100644
--- a/pkg/server/smtp/handler_test.go
+++ b/pkg/server/smtp/handler_test.go
@@ -206,6 +206,19 @@ func TestMailState(t *testing.T) {
 		t.Error(err)
 	}
 
+	// Test late EHLO, similar to RSET
+	script = []scriptStep{
+		{"EHLO localhost", 250},
+		{"EHLO localhost", 250},
+		{"MAIL FROM:<john@gmail.com>", 250},
+		{"RCPT TO:<u1@gmail.com>", 250},
+		{"EHLO localhost", 250},
+		{"MAIL FROM:<john@gmail.com>", 250},
+	}
+	if err := playSession(t, server, script); err != nil {
+		t.Error(err)
+	}
+
 	// Test RSET
 	script = []scriptStep{
 		{"HELO localhost", 250},