gomail/internal/pipeline/pipeline.go

// Package pipeline implements the inbound security pipeline: SPF, DKIM
// verification, DMARC, header/URL heuristics — each stage contributes a
// score, and the orchestrator maps the total score to a verdict (clean,
// flagged, quarantine, blocked) per the configured thresholds.
//
// Every stage is blocking and runs before the SMTP DATA response — no
// third-party spam-filtering library, entirely stdlib DNS/crypto/net/mail.
package pipeline

import (
	"context"
	"net"
	"net/mail"
	"strings"
	"time"

	"gomail/internal/config"
	"gomail/internal/db"
)

// MailContext carries everything a stage needs and accumulates results.
type MailContext struct {
	SenderIP   net.IP
	SenderHost string
	MailFrom   string
	RcptTo     string
	RawMessage []byte

	Checks     []StageResult
	TotalScore float64
	Verdict    db.MessageVerdict

	parsedMessage  *mail.Message
	parseErr       error
	parseAttempted bool
}

// ParsedMessage lazily parses RawMessage via net/mail — stages call this
// instead of parsing independently, so the (relatively expensive) header
// parse happens at most once per message regardless of how many stages need it.
func (mc *MailContext) ParsedMessage() (*mail.Message, error) {
	if !mc.parseAttempted {
		mc.parsedMessage, mc.parseErr = mail.ReadMessage(strings.NewReader(string(mc.RawMessage)))
		mc.parseAttempted = true
	}
	return mc.parsedMessage, mc.parseErr
}

// RcptDomain returns the domain portion of RcptTo.
func (mc *MailContext) RcptDomain() string {
	return domainOf(mc.RcptTo)
}

// MailFromDomain returns the domain portion of the envelope sender.
func (mc *MailContext) MailFromDomain() string {
	return domainOf(mc.MailFrom)
}

func domainOf(addr string) string {
	parts := strings.SplitN(strings.ToLower(addr), "@", 2)
	if len(parts) == 2 {
		return parts[1]
	}
	return ""
}

// StageResult is one check's outcome.
type StageResult struct {
	Stage      string
	Result     db.CheckResult
	Score      float64
	Detail     string
	DurationMs int64
}

// Stage is one pipeline check. Run must not block indefinitely — pass a
// context with a deadline and respect it for any network I/O (DNS lookups).
type Stage interface {
	Name() string
	Run(ctx context.Context, mc *MailContext) *StageResult
}

// Orchestrator runs the configured stages in order and computes the verdict.
type Orchestrator struct {
	stages []Stage
	cfg    *config.Config
}

func NewOrchestrator(cfg *config.Config, stages []Stage) *Orchestrator {
	return &Orchestrator{stages: stages, cfg: cfg}
}

// DefaultStages returns the deterministic, always-on stage set (SPF, DKIM,
// DMARC, header anomaly, URL heuristics) — no external service
// dependencies, always safe to run regardless of what's configured.
func DefaultStages() []Stage {
	return []Stage{
		&SPFStage{},
		&DKIMStage{},
		&DMARCStage{},
		&HeaderStage{},
		&URLStage{},
	}
}

// StagesFromConfig returns DefaultStages() plus any of the optional
// external-service stages (ClamAV, Rspamd, LLM) that config has an address
// configured for — each is entirely absent from the pipeline, not merely
// disabled, when its config field is empty, so an unreachable/misconfigured
// service that was never intended to be used can't accidentally affect
// delivery.
func StagesFromConfig(cfg *config.Config) []Stage {
	stages := DefaultStages()

	if cfg.Pipeline.ClamAVSocket != "" {
		stages = append(stages, &ClamAVStage{Addr: cfg.Pipeline.ClamAVSocket, Timeout: 30 * time.Second})
	}
	if cfg.Pipeline.RspamdURL != "" {
		stages = append(stages, &RspamdStage{BaseURL: cfg.Pipeline.RspamdURL, Timeout: 15 * time.Second})
	}
	if cfg.Pipeline.LLMURL != "" {
		timeout := time.Duration(cfg.Pipeline.LLMTimeoutSecs) * time.Second
		if timeout <= 0 {
			timeout = 30 * time.Second
		}
		stages = append(stages, &LLMStage{BaseURL: cfg.Pipeline.LLMURL, Model: cfg.Pipeline.LLMModel, Timeout: timeout})
	}

	return stages
}

// Run executes every stage in order, accumulating score, and computes the
// final verdict against the configured thresholds. Individual stage panics
// are not recovered here deliberately — a panicking stage is a bug that
// should surface loudly in testing, not be silently swallowed in production
// and misclassify mail.
func (o *Orchestrator) Run(ctx context.Context, mc *MailContext) {
	for _, stage := range o.stages {
		start := time.Now()
		result := stage.Run(ctx, mc)
		if result == nil {
			continue
		}
		result.DurationMs = time.Since(start).Milliseconds()
		mc.Checks = append(mc.Checks, *result)
		mc.TotalScore += result.Score
	}

	mc.Verdict = verdictFor(mc.TotalScore, o.cfg.Pipeline)
}

func verdictFor(score float64, p config.PipelineConfig) db.MessageVerdict {
	switch {
	case score >= p.ScoreBlock:
		return db.VerdictBlocked
	case score >= p.ScoreQuarantine:
		return db.VerdictQuarantine
	case score >= p.ScoreFlag:
		return db.VerdictFlagged
	default:
		return db.VerdictClean
	}
}