repos / pico

main.go

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