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