// Package smtpsrv implements chasquid's SMTP server and connection handler.
package smtpsrv

import (
	"crypto"
	"crypto/ed25519"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"flag"
	"fmt"
	"math"
	"net"
	"net/http"
	"net/url"
	"os"
	"path"
	"strings"
	"sync/atomic"
	"syscall"
	"time"

	"blitiri.com.ar/go/chasquid/internal/aliases"
	"blitiri.com.ar/go/chasquid/internal/auth"
	"blitiri.com.ar/go/chasquid/internal/courier"
	"blitiri.com.ar/go/chasquid/internal/dkim"
	"blitiri.com.ar/go/chasquid/internal/domaininfo"
	"blitiri.com.ar/go/chasquid/internal/expvarom"
	"blitiri.com.ar/go/chasquid/internal/localrpc"
	"blitiri.com.ar/go/chasquid/internal/maillog"
	"blitiri.com.ar/go/chasquid/internal/queue"
	"blitiri.com.ar/go/chasquid/internal/set"
	"blitiri.com.ar/go/chasquid/internal/trace"
	"blitiri.com.ar/go/chasquid/internal/userdb"
	"blitiri.com.ar/go/log"
)

var (
	// Reload frequency.
	// We should consider making this a proper option if there's interest in
	// changing it, but until then, it's a test-only flag for simplicity.
	reloadEvery = flag.Duration("testing__reload_every", 30*time.Second,
		"how often to reload, ONLY FOR TESTING")
)

// Backoff between Accept errors.
const (
	acceptMinBackoff = 5 * time.Millisecond
	acceptMaxBackoff = 1 * time.Second
)

// Exported variables.
var (
	acceptErrors = expvarom.NewInt("chasquid/smtpIn/acceptErrors",
		"count of errors accepting incoming connections")
	activeConnections = expvarom.NewInt("chasquid/smtpIn/activeConnections",
		"number of active incoming connections")
	connectionsOverLimit = expvarom.NewInt(
		"chasquid/smtpIn/connectionsOverLimit",
		"count of incoming connections closed due to the connection limit")
)

// Server represents an SMTP server instance.
type Server struct {
	// Main hostname, used for display only.
	Hostname string

	// Maximum data size.
	MaxDataSize int64

	// Addresses.
	addrs map[SocketMode][]string

	// Listeners (that came via systemd).
	listeners map[SocketMode][]net.Listener

	// TLS config (including loaded certificates).
	tlsConfig *tls.Config

	// Use HAProxy on incoming connections.
	HAProxyEnabled bool

	// Local domains.
	localDomains *set.String

	// User databases (per domain).
	// Authenticator.
	authr *auth.Authenticator

	// Aliases resolver.
	aliasesR *aliases.Resolver

	// Domain info database.
	dinfo *domaininfo.DB

	// Map of domain -> DKIM signers.
	dkimSigners map[string][]*dkim.Signer

	// Time before we give up on a connection, even if it's sending data.
	connTimeout time.Duration

	// Time we wait for command round-trips (excluding DATA).
	commandTimeout time.Duration

	// Queue where we put incoming mail.
	queue *queue.Queue

	// Path to the hooks.
	HookPath string

	// Maximum number of concurrent incoming connections (0 = no limit), and
	// how many there are right now.
	maxConns    int64
	activeConns atomic.Int64

	// When we last logged about reaching the connection limit (Unix time).
	connLimitLogged atomic.Int64
}

// NewServer returns a new empty Server.
func NewServer() *Server {
	authr := auth.NewAuthenticator()
	aliasesR := aliases.NewResolver(authr.Exists)
	return &Server{
		addrs:     map[SocketMode][]string{},
		listeners: map[SocketMode][]net.Listener{},

		// Disable session tickets for now, to workaround a Microsoft bug
		// causing deliverability issues.
		//
		// See https://github.com/golang/go/issues/70232 for more details.
		//
		// This doesn't impact security, it just makes the re-establishment of
		// TLS sessions a bit slower, but for a server like chasquid it's not
		// going to be significant.
		//
		// Note this is not a Go-specific problem, and affects other servers
		// too (like Postfix/OpenSSL). This is a Microsoft problem that they
		// need to fix. Unfortunately, because they're quite a big provider
		// and are not very responsive in fixing their problems, we have to do
		// a workaround here.
		// TODO: Remove this once Microsoft fixes their servers.
		tlsConfig: &tls.Config{
			SessionTicketsDisabled: true,
		},

		connTimeout:    20 * time.Minute,
		commandTimeout: 1 * time.Minute,
		maxConns:       connLimit(),
		localDomains:   &set.String{},
		authr:          authr,
		aliasesR:       aliasesR,
		dkimSigners:    map[string][]*dkim.Signer{},
	}
}

// AddCerts (TLS) to the server.
func (s *Server) AddCerts(certPath, keyPath string) error {
	cert, err := tls.LoadX509KeyPair(certPath, keyPath)
	if err != nil {
		return err
	}
	s.tlsConfig.Certificates = append(s.tlsConfig.Certificates, cert)
	return nil
}

// AddAddr adds an address for the server to listen on.
func (s *Server) AddAddr(a string, m SocketMode) {
	s.addrs[m] = append(s.addrs[m], a)
}

// AddListeners adds listeners for the server to listen on.
func (s *Server) AddListeners(ls []net.Listener, m SocketMode) {
	s.listeners[m] = append(s.listeners[m], ls...)
}

// AddDomain adds a local domain to the server.
func (s *Server) AddDomain(d string) {
	s.localDomains.Add(d)
	s.aliasesR.AddDomain(d)
}

// AddUserDB adds a userdb file as backend for the domain.
func (s *Server) AddUserDB(domain, f string) (int, error) {
	// Load the userdb, and register it unconditionally (so reload works even
	// if there are errors right now).
	udb, err := userdb.Load(f)
	s.authr.Register(domain, auth.WrapNoErrorBackend(udb))
	return udb.Len(), err
}

// AddAliasesFile adds an aliases file for the given domain.
func (s *Server) AddAliasesFile(domain, f string) (int, error) {
	return s.aliasesR.AddAliasesFile(domain, f)
}

var (
	errDecodingPEMBlock     = fmt.Errorf("error decoding PEM block")
	errUnsupportedBlockType = fmt.Errorf("unsupported block type")
	errUnsupportedKeyType   = fmt.Errorf("unsupported key type")
)

// AddDKIMSigner for the given domain and selector.
func (s *Server) AddDKIMSigner(domain, selector, keyPath string) error {
	key, err := os.ReadFile(keyPath)
	if err != nil {
		return err
	}

	block, _ := pem.Decode(key)
	if block == nil {
		return errDecodingPEMBlock
	}

	if strings.ToUpper(block.Type) != "PRIVATE KEY" {
		return fmt.Errorf("%w: %s", errUnsupportedBlockType, block.Type)
	}

	signer, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		return err
	}

	switch k := signer.(type) {
	case *rsa.PrivateKey, ed25519.PrivateKey:
		// These are supported, nothing to do.
	default:
		return fmt.Errorf("%w: %T", errUnsupportedKeyType, k)
	}

	s.dkimSigners[domain] = append(s.dkimSigners[domain], &dkim.Signer{
		Domain:   domain,
		Selector: selector,
		Signer:   signer.(crypto.Signer),
	})
	return nil
}

// SetAuthFallback sets the authentication backend to use as fallback.
func (s *Server) SetAuthFallback(be auth.Backend) {
	s.authr.Fallback = be
}

// SetAliasesConfig sets the aliases configuration options.
func (s *Server) SetAliasesConfig(suffixSep, dropChars string) {
	s.aliasesR.SuffixSep = suffixSep
	s.aliasesR.DropChars = dropChars
	s.aliasesR.ResolveHook = path.Join(s.HookPath, "alias-resolve")
}

// SetDomainInfo sets the domain info database to use.
func (s *Server) SetDomainInfo(dinfo *domaininfo.DB) {
	s.dinfo = dinfo
}

// InitQueue initializes the queue.
func (s *Server) InitQueue(path string, localC, remoteC courier.Courier) {
	q, err := queue.New(path, s.localDomains, s.aliasesR, localC, remoteC)
	if err != nil {
		log.Fatalf("Error initializing queue: %v", err)
	}

	err = q.Load()
	if err != nil {
		log.Fatalf("Error loading queue: %v", err)
	}
	s.queue = q

	http.HandleFunc("GET /debug/queue",
		func(w http.ResponseWriter, r *http.Request) {
			_, _ = w.Write([]byte(q.DumpString()))
		})
}

func (s *Server) SetQueueLimits(maxItems uint32, giveUpAfter time.Duration) {
	if maxItems > 0 {
		s.queue.MaxItems = int(maxItems)
	}
	if giveUpAfter > 0 {
		s.queue.GiveUpAfter = giveUpAfter
	}
}

func (s *Server) aliasResolveRPC(tr *trace.Trace, req url.Values) (url.Values, error) {
	rcpts, err := s.aliasesR.Resolve(tr, req.Get("Address"))
	if err != nil {
		return nil, err
	}

	v := url.Values{}
	for _, rcpt := range rcpts {
		v.Add(string(rcpt.Type), rcpt.Addr)
	}

	return v, nil
}

func (s *Server) dinfoClearRPC(tr *trace.Trace, req url.Values) (url.Values, error) {
	domain := req.Get("Domain")
	exists := s.dinfo.Clear(tr, domain)
	if !exists {
		return nil, fmt.Errorf("does not exist")
	}
	return nil, nil
}

// periodicallyReload some of the server's information that can be changed
// without the server knowing, such as aliases and the user databases.
func (s *Server) periodicallyReload() {
	if reloadEvery == nil {
		return
	}

	for range time.Tick(*reloadEvery) {
		s.Reload()
	}
}

func (s *Server) Reload() {
	// Note that any error while reloading is fatal: this way, if there is an
	// unexpected error it can be detected (and corrected) quickly, instead of
	// much later (e.g. upon restart) when it might be harder to debug.
	if err := s.aliasesR.Reload(); err != nil {
		log.Fatalf("Error reloading aliases: %v", err)
	}

	if err := s.authr.Reload(); err != nil {
		log.Fatalf("Error reloading authenticators: %v", err)
	}
}

// ListenAndServe on the addresses and listeners that were previously added.
// This function will not return.
func (s *Server) ListenAndServe() {
	if len(s.tlsConfig.Certificates) == 0 {
		// chasquid assumes there's at least one valid certificate (for things
		// like STARTTLS, user authentication, etc.), so we fail if none was
		// found.
		log.Errorf("No SSL/TLS certificates found")
		log.Errorf("Ideally there should be a certificate for each MX you act as")
		log.Fatalf("At least one valid certificate is needed")
	}

	localrpc.DefaultServer.Register("AliasResolve", s.aliasResolveRPC)
	localrpc.DefaultServer.Register("DomaininfoClear", s.dinfoClearRPC)

	go s.periodicallyReload()

	if s.maxConns > 0 {
		log.Infof("Incoming connection limit: %d", s.maxConns)
	} else {
		log.Infof("No incoming connection limit")
	}

	for m, addrs := range s.addrs {
		for _, addr := range addrs {
			l, err := net.Listen("tcp", addr)
			if err != nil {
				log.Fatalf("Error listening: %v", err)
			}

			log.Infof("Server listening on %s (%v)", addr, m)
			maillog.Listening(addr)
			go s.serve(l, m)
		}
	}

	for m, ls := range s.listeners {
		for _, l := range ls {
			log.Infof("Server listening on %s (%v, via systemd)", l.Addr(), m)
			maillog.Listening(l.Addr().String())
			go s.serve(l, m)
		}
	}

	// Never return. The serve goroutines only return if their listener is
	// closed, which doesn't happen in practice.
	select {}
}

func (s *Server) serve(l net.Listener, mode SocketMode) {
	// If this mode is expected to be TLS-wrapped, make it so.
	if mode.TLS {
		l = tls.NewListener(l, s.tlsConfig)
	}

	pdhook := path.Join(s.HookPath, "post-data")

	backoff := acceptMinBackoff
	var lastErrLog time.Time

	for {
		conn, err := l.Accept()
		if err != nil {
			if errors.Is(err, net.ErrClosed) {
				// We never close the listeners, so this is unexpected.
				log.Errorf("Listener %s closed, no longer accepting: %v",
					l.Addr(), err)
				return
			}

			// Other errors are usually transient (e.g. running out of file
			// descriptors), so we back off and retry instead of exiting,
			// which would make it trivial to take the server down.
			// Log at most once a minute to avoid flooding the logs.
			acceptErrors.Add(1)
			if time.Since(lastErrLog) >= time.Minute {
				log.Errorf("Error accepting on %s (will retry): %v",
					l.Addr(), err)
				lastErrLog = time.Now()
			}
			time.Sleep(backoff)
			backoff = min(backoff*2, acceptMaxBackoff)
			continue
		}
		backoff = acceptMinBackoff

		// Close the connection if we're over the limit. Log at most once a
		// minute (across all listeners) to avoid flooding the logs.
		if n := s.activeConns.Add(1); s.maxConns > 0 && n > s.maxConns {
			s.activeConns.Add(-1)
			connectionsOverLimit.Add(1)
			conn.Close()

			now, last := time.Now().Unix(), s.connLimitLogged.Load()
			if now-last >= 60 && s.connLimitLogged.CompareAndSwap(last, now) {
				log.Errorf("Incoming connection limit (%d) reached, "+
					"closing new connections (%d closed so far)",
					s.maxConns, connectionsOverLimit.Value())
			}
			continue
		}
		activeConnections.Add(1)

		sc := &Conn{
			hostname:       s.Hostname,
			maxDataSize:    s.MaxDataSize,
			postDataHook:   pdhook,
			conn:           conn,
			mode:           mode,
			tlsConfig:      s.tlsConfig,
			haproxyEnabled: s.HAProxyEnabled,
			onTLS:          mode.TLS,
			authr:          s.authr,
			aliasesR:       s.aliasesR,
			localDomains:   s.localDomains,
			dinfo:          s.dinfo,
			dkimSigners:    s.dkimSigners,
			deadline:       time.Now().Add(s.connTimeout),
			commandTimeout: s.commandTimeout,
			queue:          s.queue,
		}
		go func() {
			sc.Handle()
			s.activeConns.Add(-1)
			activeConnections.Add(-1)
		}()
	}
}

// connLimit returns the maximum number of concurrent incoming connections we
// should allow, or 0 for no limit.
//
// It's half of the file descriptor limit, to leave room for everything else
// that needs them (queue files, outgoing connections, hooks, etc.), so a large
// number of incoming connections doesn't prevent the rest of the server from
// working.
func connLimit() int64 {
	var lim syscall.Rlimit
	if err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, &lim); err != nil {
		return 0
	}

	// The type of lim.Cur varies across platforms (int64 or uint64). Values
	// out of range (e.g. RLIM_INFINITY) mean there's no meaningful limit.
	cur := int64(lim.Cur)
	if cur <= 0 || cur > math.MaxInt32 {
		return 0
	}
	return cur / 2
}