mail/server/server.go

package server

import (
	"context"
	"fmt"
	"io"
	"strings"
	"time"

	"github.com/emersion/go-message/mail"
	"github.com/emersion/go-milter"
	"github.com/emersion/go-smtp"
	"go.uber.org/zap"
	"golang.org/x/time/rate"
)

const ConnectionTimeout = 3 * time.Minute

// ErrRateLimit indicates that the message is rejected due to a rate limit.
var ErrRateLimit = &smtp.SMTPError{
	Code:         450,
	EnhancedCode: [3]int{4, 2, 0},
	Message:      "too many requests",
}

// ErrBadGateway indicates a downstream system failure.
var ErrBadGateway = &smtp.SMTPError{
	Code:         451,
	EnhancedCode: [3]int{4, 4, 3},
	Message:      "bad gateway",
}

// Backend provides a backend implementation of an SMTP server.
type Backend struct {
	AllowedDomains []string

	Log *zap.Logger

	RootContext context.Context

	RSpam *milter.Client

	OverallLimiter *rate.Limiter
}

// NewSession implements smtp.Backend
func (b *Backend) NewSession(conn *smtp.Conn) (smtp.Session, error) {
	ctx, cancel := context.WithTimeout(b.RootContext, ConnectionTimeout)
	defer cancel()

	if err := b.OverallLimiter.Wait(ctx); err != nil {
		return nil, ErrRateLimit
	}

	rs, err := b.RSpam.Session()
	if err != nil {
		return nil, ErrBadGateway
	}

	return &Session{
		b:     b,
		conn:  conn,
		rspam: rs,
	}, nil
}

// Session provides an implementation of an SMTP session.
type Session struct {
	b     *Backend
	conn  *smtp.Conn
	rspam *milter.ClientSession
}

// Discard currently processed message.
func (session *Session) Reset() {
	panic("not implemented") // TODO: Implement
}

// Free all resources associated with session.
func (session *Session) Logout() error {
	panic("not implemented") // TODO: Implement
}

// Authenticate the user using SASL PLAIN.
func (session *Session) AuthPlain(username string, password string) error {
	return smtp.ErrAuthUnsupported
}

// Set return path for currently processed message.
func (session *Session) Mail(from string, opts *smtp.MailOptions) error {
	action, err := session.rspam.Mail(from, nil)
	if err != nil {
		return fmt.Errorf("filter failure: %w", err)
	}

	return parseAction(action)
}

// Add recipient for currently processed message.
func (session *Session) Rcpt(to string) error {
	if err := session.validateDomain(to); err != nil {
		return &smtp.SMTPError{
			Code:         550,
			EnhancedCode: [3]int{5, 1, 2},
			Message:      err.Error(),
		}
	}

	action, err := session.rspam.Rcpt(to, nil)
	if err != nil {
		return fmt.Errorf("filter failure: %w", err)
	}

	return parseAction(action)
}

// Set currently processed message contents and send it.
//
// r must be consumed before Data returns.
func (session *Session) Data(r io.Reader) error {
	// TODO: handle the modifyAction below, whatever it is.
	_, action, err := session.rspam.BodyReadFrom(r)
	if err != nil {
		return fmt.Errorf("filter failure: %w", err)
	}

	return parseAction(action)
}

func (session *Session) validateDomain(to string) error {
	toAddr, err := mail.ParseAddress(to)
	if err != nil {
		return fmt.Errorf("failed to parse address: %w", err)
	}

	pieces := strings.Split(toAddr.Address, "@")
	if len(pieces) != 2 {
		return fmt.Errorf("invalid address")
	}

	for _, d := range session.b.AllowedDomains {
		if strings.Contains(pieces[2], d) {
			return nil
		}
	}

	return fmt.Errorf("invalid domain")
}

func parseAction(action *milter.Action) error {
	switch action.Code {
	case milter.ActAccept:
		return nil
	case milter.ActContinue:
		return nil
	case milter.ActDiscard:
		return milterReject(action)
	case milter.ActReject:
		return milterReject(action)
	case milter.ActTempFail:
		return milterReject(action)
	case milter.ActReplyCode:
		return milterReject(action)
	case milter.ActSkip:
		return milterReject(action)
	}

	return fmt.Errorf("unhandled action code %q", action.Code)
}

func milterReject(action *milter.Action) error {
	return &smtp.SMTPError{
		Code:         action.SMTPCode,
		EnhancedCode: [3]int{4, 2, 0},
		Message:      action.SMTPText,
	}
}