server: replace EnableAuth with AuthSession

backend.go

@@ -2,6 +2,8 @@ package smtp
 
 import (
 	"io"
+
+	"github.com/emersion/go-sasl"
 )
 
 var (
@@ -20,6 +22,11 @@ var (
 		EnhancedCode: EnhancedCode{5, 7, 0},
 		Message:      "Authentication not supported",
 	}
+	ErrAuthUnknownMechanism = &SMTPError{
+		Code:         504,
+		EnhancedCode: EnhancedCode{5, 7, 4},
+		Message:      "Unsupported authentication mechanism",
+	}
 )
 
 // A SMTP server backend.
@@ -87,3 +94,12 @@ type LMTPSession interface {
 type StatusCollector interface {
 	SetStatus(rcptTo string, err error)
 }
+
+// AuthSession is an add-on interface for Session. It provides support for the
+// AUTH extension.
+type AuthSession interface {
+	Session
+
+	AuthMechanisms() []string
+	Auth(mech string) (sasl.Server, error)
+}

conn.go

@@ -15,6 +15,8 @@ import (
 	"strings"
 	"sync"
 	"time"
+
+	"github.com/emersion/go-sasl"
 )
 
 // Number of errors we'll tolerate per connection before closing. Defaults to 3.
@@ -261,7 +263,7 @@ func (c *Conn) handleGreet(enhanced bool, arg string) {
 	}
 	if c.authAllowed() {
 		authCap := "AUTH"
-		for name := range c.server.auths {
+		for _, name := range c.authMechanisms() {
 			authCap += " " + name
 		}
 
@@ -792,14 +794,16 @@ func (c *Conn) handleAuth(arg string) {
 		}
 	}
 
-	newSasl, ok := c.server.auths[mechanism]
-	if !ok {
-		c.writeResponse(504, EnhancedCode{5, 7, 4}, "Unsupported authentication mechanism")
+	sasl, err := c.auth(mechanism)
+	if err != nil {
+		if smtpErr, ok := err.(*SMTPError); ok {
+			c.writeResponse(smtpErr.Code, smtpErr.EnhancedCode, smtpErr.Message)
+		} else {
+			c.writeResponse(454, EnhancedCode{4, 7, 0}, err.Error())
+		}
 		return
 	}
 
-	sasl := newSasl(c)
-
 	response := ir
 	for {
 		challenge, done, err := sasl.Next(response)
@@ -844,6 +848,36 @@ func (c *Conn) handleAuth(arg string) {
 	c.didAuth = true
 }
 
+func (c *Conn) authMechanisms() []string {
+	if authSession, ok := c.Session().(AuthSession); ok {
+		return authSession.AuthMechanisms()
+	}
+	return []string{sasl.Plain}
+}
+
+func (c *Conn) auth(mech string) (sasl.Server, error) {
+	if authSession, ok := c.Session().(AuthSession); ok {
+		return authSession.Auth(mech)
+	}
+
+	if mech != sasl.Plain {
+		return nil, ErrAuthUnknownMechanism
+	}
+
+	return sasl.NewPlainServer(func(identity, username, password string) error {
+		if identity != "" && identity != username {
+			return errors.New("identities not supported")
+		}
+
+		sess := c.Session()
+		if sess == nil {
+			panic("No session when AUTH is called")
+		}
+
+		return sess.AuthPlain(username, password)
+	}), nil
+}
+
 func (c *Conn) handleStartTLS() {
 	if _, isTLS := c.TLSConnectionState(); isTLS {
 		c.writeResponse(502, EnhancedCode{5, 5, 1}, "Already running in TLS")

server.go

@@ -10,17 +10,12 @@ import (
 	"os"
 	"sync"
 	"time"
-
-	"github.com/emersion/go-sasl"
 )
 
 var (
 	ErrServerClosed = errors.New("smtp: server already closed")
 )
 
-// A function that creates SASL servers.
-type SaslServerFactory func(conn *Conn) sasl.Server
-
 // Logger interface is used by Server to report unexpected internal errors.
 type Logger interface {
 	Printf(format string, v ...interface{})
@@ -71,10 +66,8 @@ type Server struct {
 	// The server backend.
 	Backend Backend
 
-	wg sync.WaitGroup
-
-	auths map[string]SaslServerFactory
-	done  chan struct{}
+	wg   sync.WaitGroup
+	done chan struct{}
 
 	locker    sync.Mutex
 	listeners []net.Listener
@@ -90,23 +83,7 @@ func NewServer(be Backend) *Server {
 		Backend:  be,
 		done:     make(chan struct{}, 1),
 		ErrorLog: log.New(os.Stderr, "smtp/server ", log.LstdFlags),
-		auths: map[string]SaslServerFactory{
-			sasl.Plain: func(conn *Conn) sasl.Server {
-				return sasl.NewPlainServer(func(identity, username, password string) error {
-					if identity != "" && identity != username {
-						return errors.New("identities not supported")
-					}
-
-					sess := conn.Session()
-					if sess == nil {
-						panic("No session when AUTH is called")
-					}
-
-					return sess.AuthPlain(username, password)
-				})
-			},
-		},
-		conns: make(map[*Conn]struct{}),
+		conns:    make(map[*Conn]struct{}),
 	}
 }
 
@@ -327,11 +304,3 @@ func (s *Server) Shutdown(ctx context.Context) error {
 		return err
 	}
 }
-
-// EnableAuth enables an authentication mechanism on this server.
-//
-// This function should not be called directly, it must only be used by
-// libraries implementing extensions of the SMTP protocol.
-func (s *Server) EnableAuth(name string, f SaslServerFactory) {
-	s.auths[name] = f
-}