Havline/internal/service/proxy.go

package service

import (
	"context"
	"database/sql"
	"fmt"
	"strconv"

	"github.com/havline/havline/internal/certificate"
	"github.com/havline/havline/internal/config"
	"github.com/havline/havline/internal/nginx"
	"github.com/havline/havline/internal/notify"
	"github.com/havline/havline/internal/proxy"
	"github.com/havline/havline/internal/settings"
	"github.com/havline/havline/internal/validate"
)

type ProxyService struct {
	db        *sql.DB
	cfg       config.Config
	store     *proxy.Store
	certStore *certificate.Store
	settings  *settings.Store
	nginx     *nginx.Manager
	notify    *notify.Service
	exitSync  ExitSync
	traffic   TrafficCleaner
}

// ExitSync 处理服务发布规则变化后的外部出口同步。
type ExitSync interface {
	SyncRules(ctx context.Context) error
}

// TrafficCleaner 清掉某条规则对应的流量累计（内存 + 落库）。
// 用接口注入而不是直接依赖 traffic 包：service 是编排层，不该反过来知道采集细节。
// 参数用 proxy.Endpoint 而不是拼好的字符串，键格式由 traffic 包自己决定。
type TrafficCleaner interface {
	ForgetEndpoints(ctx context.Context, endpoints []proxy.Endpoint) error
}

func NewProxyService(cfg config.Config, db *sql.DB, store *proxy.Store, certStore *certificate.Store, settingsStore *settings.Store, nginxMgr *nginx.Manager) *ProxyService {
	return &ProxyService{db: db, cfg: cfg, store: store, certStore: certStore, settings: settingsStore, nginx: nginxMgr}
}

func (s *ProxyService) SetNotify(notifySvc *notify.Service) {
	s.notify = notifySvc
}

func (s *ProxyService) SetExitSync(sync ExitSync) {
	s.exitSync = sync
}

func (s *ProxyService) SetTrafficCleaner(cleaner TrafficCleaner) {
	s.traffic = cleaner
}

func (s *ProxyService) List(ctx context.Context) ([]proxy.Rule, error) {
	return s.store.List(ctx)
}

func (s *ProxyService) Get(ctx context.Context, id int64) (proxy.Rule, error) {
	return s.store.Get(ctx, id)
}

func (s *ProxyService) Create(ctx context.Context, in proxy.CreateInput) (proxy.Rule, error) {
	if proxy.HasExit(in.Exits, proxy.ExitLocal) {
		if err := s.validateHTTPSInput(ctx, in.Hosts, in.HTTPSEnabled); err != nil {
			return proxy.Rule{}, err
		}
	}

	rule, err := s.store.Create(ctx, in)
	if err != nil {
		return proxy.Rule{}, err
	}
	if err := s.applyNginx(ctx); err != nil {
		return rule, fmt.Errorf("规则已保存，但 Nginx 重载失败：%w", err)
	}
	if err := s.syncExits(ctx); err != nil {
		rule.SyncWarning = "规则已保存，但 Cloudflare 同步失败：" + err.Error()
	}
	return rule, nil
}

func (s *ProxyService) Update(ctx context.Context, id int64, in proxy.UpdateInput) (proxy.Rule, error) {
	current, err := s.Get(ctx, id)
	if err != nil {
		return proxy.Rule{}, err
	}
	httpsEnabled := current.HTTPSEnabled
	if in.HTTPSEnabled != nil {
		httpsEnabled = *in.HTTPSEnabled
	}
	hosts := hostsFromRule(current)
	if in.Hosts != nil {
		hosts = *in.Hosts
	}
	exits := current.Exits
	if in.Exits != nil {
		exits = *in.Exits
	}
	if proxy.HasExit(exits, proxy.ExitLocal) {
		if err := s.validateHTTPSInput(ctx, hosts, httpsEnabled); err != nil {
			return proxy.Rule{}, err
		}
	}

	rule, err := s.store.Update(ctx, id, in)
	if err != nil {
		return proxy.Rule{}, err
	}
	if err := s.applyNginx(ctx); err != nil {
		return rule, fmt.Errorf("规则已保存，但 Nginx 重载失败：%w", err)
	}
	if err := s.syncExits(ctx); err != nil {
		rule.SyncWarning = "规则已保存，但 Cloudflare 同步失败：" + err.Error()
	}
	return rule, nil
}

func (s *ProxyService) syncExits(ctx context.Context) error {
	if s.exitSync == nil {
		return nil
	}
	return s.exitSync.SyncRules(ctx)
}

func (s *ProxyService) Reorder(ctx context.Context, ids []int64) error {
	return s.store.Reorder(ctx, ids)
}

func (s *ProxyService) Delete(ctx context.Context, id int64) error {
	// 先取出规则：删除后就查不到域名与端口，没法定位要清理的流量键
	rule, err := s.store.Get(ctx, id)
	if err != nil {
		return err
	}
	if err := s.store.Delete(ctx, id); err != nil {
		return err
	}
	_ = nginx.RemoveRuleCustom(s.cfg, id)
	s.forgetTraffic(ctx, rule)
	if err := s.applyNginx(ctx); err != nil {
		return fmt.Errorf("规则已删除，但 Nginx 重载失败：%w", err)
	}
	if err := s.syncExits(ctx); err != nil {
		return fmt.Errorf("规则已删除，但 Cloudflare 同步失败：%w", err)
	}
	return nil
}

// forgetTraffic 清掉规则对应端点的流量累计。失败不影响删除本身：
// 残留的统计行只是脏数据，而规则已经删了，没必要为此把删除操作报成失败。
func (s *ProxyService) forgetTraffic(ctx context.Context, rule proxy.Rule) {
	if s.traffic == nil {
		return
	}
	_ = s.traffic.ForgetEndpoints(ctx, rule.Endpoints())
}

func (s *ProxyService) ReloadAll(ctx context.Context) error {
	return s.applyNginx(ctx)
}

func (s *ProxyService) applyNginx(ctx context.Context) error {
	rules, err := s.store.ListEnabledLocal(ctx)
	if err != nil {
		return err
	}
	certs, err := s.loadCertSources(ctx)
	if err != nil {
		return err
	}
	s.nginx.SetGenerateOptions(s.loadGenerateOptions(ctx))
	_, err = s.nginx.Apply(ctx, rules, certs)
	if err != nil && s.notify != nil {
		s.notify.Alert(ctx, notify.EventNginxReloadFailure, "Nginx 重载失败", err.Error())
	}
	return err
}

func (s *ProxyService) ValidateRule(ctx context.Context, candidate proxy.Rule, existing []proxy.Rule) error {
	rules := make([]proxy.Rule, 0, len(existing)+1)
	replaced := false
	for _, r := range existing {
		if r.ID == candidate.ID {
			if candidate.LocalEnabled() {
				rules = append(rules, candidate)
			}
			replaced = true
			continue
		}
		if r.LocalEnabled() {
			rules = append(rules, r)
		}
	}
	if !replaced && candidate.LocalEnabled() {
		rules = append(rules, candidate)
	}
	certs, err := s.loadCertSources(ctx)
	if err != nil {
		return err
	}
	s.nginx.SetGenerateOptions(s.loadGenerateOptions(ctx))
	if err := s.nginx.ValidateOnly(ctx, rules, certs); err != nil {
		return err
	}
	return nil
}

func hostsFromRule(rule proxy.Rule) []string {
	hosts := make([]string, 0, len(rule.Hosts))
	for _, host := range rule.Hosts {
		if host.ListenPort != nil && *host.ListenPort > 0 && *host.ListenPort != rule.ListenPort {
			hosts = append(hosts, host.Hostname+":"+strconv.Itoa(*host.ListenPort))
			continue
		}
		hosts = append(hosts, host.Hostname)
	}
	return hosts
}

func (s *ProxyService) validateHTTPSInput(ctx context.Context, hosts []string, httpsEnabled bool) error {
	if !httpsEnabled || len(hosts) == 0 {
		return nil
	}
	certs, err := s.loadCertSources(ctx)
	if err != nil {
		return err
	}
	parsedHosts := make([]string, 0, len(hosts))
	for _, host := range hosts {
		hostname, _, err := validate.FrontendAddress(host)
		if err != nil {
			return err
		}
		parsedHosts = append(parsedHosts, hostname)
	}
	if nginx.HasCertificateForHosts(s.cfg.CertsDir(), parsedHosts, certs) {
		return nil
	}
	return nginx.HTTPSCoverageError(parsedHosts[0], certs)
}

func (s *ProxyService) loadCertSources(ctx context.Context) ([]nginx.CertSource, error) {
	records, err := s.certStore.List(ctx)
	if err != nil {
		return nil, err
	}
	sources := make([]nginx.CertSource, 0, len(records))
	for _, record := range records {
		sources = append(sources, nginx.CertSource{
			Domains:  record.Domains,
			CertPath: record.CertPath,
			KeyPath:  record.KeyPath,
		})
	}
	return sources, nil
}