magpie/internal/source/source.go

// Package source chooses where magpie reads from, with a mirror as a
// fallback when the official address can't be reached. The official source
// stays first; a mirror is only asked when it fails.
package source

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
)

const (
	npmMirror    = "https://registry.npmmirror.com"
	githubMirror = "https://gh-proxy.com"
)

// Enabled reports whether mirror fallback was turned on with
// MAGPIE_MIRRORS=on (or true, 1, cn).
func Enabled() bool {
	switch strings.ToLower(strings.TrimSpace(os.Getenv("MAGPIE_MIRRORS"))) {
	case "on", "true", "1", "cn":
		return true
	}
	return false
}

func npmBase() string {
	if v := strings.TrimRight(strings.TrimSpace(os.Getenv("MAGPIE_NPM_REGISTRY")), "/"); v != "" {
		return v
	}
	return npmMirror
}

func githubBase() string {
	if v := strings.TrimRight(strings.TrimSpace(os.Getenv("MAGPIE_GITHUB_MIRROR")), "/"); v != "" {
		return v
	}
	return githubMirror
}

// URLs is the official address followed by its mirrors, or just the
// official one when there is no mirror for it.
func URLs(raw string) []string {
	if !Enabled() {
		return []string{raw}
	}
	u, err := url.Parse(raw)
	if err != nil {
		return []string{raw}
	}
	switch strings.ToLower(u.Hostname()) {
	case "registry.npmjs.org":
		return distinct(raw, npmBase()+u.RequestURI())
	case "github.com", "api.github.com", "raw.githubusercontent.com", "codeload.github.com":
		return distinct(raw, githubBase()+"/"+raw)
	}
	return []string{raw}
}

func distinct(raw string, mirrors ...string) []string {
	out := []string{raw}
	for _, m := range mirrors {
		if m != "" && m != raw {
			out = append(out, m)
		}
	}
	return out
}

// Do sends req as usual and, when the official answer fails in a way a
// mirror could answer, tries its mirrors in order. A request that carries
// credentials is never sent to a mirror.
func Do(c *http.Client, req *http.Request) (*http.Response, error) {
	if !mirrorable(req) {
		return do(c, req, []string{req.URL.String()})
	}
	return do(c, req, URLs(req.URL.String()))
}

// mirrorable is whether req may be sent to a mirror. Credentials stay with
// the official source: a mirror must never see a token, cookie or key.
func mirrorable(req *http.Request) bool {
	if req.URL == nil || req.URL.User != nil {
		return false
	}
	for k := range req.URL.Query() {
		k = strings.ToLower(k)
		if strings.Contains(k, "token") || strings.Contains(k, "key") ||
			strings.Contains(k, "secret") || strings.Contains(k, "password") ||
			strings.Contains(k, "auth") {
			return false
		}
	}
	for _, h := range []string{
		"Authorization",
		"Proxy-Authorization",
		"Cookie",
		"X-Api-Key",
		"Api-Key",
		"Private-Token",
	} {
		if req.Header.Get(h) != "" {
			return false
		}
	}
	return true
}

// DoOfficial sends req to its official address alone. It is for a file
// whose checksum must come from the same trusted place as the file itself.
func DoOfficial(c *http.Client, req *http.Request) (*http.Response, error) {
	return do(c, req, []string{req.URL.String()})
}

func do(c *http.Client, req *http.Request, urls []string) (*http.Response, error) {
	var last error
	for i, raw := range urls {
		r := req.Clone(req.Context())
		if i > 0 {
			u, err := url.Parse(raw)
			if err != nil {
				last = err
				continue
			}
			r.URL = u
			r.Host = ""
			if r.Body != nil {
				if r.GetBody == nil {
					break
				}
				body, err := r.GetBody()
				if err != nil {
					last = err
					break
				}
				r.Body = body
			}
		}
		var cancel context.CancelFunc
		var stop func() bool
		if len(urls) > 1 {
			ctx, cancelCandidate, stopCandidate := candidateContext(r.Context(), c, i < len(urls)-1)
			if ctx != r.Context() {
				r = r.WithContext(ctx)
				cancel = cancelCandidate
				stop = stopCandidate
			}
		}
		resp, err := c.Do(r)
		if err == nil && (!retryable(resp) || i == len(urls)-1) {
			if stop != nil {
				stop()
			}
			return resp, nil
		}
		if cancel != nil {
			cancel()
		}
		if err == nil {
			resp.Body.Close()
			last = errors.Join(last, fmt.Errorf("%s: %s", raw, resp.Status))
			continue
		}
		last = errors.Join(last, fmt.Errorf("%s: %w", raw, err))
	}
	if last == nil {
		last = fmt.Errorf("no address for %s", req.URL)
	}
	return nil, last
}

// candidateContext bounds a source's try to half of the request's
// remaining budget, so an official source that hangs doesn't make the
// next one's try double the whole wait. The last source keeps all of
// the time left. The timer only guards the wait for response headers;
// stop leaves the response's context alive for its body.
func candidateContext(ctx context.Context, c *http.Client, more bool) (context.Context, context.CancelFunc, func() bool) {
	if !more {
		return ctx, func() {}, func() bool { return false }
	}
	var d time.Duration
	if deadline, ok := ctx.Deadline(); ok {
		d = time.Until(deadline) / 2
	} else if c.Timeout > 0 {
		d = c.Timeout / 2
	}
	if d <= 0 {
		return ctx, func() {}, func() bool { return false }
	}
	ctx, cancel := context.WithCancel(ctx)
	timer := time.AfterFunc(d, cancel)
	return ctx, cancel, timer.Stop
}

func retryable(resp *http.Response) bool {
	return resp.StatusCode == http.StatusTooManyRequests ||
		resp.StatusCode >= 500 ||
		(resp.StatusCode == http.StatusForbidden && resp.Header.Get("X-RateLimit-Remaining") == "0")
}