gomail/internal/imap/server.go

// Package imap implements a hand-rolled IMAP server covering the core
// command set (RFC 3501/9051 essentials): CAPABILITY, LOGIN, LOGOUT, NOOP,
// SELECT/EXAMINE, LIST, FETCH, UID FETCH, STORE, UID STORE, SEARCH, EXPUNGE,
// CLOSE, UNSELECT. No third-party IMAP library — stdlib net.Listener plus a
// small hand-written parser for IMAP's atom/quoted-string/literal syntax.
//
// Deferred to a later pass (noted here so the gap is visible, not hidden):
// IDLE, CONDSTORE/QRESYNC, SORT/THREAD, and mailbox CREATE/DELETE/RENAME.
// The core set above is enough for read/flag/delete workflows against an
// existing mailbox, which covers most mail client usage; IDLE (push) and
// folder management are the natural next additions.
package imap

import (
	"context"
	"crypto/tls"
	"fmt"
	"log/slog"
	"net"
	"sync"
	"time"

	"gomail/internal/db"
	"gomail/internal/mailstore"
	"gomail/internal/ratelimit"
)

const (
	idleTimeout    = 30 * time.Minute // IMAP clients often sit connected much longer than SMTP
	maxCommandLine = 8192
)

type Server struct {
	database *db.DB
	store    *mailstore.Store
	tlsConf  *tls.Config
	hostname string

	listeners []net.Listener
	wg        sync.WaitGroup
	sessionWG sync.WaitGroup

	connLimiter *ratelimit.Limiter // per-IP connections/min
	authLimiter *ratelimit.Limiter // per-IP LOGIN failures/min — checked before credential verification
}

func NewServer(database *db.DB, store *mailstore.Store, tlsConf *tls.Config, hostname string, connPerMin, authFailuresPerMin int) *Server {
	return &Server{
		database:    database,
		store:       store,
		tlsConf:     tlsConf,
		hostname:    hostname,
		connLimiter: ratelimit.New(connPerMin),
		authLimiter: ratelimit.New(authFailuresPerMin),
	}
}

// ListenAndServe starts the plain (:143, STARTTLS-capable) and implicit-TLS
// (:993) listeners and blocks until ctx is cancelled or a listener fails.
func (s *Server) ListenAndServe(ctx context.Context, plainAddr, tlsAddr string) error {
	specs := []struct {
		addr   string
		useTLS bool
	}{
		{plainAddr, false},
		{tlsAddr, true},
	}

	for _, spec := range specs {
		ln, err := net.Listen("tcp", spec.addr)
		if err != nil {
			s.closeAll()
			return fmt.Errorf("listen %s: %w", spec.addr, err)
		}
		if spec.useTLS {
			ln = tls.NewListener(ln, s.tlsConf)
		}
		s.listeners = append(s.listeners, ln)
		slog.Info("IMAP listener started", "addr", spec.addr, "implicit_tls", spec.useTLS)

		s.wg.Add(1)
		go func(ln net.Listener) {
			defer s.wg.Done()
			s.acceptLoop(ctx, ln)
		}(ln)
	}

	<-ctx.Done()
	return ctx.Err()
}

func (s *Server) acceptLoop(ctx context.Context, ln net.Listener) {
	for {
		conn, err := ln.Accept()
		if err != nil {
			select {
			case <-ctx.Done():
				return
			default:
				slog.Error("IMAP accept error", "err", err)
				return
			}
		}
		ip := connHost(conn.RemoteAddr())
		if !s.connLimiter.Allow(ip) {
			slog.Warn("IMAP connection rate limit exceeded, rejecting", "ip", ip)
			conn.Close()
			continue
		}

		s.sessionWG.Add(1)
		go func() {
			defer s.sessionWG.Done()
			sess := newSession(conn, s)
			sess.run(ctx)
		}()
	}
}

func (s *Server) Shutdown(gracePeriod time.Duration) {
	s.closeAll()
	done := make(chan struct{})
	go func() {
		s.sessionWG.Wait()
		close(done)
	}()
	select {
	case <-done:
		slog.Info("all IMAP sessions drained cleanly")
	case <-time.After(gracePeriod):
		slog.Warn("IMAP shutdown grace period expired — some sessions forcibly terminated")
	}
}

func (s *Server) closeAll() {
	for _, ln := range s.listeners {
		ln.Close()
	}
	s.wg.Wait()
}

// connHost extracts just the IP (no port) from a net.Addr, for use as a
// rate-limiter key.
func connHost(addr net.Addr) string {
	host, _, err := net.SplitHostPort(addr.String())
	if err != nil {
		return addr.String()
	}
	return host
}