mailgosend/internal/smtp/submission.go

package smtp

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"log"
	"net"
	"net/mail"
	"strings"
	"time"

	gosmtp "github.com/emersion/go-smtp"

	"ghb.freebede.com/nahakubuilder/mailgosend/internal/crypto"
	"ghb.freebede.com/nahakubuilder/mailgosend/internal/dkim"
	"ghb.freebede.com/nahakubuilder/mailgosend/internal/models"
	"ghb.freebede.com/nahakubuilder/mailgosend/internal/storage"
)

// SubmissionBackend implements gosmtp.Backend for ports 587/465.
// Supports: regular users, relay accounts (SMTP-only), and unauthenticated IP relay.
type SubmissionBackend struct {
	deps *Deps
}

func (b *SubmissionBackend) NewSession(c *gosmtp.Conn) (gosmtp.Session, error) {
	clientIP, _, _ := net.SplitHostPort(c.Conn().RemoteAddr().String())

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	var banned bool
	if b.deps.Brute != nil {
		banned, _ = b.deps.Brute.IsBanned(ctx, clientIP)
	}
	if banned {
		return nil, fmt.Errorf("connection refused")
	}

	return &SubmissionSession{
		deps:     b.deps,
		clientIP: clientIP,
	}, nil
}

// SubmissionSession handles one authenticated or IP-relay submission connection.
type SubmissionSession struct {
	deps         *Deps
	clientIP     string
	user         *models.User         // set after AUTH as regular user
	relayAccount *models.RelayAccount // set after AUTH as relay account
	ipRelayMode  bool                 // set when IP relay authorization succeeds (no AUTH)
	from         string
	rcpts        []string
}

func (s *SubmissionSession) AuthPlain(username, password string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	// Try regular user first.
	user, err := s.deps.DB.GetUserByEmail(ctx, username)
	if err != nil {
		log.Printf("[smtp/submission] auth lookup error %s: %v", username, err)
		return &gosmtp.SMTPError{Code: 535, Message: "authentication failed"}
	}
	if user != nil && user.Enabled {
		if err := crypto.CheckPassword(user.PasswordHash, password); err != nil {
			s.logAttempt(ctx, username, false)
			log.Printf("[smtp/submission] auth failed for user %s from %s", username, s.clientIP)
			return &gosmtp.SMTPError{Code: 535, Message: "authentication failed"}
		}
		s.user = user
		s.deps.DB.UpdateLastLogin(ctx, user.ID)
		log.Printf("[smtp/submission] auth OK for user %s from %s", username, s.clientIP)
		return nil
	}

	// Try relay account.
	relayAcc, err := s.deps.DB.GetRelayAccountByUsername(ctx, username)
	if err != nil {
		log.Printf("[smtp/submission] relay auth lookup error %s: %v", username, err)
		return &gosmtp.SMTPError{Code: 535, Message: "authentication failed"}
	}
	if relayAcc == nil || !relayAcc.Enabled {
		s.logAttempt(ctx, username, false)
		log.Printf("[smtp/submission] auth failed for relay %s from %s", username, s.clientIP)
		return &gosmtp.SMTPError{Code: 535, Message: "authentication failed"}
	}
	if err := crypto.CheckPassword(relayAcc.PasswordHash, password); err != nil {
		s.logAttempt(ctx, username, false)
		log.Printf("[smtp/submission] auth failed for relay %s from %s", username, s.clientIP)
		return &gosmtp.SMTPError{Code: 535, Message: "authentication failed"}
	}
	s.relayAccount = relayAcc
	log.Printf("[smtp/submission] relay auth OK for %s from %s", username, s.clientIP)
	return nil
}

func (s *SubmissionSession) Mail(from string, opts *gosmtp.MailOptions) error {
	addr, err := mail.ParseAddress(from)
	if err != nil {
		return &gosmtp.SMTPError{Code: 501, EnhancedCode: gosmtp.EnhancedCode{5, 1, 7}, Message: "invalid sender"}
	}
	fromEmail := strings.ToLower(addr.Address)

	if s.user == nil && s.relayAccount == nil {
		// Unauthenticated — check relay account IP whitelist first, then domain-level IP rules.
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()

		acc, err := s.deps.DB.FindRelayAccountByIP(ctx, s.clientIP, fromEmail)
		if err != nil {
			log.Printf("[smtp/submission] relay account ip check from %s: %v", s.clientIP, err)
			return &gosmtp.SMTPError{Code: 451, EnhancedCode: gosmtp.EnhancedCode{4, 3, 0}, Message: "temporary error"}
		}
		if acc != nil {
			// Treat as authenticated relay account — no password required.
			s.relayAccount = acc
			s.from = addr.Address
			log.Printf("[smtp/submission] ip-relay (account %s) from %s as %s", acc.Username, s.clientIP, fromEmail)
			return nil
		}

		allowed, err := s.deps.DB.CheckIPRelay(ctx, s.clientIP, fromEmail)
		if err != nil {
			log.Printf("[smtp/submission] ip relay check error from %s: %v", s.clientIP, err)
			return &gosmtp.SMTPError{Code: 451, EnhancedCode: gosmtp.EnhancedCode{4, 3, 0}, Message: "temporary error"}
		}
		if !allowed {
			return &gosmtp.SMTPError{Code: 530, Message: "authentication required"}
		}
		s.ipRelayMode = true
		s.from = addr.Address
		return nil
	}

	if s.relayAccount != nil {
		// Relay account — validate sender against allowed send-as patterns.
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()

		allowed, err := s.deps.DB.IsRelaySenderAllowed(ctx, s.relayAccount.ID, fromEmail)
		if err != nil {
			log.Printf("[smtp/submission] relay sender check error for account %d: %v", s.relayAccount.ID, err)
			return &gosmtp.SMTPError{Code: 451, EnhancedCode: gosmtp.EnhancedCode{4, 3, 0}, Message: "temporary error"}
		}
		if !allowed {
			return &gosmtp.SMTPError{Code: 553, EnhancedCode: gosmtp.EnhancedCode{5, 1, 8}, Message: "sender address not permitted for this relay account"}
		}
		s.from = addr.Address
		return nil
	}

	// Regular user — sender must be own email or an alias they own.
	if !strings.EqualFold(fromEmail, s.user.Email) {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()

		resolved, err := s.deps.DB.ResolveEmail(ctx, fromEmail)
		if err != nil || resolved == nil || resolved.ID != s.user.ID {
			return &gosmtp.SMTPError{Code: 553, EnhancedCode: gosmtp.EnhancedCode{5, 1, 8}, Message: "sender not owned by authenticated user"}
		}
	}

	s.from = addr.Address
	return nil
}

func (s *SubmissionSession) Rcpt(to string, opts *gosmtp.RcptOptions) error {
	if s.user == nil && s.relayAccount == nil && !s.ipRelayMode {
		return &gosmtp.SMTPError{Code: 530, Message: "authentication required"}
	}

	addr, err := mail.ParseAddress(to)
	if err != nil {
		return &gosmtp.SMTPError{Code: 501, EnhancedCode: gosmtp.EnhancedCode{5, 1, 3}, Message: "invalid recipient"}
	}

	s.rcpts = append(s.rcpts, addr.Address)
	return nil
}

func (s *SubmissionSession) Data(r io.Reader) error {
	if s.user == nil && s.relayAccount == nil && !s.ipRelayMode {
		return &gosmtp.SMTPError{Code: 530, Message: "authentication required"}
	}
	if len(s.rcpts) == 0 {
		return &gosmtp.SMTPError{Code: 503, Message: "no recipients"}
	}

	raw, err := io.ReadAll(io.LimitReader(r, s.deps.Cfg.MaxMessageSize+1))
	if err != nil {
		return fmt.Errorf("read submission data: %w", err)
	}
	if int64(len(raw)) > s.deps.Cfg.MaxMessageSize {
		return &gosmtp.SMTPError{Code: 552, EnhancedCode: gosmtp.EnhancedCode{5, 3, 4}, Message: "message too large"}
	}

	_, err = mail.ReadMessage(bytes.NewReader(raw))
	if err != nil {
		return &gosmtp.SMTPError{Code: 550, EnhancedCode: gosmtp.EnhancedCode{5, 6, 0}, Message: "malformed message"}
	}

	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	senderDomain := domainOf(s.from)
	raw = s.signDKIM(ctx, raw, senderDomain)
	msgID := extractMsgID(raw)

	dom, err := s.deps.DB.GetDomain(ctx, senderDomain)
	var domainID int64
	if err == nil && dom != nil {
		domainID = dom.ID
	}

	queueKey, err := s.deps.Crypt.DeriveKeyGlobal("queue")
	if err != nil {
		return fmt.Errorf("queue key: %w", err)
	}
	rawEnc, err := crypto.Encrypt(queueKey, raw)
	if err != nil {
		return fmt.Errorf("encrypt for queue: %w", err)
	}

	for _, rcpt := range s.rcpts {
		_, err := s.deps.DB.EnqueueMessage(ctx, domainID, s.from, rcpt, msgID, rawEnc, s.deps.Cfg.QueueMaxAgeHours)
		if err != nil {
			log.Printf("[smtp/submission] enqueue to %s: %v", rcpt, err)
			return &gosmtp.SMTPError{Code: 451, EnhancedCode: gosmtp.EnhancedCode{4, 3, 0}, Message: "queue error"}
		}
		log.Printf("[smtp/submission] queued %s → %s", s.from, rcpt)
	}

	// Save Sent copy only for regular users (relay accounts have no mailboxes).
	if s.user != nil {
		s.saveSentCopy(ctx, raw)
	}

	return nil
}

func (s *SubmissionSession) Reset() {
	s.from = ""
	s.rcpts = s.rcpts[:0]
	s.ipRelayMode = false
}

func (s *SubmissionSession) Logout() error { return nil }

// signDKIM signs the message with the sender domain's DKIM key if available.
func (s *SubmissionSession) signDKIM(ctx context.Context, raw []byte, senderDomain string) []byte {
	dom, err := s.deps.DB.GetDomain(ctx, senderDomain)
	if err != nil || dom == nil || dom.DKIMPrivateEnc == nil {
		return raw
	}

	privPEM, err := s.deps.Crypt.DecryptGlobal("dkim", dom.DKIMPrivateEnc)
	if err != nil {
		log.Printf("[smtp/submission] dkim key decrypt for %s: %v", senderDomain, err)
		return raw
	}

	signer, err := dkim.NewSigner(string(privPEM), senderDomain, dom.DKIMSelector)
	if err != nil {
		log.Printf("[smtp/submission] dkim signer for %s: %v", senderDomain, err)
		return raw
	}

	header, err := signer.Sign(raw)
	if err != nil {
		log.Printf("[smtp/submission] dkim sign for %s: %v", senderDomain, err)
		return raw
	}

	return append([]byte(header+"\r\n"), raw...)
}

// saveSentCopy stores a copy in the user's Sent mailbox (best-effort).
func (s *SubmissionSession) saveSentCopy(ctx context.Context, raw []byte) {
	mbox, err := s.deps.DB.GetMailboxByType(ctx, s.user.ID, models.MailboxSent)
	if err != nil || mbox == nil {
		return
	}

	msg, err := mail.ReadMessage(strings.NewReader(string(raw)))
	if err != nil {
		return
	}

	subject := msg.Header.Get("Subject")
	msgDate, _ := mail.ParseDate(msg.Header.Get("Date"))
	if msgDate.IsZero() {
		msgDate = time.Now().UTC()
	}

	incoming := &storage.IncomingMessage{
		Raw:       raw,
		FromEmail: s.from,
		ToList:    strings.Join(s.rcpts, ", "),
		Subject:   subject,
		Date:      msgDate,
		MessageID: extractMsgID(raw),
		SpamScore: 0,
	}

	if _, err := s.deps.Store.SaveIncoming(ctx, s.user.ID, mbox.ID, incoming); err != nil {
		log.Printf("[smtp/submission] save sent copy: %v", err)
	}
}

// ---- helpers ----

func (s *SubmissionSession) logAttempt(ctx context.Context, email string, success bool) {
	if s.deps.Brute != nil {
		s.deps.Brute.RecordAttempt(ctx, s.clientIP, email, success) //nolint:errcheck
	}
}

func domainOf(email string) string {
	at := strings.LastIndex(email, "@")
	if at < 0 {
		return ""
	}
	return strings.ToLower(email[at+1:])
}

func extractMsgID(raw []byte) string {
	msg, err := mail.ReadMessage(bytes.NewReader(raw))
	if err != nil {
		return ""
	}
	return msg.Header.Get("Message-ID")
}