PixivFE/main.go

package main

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"runtime"
	"sync"
	"syscall"
	"time"

	"codeberg.org/vnpower/pixivfe/v2/audit"
	"codeberg.org/vnpower/pixivfe/v2/config"
	"codeberg.org/vnpower/pixivfe/v2/core"
	"codeberg.org/vnpower/pixivfe/v2/i18n"
	"codeberg.org/vnpower/pixivfe/v2/server/middleware"
	"codeberg.org/vnpower/pixivfe/v2/server/proxy_checker"
	"codeberg.org/vnpower/pixivfe/v2/server/router"
	"codeberg.org/vnpower/pixivfe/v2/server/template"
)

func main() {
	if err := config.GlobalConfig.LoadConfig(); err != nil {
		panic(fmt.Errorf("failed to load configuration: %w", err))
	}

	auditor, err := audit.Init(config.GlobalConfig.InDevelopment)
	if err != nil {
		panic(fmt.Errorf("failed to initialize audit logger: %v", err))
	}
	auditor.SugaredLogger.Info("Audit logger initialized")

	i18n.Init()
	template.Init(config.GlobalConfig.InDevelopment, "assets/views")

	// Conditionally initialize and start the proxy checker
	if config.GlobalConfig.ProxyCheckEnabled {
		go func() {
			var wgFirstCheck sync.WaitGroup
			wgFirstCheck.Add(1)

			startTime := time.Now() // Record the start time for logging

			go func() {
				ctx, cancel := context.WithTimeout(context.Background(), config.GlobalConfig.ProxyCheckTimeout)
				proxy_checker.CheckProxies(ctx, auditor)
				cancel()
				wgFirstCheck.Done()
				if config.GlobalConfig.ProxyCheckInterval != 0 {
					for {
						time.Sleep(config.GlobalConfig.ProxyCheckInterval)
						ctx, cancel := context.WithTimeout(context.Background(), config.GlobalConfig.ProxyCheckTimeout)
						proxy_checker.CheckProxies(ctx, auditor)
						cancel()
					}
				}
			}()

			// Wait for the first proxy check to complete
			auditor.SugaredLogger.Info("Waiting for initial image proxy check to complete...")
			wgFirstCheck.Wait()

			d := time.Since(startTime)
			truncated := d.Truncate(time.Millisecond)

			auditor.SugaredLogger.Infof("Initial image proxy check completed in %v.", truncated)
		}()
	} else {
		auditor.SugaredLogger.Infof("ProxyCheckEnabled set to %t, skipping image proxy checker initialization.", config.GlobalConfig.ProxyCheckEnabled)
	}

	// Initialize cache
	core.InitCache(auditor)
	auditor.SugaredLogger.Info("API response cache initialized.")

	auditor.SugaredLogger.Info("Starting server...")

	router := router.DefineRoutes(auditor)
	// the first middleware is the most outer / first executed one
	router.Use(middleware.ProvideUserContext)  // needed for everything else
	router.Use(middleware.SetLocaleFromCookie) // needed for i18n.*()
	router.Use(middleware.LogRequest(auditor)) // all pages need this
	router.Use(middleware.SetPrivacyHeaders)   // all pages need this
	router.Use(middleware.SetRedirectHeader)   // needed to prevent a cached page from being served upon a user action
	router.Use(middleware.HandleError)         // if the inner handler fails, this shows the error page instead
	router.Use(middleware.InitializeRateLimiter())

	// watch and compile sass when in development mode
	if config.GlobalConfig.InDevelopment {
		go run_sass(auditor)
	}

	// Create http.Server instance
	server := &http.Server{
		Handler: router,
	}

	// Listen
	listener := chooseListener(auditor)

	// Start main server
	go func() {
		if err := server.Serve(listener); err != nil && err != http.ErrServerClosed {
			auditor.SugaredLogger.Fatalf("Starting http.Server failed: %v", err)
		}
	}()

	// Set up graceful shutdown
	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit

	auditor.SugaredLogger.Info("Server is shutting down...")
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	if err := server.Shutdown(ctx); err != nil {
		auditor.SugaredLogger.Fatalf("Server forced to shutdown: %v", err)
	}

	auditor.SugaredLogger.Info("Server exiting")
}

func run_sass(auditor *audit.Auditor) {
	cmd := exec.Command("sass", "--watch", "assets/css")
	cmd.Stdout = os.Stderr // Sass quirk
	cmd.Stderr = os.Stderr
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true, Pdeathsig: syscall.SIGHUP}
	runtime.LockOSThread() // Go quirk https://github.com/golang/go/issues/27505
	err := cmd.Run()
	if err != nil {
		auditor.SugaredLogger.Errorf("When running sass: %v", err)
	}
}

func chooseListener(auditor *audit.Auditor) net.Listener {
	var listener net.Listener

	// Check if we should use a Unix domain socket
	if config.GlobalConfig.UnixSocket != "" {
		unixAddr := config.GlobalConfig.UnixSocket
		unixListener, err := net.Listen("unix", unixAddr)
		if err != nil {
			// Panic with the error description if unable to listen on Unix socket
			panic(fmt.Errorf("failed to listen on Unix socket %q: %w", unixAddr, err))
		}

		// Assign the listener and log where we are listening
		listener = unixListener
		auditor.SugaredLogger.Infof("Listening on Unix domain socket: %v", unixAddr)

	} else {
		// Otherwise, fall back to TCP listener
		addr := net.JoinHostPort(config.GlobalConfig.Host, config.GlobalConfig.Port)
		tcpListener, err := net.Listen("tcp", addr)
		if err != nil {
			// Panic with the error description if unable to listen on TCP
			auditor.SugaredLogger.Panicf("Failed to start TCP listener on %v: %v", addr, err)
		}

		// Assign the TCP listener
		listener = tcpListener
		addr = tcpListener.Addr().String()

		// Extract the host and port for logging
		_, port, err := net.SplitHostPort(addr)
		if err != nil {
			// Panic in case of invalid split into host and port
			panic(fmt.Errorf("error parsing listener address %q: %w", addr, err))
		}

		// Log the address and convenient URL for local development
		auditor.SugaredLogger.Infof("Listening on %v. Accessible at: http://pixivfe.localhost:%v/", addr, port)
	}

	return listener
}