Prism/cmd/prism/main.go

package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"syscall"
	"time"

	"github.com/robfig/cron/v3"

	"github.com/prism/proxy/internal/api"
	"github.com/prism/proxy/internal/cache"
	"github.com/prism/proxy/internal/configbuilder"
	"github.com/prism/proxy/internal/core"
	"github.com/prism/proxy/internal/geodata"
	"github.com/prism/proxy/internal/health"
	applog "github.com/prism/proxy/internal/log"
	"github.com/prism/proxy/internal/metrics"
	"github.com/prism/proxy/internal/mihomoapi"
	"github.com/prism/proxy/internal/service"
	"github.com/prism/proxy/internal/store"
	"github.com/prism/proxy/internal/subscription"
	"github.com/prism/proxy/internal/trafficlog"
)

func main() {
	dataDir := flag.String("data", "data", "data directory")
	flag.Parse()

	log.Printf("Prism starting, data dir: %s", *dataDir)

	if err := os.MkdirAll(*dataDir, 0o755); err != nil {
		log.Fatal(err)
	}

	st, err := store.Open(*dataDir)
	if err != nil {
		log.Fatal(err)
	}
	defer st.Close()
	log.Printf("database ready")

	if n, err := st.BackfillNodeCountries(context.Background()); err != nil {
		log.Printf("country backfill: %v", err)
	} else if n > 0 {
		log.Printf("country tags backfilled for %d nodes", n)
	}

	core.InitDataDir(*dataDir)

	c := cache.New()
	engine := core.NewEngine()
	builder := configbuilder.New(st, *dataDir)
	logCollector := applog.NewCollector(st, 5000, true)
	subSvc := subscription.NewService(st)

	app := &service.App{
		Store:        st,
		Cache:        c,
		Engine:       engine,
		Builder:      builder,
		Subscription: subSvc,
		Logs:         logCollector,
		DataDir:      *dataDir,
	}

	ctx := context.Background()
	apiBase, secret, _ := app.MihomoAPIBase(ctx)
	hc := health.NewChecker(st, c, apiBase, secret)
	mc := metrics.NewCollector(st, c, apiBase, secret)
	trafficCollector := trafficlog.NewCollector(st, 5000)
	trafficCollector.Start(ctx, mihomoapi.New(apiBase, secret))

	settings, _ := st.GetSettings(ctx)
	applyGeoSettingsOnStart(*dataDir, settings)
	if settings["geo_auto_download"] != "false" {
		go func() {
			log.Printf("checking geo data in background...")
			cfg := geodata.ConfigFromSettings(settings)
			if updated, err := geodata.Ensure(*dataDir, cfg, false); err != nil {
				log.Printf("geo data: %v", err)
				logCollector.Add("warn", "prism", "geo data: "+err.Error())
			} else if len(updated) > 0 {
				log.Printf("geo data downloaded: %v", updated)
			}
		}()
	}

	logCollector.Add("info", "prism", "Prism started")
	log.Printf("loading proxy engine...")

	if err := app.ReloadEngine(ctx); err != nil {
		log.Printf("initial engine reload: %v (will retry after config)", err)
	} else {
		log.Printf("proxy engine ready")
	}
	go func() {
		warmCtx, cancel := context.WithTimeout(context.Background(), 8*time.Second)
		defer cancel()
		_, _ = mc.Refresh(warmCtx)
		app.RefreshEngineStatusAsync()
	}()

	srv := api.NewServer(app, hc, mc, subSvc, logCollector, trafficCollector)
	router := srv.Router()

	apiPort := settings["api_port"]
	if apiPort == "" {
		apiPort = "8080"
	}

	httpServer := &http.Server{
		Addr:    ":" + apiPort,
		Handler: router,
	}

	cr := cron.New()
	_, _ = cr.AddFunc("@every 1m", func() {
	_, _ = mc.Refresh(context.Background())
	})
	_, _ = cr.AddFunc("@every 5m", func() {
		settings, _ := st.GetSettings(context.Background())
		sec, _ := strconv.Atoi(settings["health_interval_sec"])
		if sec <= 0 {
			sec = 300
		}
		// health checker runs on its own schedule approximation
		_ = hc.TestAll(context.Background(), nil)
	})
	_, _ = cr.AddFunc("@every 10m", func() {
		syncSubscriptions(app, subSvc)
	})
	cr.Start()

	go func() {
		log.Printf("Prism API listening on :%s", apiPort)
		if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			log.Fatal(err)
		}
	}()

	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
	<-sig

	shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_ = httpServer.Shutdown(shutdownCtx)
	cr.Stop()
	engine.Shutdown()
}

func syncSubscriptions(app *service.App, subSvc *subscription.Service) {
	ctx := context.Background()
	settings, err := app.Store.GetSettings(ctx)
	if err != nil || settings["sub_sync_enabled"] != "true" {
		return
	}
	subs, err := app.Store.ListSubscriptions(ctx)
	if err != nil {
		return
	}
	now := time.Now()
	for _, sub := range subs {
		if !sub.Enabled {
			continue
		}
		if sub.LastFetchAt != nil && now.Sub(*sub.LastFetchAt) < time.Duration(sub.IntervalSec)*time.Second {
			continue
		}
		if err := subSvc.FetchAndStore(ctx, sub.ID); err != nil {
			app.Logs.Add("error", "subscription", fmt.Sprintf("sync %s failed: %v", sub.Name, err))
			continue
		}
	}
	_ = app.ReloadEngine(ctx)
}