tiny-proxy/main.go

package main

import (
	"context"
	"flag"
	"log/slog"
	"net"
	"net/http"
	"os"
	"time"

	"tiny-proxy/internal/config"
	"tiny-proxy/internal/distro"
	"tiny-proxy/internal/docker"
	"tiny-proxy/internal/github"
	"tiny-proxy/internal/maven"
	"tiny-proxy/internal/metrics"
	"tiny-proxy/internal/proxycache"
	"tiny-proxy/internal/storage"
	"tiny-proxy/internal/tracing"
	"tiny-proxy/internal/ui"
	"tiny-proxy/internal/upstream"
)

func main() {
	configPath := flag.String("config", "config.yaml", "path to config file")
	flag.Parse()

	cfg, err := config.Load(*configPath)
	if err != nil {
		slog.Error("load config", "err", err)
		os.Exit(1)
	}
	// From here on every log record carries the trace id of the request that
	// caused it, so one grep follows a client request through the upstreams
	// it fanned out to.
	slog.SetDefault(tracing.NewLogger(os.Stderr, cfg.Log.Level, cfg.Log.Format))

	ctx := context.Background()
	reg := metrics.NewRegistry()
	var store storage.Store
	switch cfg.Storage.Type {
	case "none":
		// No cache: every request is proxied straight to the upstream.
		store = storage.NullStore{}
	case "s3":
		store, err = storage.NewS3(ctx, cfg.Storage.S3)
	default:
		var local *storage.LocalFS
		local, err = storage.NewLocalFS(cfg.Storage.LocalFS.Root)
		if local != nil {
			local.OnCorrupt = func(string) { reg.IncCorruptEntry() }
			store = local
		}
	}
	if err != nil {
		slog.Error("init storage", "err", err)
		os.Exit(1)
	}

	// One proxy/auth-aware HTTP client per vrepo upstream (maven, debian and
	// pacman vrepos may have several, tried in order).
	clients := map[string][]*http.Client{}
	for _, v := range cfg.VRepos {
		// Registering up front keeps every configured series present from
		// the first scrape — a mirror that has served nothing yet reads as
		// zero rather than as a gap — and lets each client hold its own
		// counters directly instead of looking them up per request.
		reg.RegisterVRepo(v.Name, string(v.Type))
		for _, u := range v.AllUpstreams() {
			c, err := upstream.NewClient(u, upstream.Options{
				VRepo:   v.Name,
				Log:     cfg.Log.LogUpstream(),
				Metrics: reg,
				Stats:   reg.RegisterUpstream(v.Name, u.URL),
			})
			if err != nil {
				slog.Error("init upstream client", "vrepo", v.Name, "url", u.URL, "err", err)
				os.Exit(1)
			}
			clients[v.Name] = append(clients[v.Name], c)
		}
	}
	// Docker, github and gitea vrepos are single-upstream.
	first := map[string]*http.Client{}
	for name, cs := range clients {
		first[name] = cs[0]
	}

	cache := proxycache.New(store)
	cache.Metrics = reg
	mux := http.NewServeMux()
	mux.Handle("/v2/", metrics.CountRequests(reg, "docker", docker.NewHandler(cfg.VRepos, first, cache)))
	mux.Handle("/maven/", metrics.CountRequests(reg, "maven", maven.NewHandler(cfg.VRepos, clients, cache)))
	mux.Handle("/github/", metrics.CountRequests(reg, "github", github.NewHandler(cfg.VRepos, first, cache)))
	mux.Handle("/gitea/", metrics.CountRequests(reg, "gitea", github.NewHandler(cfg.VRepos, first, cache)))
	mux.Handle("/debian/", metrics.CountRequests(reg, "debian", distro.NewHandler(distro.Debian, cfg.VRepos, clients, cache)))
	mux.Handle("/pacman/", metrics.CountRequests(reg, "pacman", distro.NewHandler(distro.Pacman, cfg.VRepos, clients, cache)))
	mux.Handle("/ui/", ui.NewHandler(cfg.VRepos, store, clients).Routes())
	mux.HandleFunc("GET /{$}", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, "/ui/", http.StatusFound)
	})
	mux.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("ok\n"))
	})
	mux.Handle("GET /metrics", metrics.NewHandler(reg, cfg.Metrics.Token))

	srv := &http.Server{
		Addr: cfg.Listen,
		// Wrapping the whole mux, so /ui/, /metrics and /healthz get trace
		// ids and access logs too.
		Handler:           tracing.Middleware(cfg.Log.LogRequests(), reg, mux),
		ReadHeaderTimeout: 10 * time.Second,
	}
	for _, v := range cfg.VRepos {
		urls := make([]string, 0, len(v.AllUpstreams()))
		for _, u := range v.AllUpstreams() {
			urls = append(urls, u.URL)
		}
		slog.Info("vrepo configured", "name", v.Name, "type", v.Type, "upstreams", urls, "auth", v.Auth.Mode)
	}
	// Serving on our own listener rather than ListenAndServe so client
	// connections can be counted at the socket: the client byte totals then
	// include request lines, headers and chunked framing instead of payload
	// alone. Bound before the ready line, so a taken port reports as a
	// failure to start rather than after announcing success.
	ln, err := net.Listen("tcp", cfg.Listen)
	if err != nil {
		slog.Error("listen", "addr", cfg.Listen, "err", err)
		os.Exit(1)
	}
	slog.Info("listening", "addr", cfg.Listen, "storage", cfg.Storage.Type,
		"log_level", cfg.Log.Level, "log_format", cfg.Log.Format)
	if err := srv.Serve(reg.CountConnections(ln)); err != nil {
		slog.Error("server stopped", "err", err)
		os.Exit(1)
	}
}