main pico / cmd / pgs / cdn / main.go
Antonio Mika  ·  2026-10-25
  1package main
  2
  3import (
  4	"context"
  5	"fmt"
  6	"io"
  7	"log/slog"
  8	"net"
  9	"net/http"
 10	"net/url"
 11	"strings"
 12
 13	"github.com/picosh/pico/pkg/apps/pgs"
 14	"github.com/picosh/pico/pkg/shared"
 15	"github.com/prometheus/client_golang/prometheus/promhttp"
 16)
 17
 18func main() {
 19	pipeEnabled := shared.GetEnv("PICO_PIPE_ENABLED", "true")
 20	withPipe := strings.ToLower(pipeEnabled) == "true"
 21	logger := shared.CreateLogger("pgs-cdn", withPipe)
 22	ctx := context.Background()
 23	drain := pgs.CreateSubCacheDrain(ctx, logger)
 24	pubsub := pgs.NewPubsubPipe(drain)
 25	defer func() {
 26		_ = pubsub.Close()
 27	}()
 28	cfg := pgs.NewConfig(logger, nil, nil, drain)
 29	proxy := newProxyServe(cfg.Logger)
 30	httpCache, cache := pgs.NewPgsHttpCache(cfg, proxy)
 31	// The CDN can't look users up, but origin fails every site route for a
 32	// site that doesn't exist, so a success proves the owner is real.
 33	cacher := &cachedHttp{
 34		Logger: cfg.Logger,
 35		Cache:  pgs.NewSiteMetrics(cfg, "pgs-cdn", true).Handler(httpCache),
 36	}
 37
 38	go pgs.CacheMgmt(ctx, cfg, cache, nil)
 39
 40	portStr := fmt.Sprintf(":%s", cfg.WebPort)
 41	cfg.Logger.Info(
 42		"starting server on port",
 43		"port", cfg.WebPort,
 44		"domain", cfg.Domain,
 45	)
 46	err := http.ListenAndServe(portStr, cacher)
 47	cfg.Logger.Error("listen and serve", "err", err)
 48}
 49
 50type proxyServe struct {
 51	Logger    *slog.Logger
 52	transport *http.Transport
 53}
 54
 55func newProxyServe(logger *slog.Logger) *proxyServe {
 56	defaultTransport := http.DefaultTransport.(*http.Transport)
 57	oldDial := defaultTransport.DialContext
 58	transport := &http.Transport{
 59		DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
 60			return oldDial(ctx, "tcp", "ash.pgs.sh:443")
 61		},
 62	}
 63	return &proxyServe{Logger: logger, transport: transport}
 64}
 65
 66// Headers that should be stripped from upstream responses because the CDN's
 67// cache layer will add its own versions.
 68var stripHeaders = map[string]bool{
 69	"age":          true,
 70	"cache-status": true,
 71	"date":         true,
 72}
 73
 74func (p *proxyServe) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 75	// Dial ash.pgs.sh but keep the original Host header so ash.pgs.sh
 76	// can route the request to the correct subdomain (zmx.sh, etc.).
 77	// Without the Host header, ash.pgs.sh serves its default vhost (HTML).
 78	target, _ := url.Parse("https://ash.pgs.sh" + req.URL.Path)
 79	if req.URL.RawQuery != "" {
 80		target.RawQuery = req.URL.RawQuery
 81	}
 82	p.Logger.Info("proxying request to ash.pgs.sh", "url", target.String())
 83
 84	proxyReq := req.Clone(req.Context())
 85	proxyReq.URL.Scheme = target.Scheme
 86	proxyReq.URL.Host = target.Host
 87	proxyReq.URL.Path = target.Path
 88	proxyReq.URL.RawQuery = target.RawQuery
 89	proxyReq.RequestURI = ""
 90	// Prevent the upstream from returning a compressed body. The CDN cache
 91	// stores a single representation per URL; Caddy handles per-client
 92	// encoding on the way out. If we forward Accept-Encoding, origin may
 93	// return zstd-compressed bytes that get cached and then served to
 94	// clients that never requested zstd.
 95	proxyReq.Header.Del("Accept-Encoding")
 96	// Preserve the original Host header so ash.pgs.sh routes correctly.
 97	proxyReq.Host = req.Host
 98
 99	resp, err := p.transport.RoundTrip(proxyReq)
100	if err != nil {
101		http.Error(w, err.Error(), http.StatusBadGateway)
102		return
103	}
104	defer func() {
105		_ = resp.Body.Close()
106	}()
107
108	// Copy headers from upstream, but strip cache-related headers that the
109	// CDN's cache layer will regenerate
110	for k, vals := range resp.Header {
111		if stripHeaders[strings.ToLower(k)] {
112			continue
113		}
114		w.Header()[k] = vals
115	}
116	w.WriteHeader(resp.StatusCode)
117	_, _ = io.Copy(w, resp.Body)
118}
119
120type cachedHttp struct {
121	Logger *slog.Logger
122	Cache  http.Handler
123}
124
125func (c *cachedHttp) ServeHTTP(writer http.ResponseWriter, req *http.Request) {
126	if req.URL.Path == "/_metrics" {
127		promhttp.Handler().ServeHTTP(writer, req)
128		return
129	}
130
131	if req.URL.Path == "/check" {
132		c.Logger.Info("proxying `/check` request to ash.pgs.sh", "query", req.URL.RawQuery)
133		req, _ := http.NewRequest("GET", "https://ash.pgs.sh/check?"+req.URL.RawQuery, nil)
134		req.Host = "pgs.sh"
135		// reqDump, _ := httputil.DumpRequestOut(req, true)
136		// fmt.Printf("REQUEST:\n%s", string(reqDump))
137
138		resp, err := http.DefaultClient.Do(req)
139		if err != nil {
140			c.Logger.Error("check request", "err", err)
141		}
142		defer func() {
143			_ = resp.Body.Close()
144		}()
145		writer.WriteHeader(resp.StatusCode)
146		_, _ = io.Copy(writer, resp.Body)
147		return
148	}
149
150	c.Cache.ServeHTTP(writer, req)
151}