From c096f018d64518a3d891bc87f29b49d0d7c8dd1a Mon Sep 17 00:00:00 2001
From: Timur Makarchuk <Timur.makarchuk@gmail.com>
Date: Sat, 10 Apr 2021 23:45:09 +0300
Subject: [PATCH] Add support for AUTH, closes #62

* Add PLAIN and LOGIN auth support
---
 pkg/server/smtp/handler.go      | 73 +++++++++++++++++++++++++++++++++
 pkg/server/smtp/handler_test.go | 40 ++++++++++++++++++
 2 files changed, 113 insertions(+)

diff --git a/pkg/server/smtp/handler.go b/pkg/server/smtp/handler.go
index 41db786..5171be0 100644
--- a/pkg/server/smtp/handler.go
+++ b/pkg/server/smtp/handler.go
@@ -29,12 +29,25 @@ const (
 	GREET State = iota
 	// READY State: Got HELO, waiting for MAIL
 	READY
+	// LOGIN State: Got AUTH LOGIN command, expecting Username
+	LOGIN
+	// PASSWORD State: Got Username, expecting password
+	PASSWORD
 	// MAIL State: Got MAIL, accepting RCPTs
 	MAIL
 	// DATA State: Got DATA, waiting for "."
 	DATA
 	// QUIT State: Client requested end of session
 	QUIT
+
+	// Messages sent to user during LOGIN auth procedure
+	// Can vary, but values are taken directly from spec
+	// https://tools.ietf.org/html/draft-murchison-sasl-login-00
+
+	//usernameChallenge sent when inviting user to provide username. Is base64 encoded string `User Name`
+	usernameChallenge = "VXNlciBOYW1lAA=="
+	//passwordChallenge sent when inviting user to provide password. Is base64 encoded string `Password`
+	passwordChallenge = "UGFzc3dvcmQA"
 )
 
 // fromRegex captures the from address and optional BODY=8BITMIME clause.  Matches FROM, while
@@ -76,6 +89,7 @@ var commands = map[string]bool{
 	"QUIT":     true,
 	"TURN":     true,
 	"STARTTLS": true,
+	"AUTH":     true,
 }
 
 // Session holds the state of an SMTP session
@@ -153,6 +167,16 @@ func (s *Server) startSession(id int, conn net.Conn) {
 		}
 		line, err := ssn.readLine()
 		if err == nil {
+			//Handle LOGIN/PASSWORD states here, because they don't expect a command
+			switch ssn.state {
+			case LOGIN:
+				ssn.loginHandler(line)
+				continue
+			case PASSWORD:
+				ssn.passwordHandler(line)
+				continue
+			}
+
 			if cmd, arg, ok := ssn.parseCmd(line); ok {
 				// Check against valid SMTP commands
 				if cmd == "" {
@@ -260,6 +284,7 @@ func (s *Session) greetHandler(cmd string, arg string) {
 		// features before SIZE per RFC
 		s.send("250-" + readyBanner)
 		s.send("250-8BITMIME")
+		s.send("250-AUTH PLAIN LOGIN")
 		if s.Server.config.TLSEnabled && s.Server.tlsConfig != nil && s.tlsState == nil {
 			s.send("250-STARTTLS")
 		}
@@ -281,7 +306,28 @@ func parseHelloArgument(arg string) (string, error) {
 	return domain, nil
 }
 
+func (s *Session) loginHandler(line string) {
+	if len(line) == 0 {
+		s.send("500 invalid Username")
+		s.enterState(READY)
+		return
+	}
+	s.send(fmt.Sprintf("334 %v", passwordChallenge))
+	s.enterState(PASSWORD)
+}
+
+func (s *Session) passwordHandler(line string) {
+	if len(line) == 0 {
+		s.send("500 invalid Password")
+		s.enterState(READY)
+		return
+	}
+	s.send("235 Authentication successful")
+	s.enterState(READY)
+}
+
 // READY state -> waiting for MAIL
+// AUTH can change
 func (s *Session) readyHandler(cmd string, arg string) {
 	if cmd == "STARTTLS" {
 		if !s.Server.config.TLSEnabled {
@@ -305,6 +351,33 @@ func (s *Session) readyHandler(cmd string, arg string) {
 		s.tlsState = new(tls.ConnectionState)
 		*s.tlsState = tlsConn.ConnectionState()
 		s.enterState(GREET)
+	} else if cmd == "AUTH" {
+		args := strings.SplitN(arg, " ", 3)
+		authMethod := args[0]
+		switch authMethod {
+		case "PLAIN":
+			{
+				if len(args) != 2 {
+					s.send("500 Bad auth arguments")
+					s.logger.Warn().Msgf("Bad auth attempt: %q", arg)
+					return
+				}
+				s.logger.Info().Msgf("Accepting credentials: %q", args[1])
+				s.send("235 2.7.0 Authentication successful")
+				return
+			}
+		case "LOGIN":
+			{
+				s.send(fmt.Sprintf("334 %v", usernameChallenge))
+				s.enterState(LOGIN)
+				return
+			}
+		default:
+			{
+				s.send(fmt.Sprintf("500 Unsupported AUTH method: %v", authMethod))
+				return
+			}
+		}
 	} else if cmd == "MAIL" {
 		// Capture group 1: from address.  2: optional params.
 		m := fromRegex.FindStringSubmatch(arg)
diff --git a/pkg/server/smtp/handler_test.go b/pkg/server/smtp/handler_test.go
index bcb61d8..08393d1 100644
--- a/pkg/server/smtp/handler_test.go
+++ b/pkg/server/smtp/handler_test.go
@@ -108,6 +108,46 @@ func TestEmptyEnvelope(t *testing.T) {
 	}
 }
 
+// Test AUTH
+func TestAuth(t *testing.T) {
+	ds := test.NewStore()
+	server, logbuf, teardown := setupSMTPServer(ds)
+	defer teardown()
+
+	//PLAIN AUTH
+	script := []scriptStep{
+		{"EHLO localhost", 250},
+		{"AUTH PLAIN aW5idWNrZXQ6cGFzc3dvcmQK", 235},
+		{"RSET", 250},
+		{"AUTH GSSAPI aW5idWNrZXQ6cGFzc3dvcmQK", 500},
+		{"RSET", 250},
+		{"AUTH PLAIN", 500},
+		{"RSET", 250},
+		{"AUTH PLAIN aW5idWNrZXQ6cG Fzc3dvcmQK", 500},
+	}
+	if err := playSession(t, server, script); err != nil {
+		t.Error(err)
+	}
+
+	//LOGIN AUTH
+	script = []scriptStep{
+		{"EHLO localhost", 250},
+		{"AUTH LOGIN", 334},
+		{"USERNAME", 334},
+		{"PASSWORD", 235},
+	}
+	if err := playSession(t, server, script); err != nil {
+		t.Error(err)
+	}
+
+	if t.Failed() {
+		// Wait for handler to finish logging
+		time.Sleep(2 * time.Second)
+		// Dump buffered log data if there was a failure
+		_, _ = io.Copy(os.Stderr, logbuf)
+	}
+}
+
 // Test commands in READY state
 func TestReadyState(t *testing.T) {
 	ds := test.NewStore()