teldrive/cmd/run.go

package cmd

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"regexp"
	"strings"
	"time"

	"github.com/go-chi/chi/v5"
	chimiddleware "github.com/go-chi/chi/v5/middleware"
	"github.com/go-chi/cors"
	"github.com/go-co-op/gocron"
	"github.com/spf13/cobra"
	"github.com/tgdrive/teldrive/internal/api"
	"github.com/tgdrive/teldrive/internal/appcontext"
	"github.com/tgdrive/teldrive/internal/auth"
	"github.com/tgdrive/teldrive/internal/cache"
	"github.com/tgdrive/teldrive/internal/chizap"
	"github.com/tgdrive/teldrive/internal/config"
	"github.com/tgdrive/teldrive/internal/database"
	"github.com/tgdrive/teldrive/internal/duration"
	"github.com/tgdrive/teldrive/internal/logging"
	"github.com/tgdrive/teldrive/internal/middleware"
	"github.com/tgdrive/teldrive/internal/tgc"
	"github.com/tgdrive/teldrive/internal/tgstorage"
	"github.com/tgdrive/teldrive/ui"

	"github.com/tgdrive/teldrive/pkg/cron"
	"github.com/tgdrive/teldrive/pkg/services"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"gorm.io/gorm"
)

func NewRun() *cobra.Command {
	var cfg config.ServerCmdConfig
	cmd := &cobra.Command{
		Use:   "run",
		Short: "Start Teldrive Server",
		Run: func(cmd *cobra.Command, args []string) {
			runApplication(cmd.Context(), &cfg)

		},
		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
			loader := config.NewConfigLoader()
			if err := loader.InitializeConfig(cmd); err != nil {
				return err
			}
			if err := loader.Load(&cfg); err != nil {
				return err
			}
			if err := checkRequiredRunFlags(&cfg); err != nil {
				return err
			}
			return nil
		},
	}
	addServerFlags(cmd, &cfg)
	return cmd
}
func addServerFlags(cmd *cobra.Command, cfg *config.ServerCmdConfig) {

	flags := cmd.Flags()

	config.AddCommonFlags(flags, cfg)

	// Server config
	flags.IntVarP(&cfg.Server.Port, "server-port", "p", 8080, "Server port")
	duration.DurationVar(flags, &cfg.Server.GracefulShutdown, "server-graceful-shutdown", 10*time.Second, "Server graceful shutdown timeout")
	flags.BoolVar(&cfg.Server.EnablePprof, "server-enable-pprof", false, "Enable Pprof Profiling")
	duration.DurationVar(flags, &cfg.Server.ReadTimeout, "server-read-timeout", 1*time.Hour, "Server read timeout")
	duration.DurationVar(flags, &cfg.Server.WriteTimeout, "server-write-timeout", 1*time.Hour, "Server write timeout")

	// CronJobs config
	flags.BoolVar(&cfg.CronJobs.Enable, "cronjobs-enable", true, "Run cron jobs")
	duration.DurationVar(flags, &cfg.CronJobs.CleanFilesInterval, "cronjobs-clean-files-interval", 1*time.Hour, "Clean files interval")
	duration.DurationVar(flags, &cfg.CronJobs.CleanUploadsInterval, "cronjobs-clean-uploads-interval", 12*time.Hour, "Clean uploads interval")
	duration.DurationVar(flags, &cfg.CronJobs.FolderSizeInterval, "cronjobs-folder-size-interval", 2*time.Hour, "Folder size update  interval")

	// Cache config
	flags.IntVar(&cfg.Cache.MaxSize, "cache-max-size", 10*1024*1024, "Max Cache max size (memory)")
	flags.StringVar(&cfg.Cache.RedisAddr, "cache-redis-addr", "", "Redis address")
	flags.StringVar(&cfg.Cache.RedisPass, "cache-redis-pass", "", "Redis password")

	// JWT config
	flags.StringVar(&cfg.JWT.Secret, "jwt-secret", "", "JWT secret key")
	duration.DurationVar(flags, &cfg.JWT.SessionTime, "jwt-session-time", (30*24)*time.Hour, "JWT session duration")
	flags.StringSliceVar(&cfg.JWT.AllowedUsers, "jwt-allowed-users", []string{}, "Allowed users")

	// Telegram config
	flags.StringVar(&cfg.TG.StorageFile, "tg-storage-file", "", "Sqlite Storage file path")
	flags.BoolVar(&cfg.TG.RateLimit, "tg-rate-limit", true, "Enable rate limiting for telegram client")
	flags.IntVar(&cfg.TG.RateBurst, "tg-rate-burst", 5, "Limiting burst for telegram client")
	flags.IntVar(&cfg.TG.Rate, "tg-rate", 100, "Limiting rate for telegram client")
	flags.StringVar(&cfg.TG.Proxy, "tg-proxy", "", "HTTP OR SOCKS5 proxy URL")
	flags.BoolVar(&cfg.TG.DisableStreamBots, "tg-disable-stream-bots", false, "Disable Stream bots")
	flags.BoolVar(&cfg.TG.Ntp, "tg-ntp", false, "Use NTP server time")
	flags.BoolVar(&cfg.TG.EnableLogging, "tg-enable-logging", false, "Enable telegram client logging")
	flags.Int64Var(&cfg.TG.PoolSize, "tg-pool-size", 8, "Telegram Session pool size")

	// Telegram Uploads config
	flags.StringVar(&cfg.TG.Uploads.EncryptionKey, "tg-uploads-encryption-key", "", "Uploads encryption key")
	flags.IntVar(&cfg.TG.Uploads.Threads, "tg-uploads-threads", 8, "Uploads threads")
	flags.IntVar(&cfg.TG.Uploads.MaxRetries, "tg-uploads-max-retries", 10, "Uploads Retries")
	duration.DurationVar(flags, &cfg.TG.ReconnectTimeout, "tg-reconnect-timeout", 5*time.Minute, "Reconnect Timeout")
	duration.DurationVar(flags, &cfg.TG.Uploads.Retention, "tg-uploads-retention", (24*7)*time.Hour, "Uploads retention duration")
	flags.IntVar(&cfg.TG.Stream.MultiThreads, "tg-stream-multi-threads", 0, "Stream multi-threads")
	flags.IntVar(&cfg.TG.Stream.Buffers, "tg-stream-buffers", 8, "No of Stream buffers")
	duration.DurationVar(flags, &cfg.TG.Stream.ChunkTimeout, "tg-stream-chunk-timeout", 20*time.Second, "Chunk Fetch Timeout")

}

func checkRequiredRunFlags(cfg *config.ServerCmdConfig) error {
	var missingFields []string

	if cfg.DB.DataSource == "" {
		missingFields = append(missingFields, "db-data-source")
	}
	if cfg.JWT.Secret == "" {
		missingFields = append(missingFields, "jwt-secret")
	}
	if cfg.TG.AppHash == "" {
		missingFields = append(missingFields, "tg-app-hash")
	}
	if cfg.TG.AppId == 0 {
		missingFields = append(missingFields, "tg-app-id")
	}

	if len(missingFields) > 0 {
		return fmt.Errorf("required configuration values not set: %s", strings.Join(missingFields, ", "))
	}

	return nil
}

func findAvailablePort(startPort int) (int, error) {
	for port := startPort; port < startPort+100; port++ {
		addr := fmt.Sprintf(":%d", port)
		listener, err := net.Listen("tcp", addr)
		if err != nil {
			continue
		}
		listener.Close()
		return port, nil
	}
	return 0, fmt.Errorf("no available ports found between %d and %d", startPort, startPort+100)
}

func runApplication(ctx context.Context, conf *config.ServerCmdConfig) {
	lvl, err := zapcore.ParseLevel(conf.Log.Level)
	if err != nil {
		lvl = zapcore.InfoLevel
	}
	logging.SetConfig(&logging.Config{
		Level:    lvl,
		FilePath: conf.Log.File,
	})

	lg := logging.DefaultLogger().Sugar()

	defer lg.Sync()

	port, err := findAvailablePort(conf.Server.Port)
	if err != nil {
		lg.Fatalw("failed to find available port", "err", err)
	}
	if port != conf.Server.Port {
		lg.Infof("Port %d is occupied, using port %d instead", conf.Server.Port, port)
		conf.Server.Port = port
	}

	scheduler := gocron.NewScheduler(time.UTC)

	cacher := cache.NewCache(ctx, &conf.Cache)

	db, err := database.NewDatabase(&conf.DB, lg)

	if err != nil {
		lg.Fatalw("failed to create database", "err", err)
	}

	err = database.MigrateDB(db)

	if err != nil {
		lg.Fatalw("failed to migrate database", "err", err)
	}

	tgdb, err := tgstorage.NewDatabase(conf.TG.StorageFile)
	if err != nil {
		lg.Fatalw("failed to create tg db", "err", err)
	}

	err = tgstorage.MigrateDB(tgdb)
	if err != nil {
		lg.Fatalw("failed to migrate tg db", "err", err)
	}

	worker := tgc.NewBotWorker()

	srv := setupServer(conf, db, cacher, tgdb, worker)

	cron.StartCronJobs(scheduler, db, conf)

	go func() {
		lg.Infof("Server started at http://localhost:%d", conf.Server.Port)
		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			lg.Errorw("failed to start server", "err", err)
		}
	}()

	<-ctx.Done()

	lg.Info("Shutting down server...")

	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), conf.Server.GracefulShutdown)

	defer shutdownCancel()

	if err := srv.Shutdown(shutdownCtx); err != nil {
		lg.Errorw("server shutdown failed", "err", err)
	}

	scheduler.Stop()

	lg.Info("Server stopped")
}

func setupServer(cfg *config.ServerCmdConfig, db *gorm.DB, cache cache.Cacher, tgdb *gorm.DB, worker *tgc.BotWorker) *http.Server {

	lg := logging.DefaultLogger()

	apiSrv := services.NewApiService(db, cfg, cache, tgdb, worker)

	srv, err := api.NewServer(apiSrv, auth.NewSecurityHandler(db, cache, &cfg.JWT))

	if err != nil {
		lg.Fatal("failed to create server", zap.Error(err))
	}

	extendedSrv := services.NewExtendedMiddleware(srv, services.NewExtendedService(apiSrv))

	mux := chi.NewRouter()

	mux.Use(chimiddleware.Recoverer)
	mux.Use(cors.Handler(cors.Options{
		AllowedOrigins: []string{"*"},
		AllowedMethods: []string{"GET", "POST", "PUT", "DELETE", "OPTIONS", "PATCH", "HEAD"},
		AllowedHeaders: []string{"Accept", "Authorization", "Content-Type"},
		MaxAge:         86400,
	}))
	mux.Use(chimiddleware.RealIP)
	mux.Use(middleware.InjectLogger(lg))
	mux.Use(chizap.ChizapWithConfig(lg, &chizap.Config{
		TimeFormat: time.RFC3339,
		UTC:        true,
		SkipPathRegexps: []*regexp.Regexp{
			regexp.MustCompile(`^/(assets|images|docs)/.*`),
		},
	}))
	mux.Use(appcontext.Middleware)
	mux.Mount("/api/", http.StripPrefix("/api", extendedSrv))
	mux.Handle("/*", middleware.SPAHandler(ui.StaticFS))

	return &http.Server{
		Addr:              fmt.Sprintf(":%d", cfg.Server.Port),
		Handler:           mux,
		ReadTimeout:       cfg.Server.ReadTimeout,
		WriteTimeout:      cfg.Server.WriteTimeout,
		ReadHeaderTimeout: 10 * time.Second,
		IdleTimeout:       60 * time.Second,
	}
}