dns-autoresolver/cmd/server/main.go

package main

import (
	"context"
	"errors"
	"log"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"

	"github.com/vasyakrg/dns-autoresolver/internal/api"
	"github.com/vasyakrg/dns-autoresolver/internal/auth"
	"github.com/vasyakrg/dns-autoresolver/internal/config"
	"github.com/vasyakrg/dns-autoresolver/internal/crypto"
	"github.com/vasyakrg/dns-autoresolver/internal/metrics"
	"github.com/vasyakrg/dns-autoresolver/internal/notify"
	"github.com/vasyakrg/dns-autoresolver/internal/provider/registry"
	"github.com/vasyakrg/dns-autoresolver/internal/provider/selectel"
	"github.com/vasyakrg/dns-autoresolver/internal/scheduler"
	"github.com/vasyakrg/dns-autoresolver/internal/service"
	"github.com/vasyakrg/dns-autoresolver/internal/store"
	"github.com/vasyakrg/dns-autoresolver/internal/web"
)

// sessionTTL is how long a login session cookie remains valid before the
// user must re-authenticate.
const sessionTTL = 720 * time.Hour

// schedulerTick is how often the in-process scheduler checks for due project
// schedules. Individual projects only actually run when their own
// schedules.interval_seconds has elapsed (see internal/store ListDueSchedules) —
// this is just the polling granularity.
const schedulerTick = time.Minute

// shutdownTimeout bounds how long graceful shutdown waits for in-flight HTTP
// requests to finish before forcing the listener closed.
const shutdownTimeout = 10 * time.Second

// isAPIPath reports whether path must be routed to the API router rather
// than the SPA. "/api" (no trailing slash) counts as an API path too —
// only strings.HasPrefix(path, "/api/") would otherwise miss it and fall
// through to the SPA fallback.
func isAPIPath(path string) bool {
	return path == "/api" || strings.HasPrefix(path, "/api/")
}

// buildMux wires the public /healthz + /metrics endpoints, the API router,
// and the embedded SPA. /healthz and /metrics are intentionally auth-free —
// /healthz is a liveness probe (always 200 while the process serves), and
// metricsHandler only ever exposes aggregate counters/gauges.
func buildMux(metricsHandler http.Handler, apiRouter http.Handler, webHandler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch {
		case r.URL.Path == "/healthz":
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte("ok"))
		case r.URL.Path == "/metrics":
			metricsHandler.ServeHTTP(w, r)
		case isAPIPath(r.URL.Path):
			apiRouter.ServeHTTP(w, r)
		case webHandler != nil:
			webHandler.ServeHTTP(w, r)
		default:
			http.NotFound(w, r)
		}
	})
}

// healthcheck performs an in-process liveness probe used as the container
// HEALTHCHECK — distroless images have no curl/wget. It GETs /healthz on the
// configured listen address and maps 200 -> 0, anything else -> 1.
func healthcheck() int {
	addr := os.Getenv("DNS_AR_LISTEN")
	if addr == "" {
		addr = ":8080"
	}
	// ":8080" -> "127.0.0.1:8080"
	if strings.HasPrefix(addr, ":") {
		addr = "127.0.0.1" + addr
	}
	c := &http.Client{Timeout: 3 * time.Second}
	resp, err := c.Get("http://" + addr + "/healthz")
	if err != nil {
		return 1
	}
	defer resp.Body.Close()
	if resp.StatusCode == http.StatusOK {
		return 0
	}
	return 1
}

func main() {
	if len(os.Args) > 1 && os.Args[1] == "-healthcheck" {
		os.Exit(healthcheck())
	}

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	cfg, err := config.Load()
	if err != nil {
		log.Fatalf("config: %v", err)
	}
	if err := store.Migrate(ctx, cfg.DBDSN); err != nil {
		log.Fatalf("migrate: %v", err)
	}
	pool, err := pgxpool.New(ctx, cfg.DBDSN)
	if err != nil {
		log.Fatalf("pool: %v", err)
	}
	defer pool.Close()

	cipher, err := crypto.NewCipher(cfg.EncKey)
	if err != nil {
		log.Fatalf("cipher: %v", err)
	}
	st := store.New(pool)
	sessions := auth.NewSessions(st, sessionTTL)

	reg := registry.New()
	reg.Register(selectel.New())

	svc := service.New(st, st, reg, cipher)

	m := metrics.New()
	dispatcher := notify.NewDispatcher(st, cipher)

	a := &api.API{
		Svc: svc, Store: st, Cipher: cipher, Reg: reg, Auth: st, Sessions: sessions,
		Schedule: st, Dispatch: dispatcher,
	}
	apiRouter := api.NewRouter(a)

	webHandler, err := web.Handler()
	if err != nil {
		log.Printf("web: static UI unavailable: %v", err)
	}

	// The scheduler only checks and notifies — it never applies zone changes
	// (Apply stays a manual, explicit API call). Its own errors are logged
	// internally and never stop the loop; ctx cancellation (signal) is the
	// only thing that ends Run.
	sched := scheduler.New(st, svc, dispatcher, m)
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		sched.Run(ctx, schedulerTick)
	}()

	mux := buildMux(m.Handler(), apiRouter, webHandler)

	srv := &http.Server{Addr: cfg.ListenAddr, Handler: mux}

	serveErr := make(chan error, 1)
	go func() {
		log.Printf("listening on %s", cfg.ListenAddr)
		if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
			serveErr <- err
			return
		}
		serveErr <- nil
	}()

	select {
	case <-ctx.Done():
		log.Printf("shutdown signal received, draining connections (timeout %s)", shutdownTimeout)
		shutdownCtx, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
		defer cancel()
		if err := srv.Shutdown(shutdownCtx); err != nil {
			log.Printf("server: graceful shutdown failed: %v", err)
		}
		<-serveErr
		// Wait for the in-flight scheduler RunOnce (interrupted by the
		// cancelled ctx passed into checker.Check) to finish before exiting,
		// so we never kill the process mid-write of a check/notify status.
		wg.Wait()
		log.Printf("server stopped")
	case err := <-serveErr:
		if err != nil {
			// Unwind scheduler.Run via ctx.Done and wait for the in-flight
			// RunOnce to finish before exiting, so an unexpected serve error
			// mid-run doesn't kill the process during a check/notify write —
			// same hazard the ctx.Done branch above already guards against.
			stop()
			wg.Wait()
			log.Fatalf("server: %v", err)
		}
	}
}