Vellum/cmd/vellum/main.go

// Command vellum is the unified media server: one process, one database,
// every media type.
package main

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/mithral/vellum/internal/auth"
	"github.com/mithral/vellum/internal/core"
	"github.com/mithral/vellum/internal/httpapi"
	"github.com/mithral/vellum/internal/jobs"
	"github.com/mithral/vellum/internal/library"
	"github.com/mithral/vellum/internal/media"
	"github.com/mithral/vellum/internal/metadata"
	"github.com/mithral/vellum/internal/modules/comics"
	"github.com/mithral/vellum/internal/modules/ebooks"
	"github.com/mithral/vellum/internal/search"
	"github.com/mithral/vellum/migrations"
)

// version is overridden at build time with -ldflags "-X main.version=...".
var version = "dev"

func main() {
	if err := run(); err != nil {
		fmt.Fprintln(os.Stderr, "vellum: fatal:", err)
		os.Exit(1)
	}
}

func run() error {
	cfg, err := core.LoadConfig()
	if err != nil {
		return err
	}
	log := core.NewLogger(cfg.LogLevel, cfg.LogFormat)
	log.Info("starting vellum", "version", version, "listen", cfg.Listen)

	if err := cfg.Validate(); err != nil {
		return err
	}

	db, err := core.OpenDB(cfg.DBPath)
	if err != nil {
		return err
	}
	defer db.Close()

	migCtx, cancelMig := context.WithTimeout(context.Background(), 60*time.Second)
	err = core.Migrate(migCtx, db, migrations.FS, log)
	cancelMig()
	if err != nil {
		return err
	}

	api := httpapi.New(db, log, version)

	authenticator := auth.NewAuthenticator(db, cfg.CookieName, cfg.SessionTTL, cfg.CookieSecure, log)
	if !cfg.CookieSecure {
		log.Warn("VELLUM_COOKIE_SECURE is false; session cookies will be sent over plain HTTP. " +
			"Set it to true when serving over TLS.")
	}

	if err := auth.EnsureAdminFromEnv(context.Background(), db, log, cfg.AdminUser, cfg.AdminPassword); err != nil {
		return fmt.Errorf("bootstrap admin: %w", err)
	}

	authHandlers := auth.NewHandlers(authenticator, log)
	authHandlers.Register(api.Mux())
	api.SetAuth(authenticator)

	library.NewHandlers(db, log, authenticator).Register(api.Mux())
	media.NewHandlers(db, log, authenticator, cfg.ConfigDir).Register(api.Mux())
	ebooks.NewHandlers(db, log, authenticator).Register(api.Mux())
	comics.NewHandlers(db, log, authenticator).Register(api.Mux())
	search.NewHandlers(db, log, authenticator).Register(api.Mux())

	// Background work: a table-backed queue in the same database. Providers
	// are registered primary-first; a fallback is only consulted when the
	// primary returns nothing.
	queue := jobs.New(db, log)
	metadata.Register(metadata.NewOpenLibrary())
	metadata.Register(metadata.NewGoogleBooks())
	enricher := metadata.NewEnricher(db, log, cfg.ConfigDir)
	queue.Register(metadata.JobKind, enricher.Handle)
	metadata.NewHandlers(db, log, authenticator, queue).Register(api.Mux())

	jobCtx, cancelJobs := context.WithCancel(context.Background())
	defer cancelJobs()
	go queue.Run(jobCtx)

	srv := &http.Server{
		Addr:    cfg.Listen,
		Handler: api.Handler(),
		// ReadHeaderTimeout protects against slowloris. WriteTimeout is
		// deliberately zero: streaming large media files must not be cut off.
		ReadHeaderTimeout: 10 * time.Second,
		ReadTimeout:       cfg.ReadTimeout,
		WriteTimeout:      cfg.WriteTimeout,
		IdleTimeout:       cfg.IdleTimeout,
		ErrorLog:          nil,
	}

	errCh := make(chan error, 1)
	go func() {
		log.Info("listening", "addr", cfg.Listen, "data_dir", cfg.DataDir)
		if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
			errCh <- err
		}
	}()

	stop := make(chan os.Signal, 1)
	signal.Notify(stop, os.Interrupt, syscall.SIGTERM)

	select {
	case err := <-errCh:
		return fmt.Errorf("http server: %w", err)
	case sig := <-stop:
		log.Info("shutting down", "signal", sig.String())
	}

	shutdownCtx, cancel := context.WithTimeout(context.Background(), cfg.ShutdownGrace)
	defer cancel()
	if err := srv.Shutdown(shutdownCtx); err != nil {
		return fmt.Errorf("graceful shutdown: %w", err)
	}
	cancelJobs()
	log.Info("stopped")
	return nil
}