sub-store/cmd/server.go

package cmd

import (
	"context"
	"errors"
	"fmt"
	"os"
	"os/signal"
	"path/filepath"
	"syscall"
	"time"

	"github.com/gofiber/fiber/v3"
	"github.com/jmoiron/sqlx"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"

	"github.com/peterqiu0516/sub-store/internal/config"
	"github.com/peterqiu0516/sub-store/internal/database"
	"github.com/peterqiu0516/sub-store/internal/geoip"
	"github.com/peterqiu0516/sub-store/internal/handler"
	"github.com/peterqiu0516/sub-store/internal/middleware"
)

var configFile string

var serveCmd = &cobra.Command{
	Use:   "serve",
	Short: "Start the Sub-Store HTTP server",
	RunE: func(cmd *cobra.Command, args []string) error {
		cfg, err := config.Load(configFile)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		return RunServer(cfg)
	},
}

func init() {
	rootCmd.AddCommand(serveCmd)
	serveCmd.Flags().StringVarP(&configFile, "config", "c", "", "config file path")
}

// RunServer opens the database, applies migrations (idempotent — goose tracks
// applied versions in goose_db_version), then starts the HTTP server.
func RunServer(cfg *config.Config) error {
	db, err := database.InitDB(cfg.Database.Path)
	if err != nil {
		return err
	}
	defer db.Close()

	// Auto-migrate on every startup; goose is idempotent so this is safe to call repeatedly.
	if err := database.RunMigrations(db); err != nil {
		return err
	}

	// Ensure GeoIP database exists (download in background if missing)
	mmdbPath := filepath.Join(filepath.Dir(cfg.Database.Path), "GeoLite2-City.mmdb")
	geoip.SetMmdbPath(mmdbPath)
	if _, err := os.Stat(mmdbPath); err != nil {
		logrus.Info("GeoIP database not found, downloading in background...")
		go func() {
			if err := geoip.DownloadMmdb(mmdbPath); err != nil {
				logrus.WithError(err).Warn("GeoIP download failed, name-based detection will be used as fallback")
			} else {
				logrus.Info("GeoIP database downloaded successfully")
				// Reload the reader
				geoip.SetMmdbPath(mmdbPath)
			}
		}()
	} else {
		logrus.Info("GeoIP database found")
	}

	return startFiber(cfg, db)
}

func startFiber(cfg *config.Config, db *sqlx.DB) error {
	logrus.SetFormatter(&logrus.TextFormatter{
		FullTimestamp: true,
	})
	logrus.SetLevel(logrus.InfoLevel)

	// Start cache cleaner — review-resolution #32
	cacheCtx, cacheCancel := context.WithCancel(context.Background())
	defer cacheCancel()
	cacheInterval := cfg.Fetcher.CacheTTL
	if cacheInterval <= 0 {
		cacheInterval = 300 * time.Second
	}
	database.StartCacheCleaner(cacheCtx, db, cacheInterval)

	app := fiber.New(fiber.Config{
		ReadTimeout:  cfg.Server.ReadTimeout,
		WriteTimeout: cfg.Server.WriteTimeout,
		BodyLimit:    cfg.Server.BodyLimit,
		ErrorHandler: func(c fiber.Ctx, err error) error {
			code := fiber.StatusInternalServerError
			message := "Internal server error"
			var fiberErr *fiber.Error
			if errors.As(err, &fiberErr) {
				code = fiberErr.Code
				message = fiberErr.Message
			}
			logrus.WithError(err).WithFields(logrus.Fields{
				"method": c.Method(),
				"path":   c.Path(),
			}).Error("request error")
			return c.Status(code).JSON(fiber.Map{
				"status": "failed",
				"error":  fiber.Map{"code": code, "message": message},
			})
		},
	})

	// Global middleware
	app.Use(middleware.SecurityHeaders())
	app.Use(middleware.CORS("*"))
	app.Use(middleware.HandleOptions())
	app.Use(middleware.DownloadHostIsolation(cfg.Auth.DownloadHosts))

	// Register routes
	handler.RegisterRoutes(app, cfg, db)

	// Health check
	app.Get("/health", func(c fiber.Ctx) error {
		return c.JSON(fiber.Map{"status": "ok"})
	})

	// Start server
	addr := fmt.Sprintf("%s:%d", cfg.Server.Host, cfg.Server.Port)
	logrus.Infof("Starting %s on %s", cfg.App.Name, addr)

	go func() {
		if err := app.Listen(addr); err != nil {
			logrus.WithError(err).Error("server failed")
		}
	}()

	// Graceful shutdown — review-resolution #34
	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit

	logrus.Info("Shutting down server...")
	cacheCancel()

	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer shutdownCancel()
	go func() {
		<-shutdownCtx.Done()
	}()

	if err := app.Shutdown(); err != nil {
		logrus.WithError(err).Error("shutdown error")
	}
	logrus.Info("Server stopped")
	return nil
}