1package router
  2
  3import (
  4	"context"
  5	"errors"
  6	"fmt"
  7	"net"
  8	"net/http"
  9	"net/http/pprof"
 10	"strings"
 11	"time"
 12
 13	"github.com/hashicorp/golang-lru/v2/expirable"
 14	"github.com/picosh/pico/pkg/shared"
 15)
 16
 17// CreatePProfRoutesMux serves the profiler on mux.
 18func CreatePProfRoutesMux(mux *http.ServeMux) {
 19	for _, method := range []string{"GET", "POST"} {
 20		mux.HandleFunc(method+" /debug/pprof/", pprof.Index)
 21		mux.HandleFunc(method+" /debug/pprof/cmdline", pprof.Cmdline)
 22		mux.HandleFunc(method+" /debug/pprof/profile", pprof.Profile)
 23		mux.HandleFunc(method+" /debug/pprof/symbol", pprof.Symbol)
 24		mux.HandleFunc(method+" /debug/pprof/trace", pprof.Trace)
 25	}
 26}
 27
 28// Cors adds CORS headers to h's responses.
 29func Cors(h http.HandlerFunc) http.HandlerFunc {
 30	return func(w http.ResponseWriter, r *http.Request) {
 31		CorsHeaders(w.Header())
 32		h(w, r)
 33	}
 34}
 35
 36// CorsPreflight answers a CORS preflight request.
 37func CorsPreflight(w http.ResponseWriter, r *http.Request) {
 38	CorsHeaders(w.Header())
 39	w.WriteHeader(http.StatusOK)
 40}
 41
 42func GetSubdomainFromRequest(r *http.Request, domain, space string) string {
 43	hostDomain := strings.ToLower(strings.Split(r.Host, ":")[0])
 44	appDomain := strings.ToLower(strings.Split(domain, ":")[0])
 45
 46	if hostDomain != appDomain && strings.Contains(hostDomain, ".") {
 47		if strings.Contains(hostDomain, appDomain) {
 48			subdomain := strings.TrimSuffix(hostDomain, fmt.Sprintf(".%s", appDomain))
 49			return subdomain
 50		} else {
 51			subdomain := GetCustomDomain(hostDomain, space)
 52			return subdomain
 53		}
 54	}
 55
 56	return ""
 57}
 58
 59// NewHostRouter serves requests for users' sites with sites, and the rest
 60// with main. A site is a subdomain of domain, or a custom domain whose TXT
 61// record for txtPrefix names one. GetSubdomain returns the site's name. A
 62// nil sites serves everything with main.
 63func NewHostRouter(domain, txtPrefix string, main, sites http.Handler) http.Handler {
 64	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 65		if sites == nil {
 66			main.ServeHTTP(w, r)
 67			return
 68		}
 69		subdomain := GetSubdomainFromRequest(r, domain, txtPrefix)
 70		if subdomain == "" {
 71			main.ServeHTTP(w, r)
 72			return
 73		}
 74		sites.ServeHTTP(w, r.WithContext(WithSubdomain(r.Context(), subdomain)))
 75	})
 76}
 77
 78type ctxSubdomainKey struct{}
 79
 80// WithSubdomain marks ctx as a request for the site subdomain.
 81func WithSubdomain(ctx context.Context, subdomain string) context.Context {
 82	return context.WithValue(ctx, ctxSubdomainKey{}, subdomain)
 83}
 84
 85// GetSubdomain is the site a request is for, or "" for the main domain.
 86func GetSubdomain(r *http.Request) string {
 87	subdomain, _ := r.Context().Value(ctxSubdomainKey{}).(string)
 88	return subdomain
 89}
 90
 91var txtCache = expirable.NewLRU[string, string](2048, nil, shared.CacheTimeout)
 92
 93// txtMissCache remembers hosts without a TXT record for a short time, so
 94// unknown hosts don't cost a lookup per request and a record that was just
 95// added starts working quickly.
 96var txtMissCache = expirable.NewLRU[string, struct{}](2048, nil, 30*time.Second)
 97
 98var lookupTXT = net.LookupTXT
 99
100// ErrNoTXTRecord is returned when a host has no TXT record for a service.
101var ErrNoTXTRecord = errors.New("no TXT record")
102
103// TXTRecordName is the record that points host at a site on space.
104func TXTRecordName(host, space string) string {
105	return fmt.Sprintf("_%s.%s", space, host)
106}
107
108// LookupTXTRecords asks DNS for host's records on space, bypassing the
109// cache. No records returns ErrNoTXTRecord.
110func LookupTXTRecords(host, space string) ([]string, error) {
111	records, err := lookupTXT(TXTRecordName(host, space))
112	var dnsErr *net.DNSError
113	if errors.As(err, &dnsErr) && dnsErr.IsNotFound {
114		return nil, ErrNoTXTRecord
115	}
116	if err != nil {
117		return nil, err
118	}
119	if len(records) == 0 {
120		return nil, ErrNoTXTRecord
121	}
122	return records, nil
123}
124
125// LookupCustomDomain returns the site host's first record on space names,
126// bypassing the cache.
127func LookupCustomDomain(host, space string) (string, error) {
128	records, err := LookupTXTRecords(host, space)
129	if err != nil {
130		return "", err
131	}
132	return strings.TrimSpace(records[0]), nil
133}
134
135func GetCustomDomain(host string, space string) string {
136	txt := TXTRecordName(host, space)
137	if record, found := txtCache.Get(txt); found {
138		return record
139	}
140	if _, missed := txtMissCache.Get(txt); missed {
141		return ""
142	}
143
144	record, err := LookupCustomDomain(host, space)
145	if errors.Is(err, ErrNoTXTRecord) {
146		txtMissCache.Add(txt, struct{}{})
147	}
148	if err != nil {
149		return ""
150	}
151	txtCache.Add(txt, record)
152	return record
153}
154
155func GetApiToken(r *http.Request) string {
156	authHeader := r.Header.Get("authorization")
157	if authHeader == "" {
158		return ""
159	}
160	return strings.TrimPrefix(authHeader, "Bearer ")
161}