proxy.go

package proxy

import (
	"context"
	"net/http"
	"net/http/httputil"
	"net/url"

	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
)

// NewGCRProxy returns a http.Handler that proxies requests to gcr.io
// It will automatically set the relevant auth headers
func NewGCRProxy() (http.Handler, error) {
	url, err := url.Parse("https://gcr.io")
	if err != nil {
		return nil, err
	}

	tok, err := google.DefaultTokenSource(context.Background(),
		"https://www.googleapis.com/auth/cloud-platform",
		"https://www.googleapis.com/auth/devstorage.read_write",
	)
	if err != nil {
		return nil, err
	}

	reverse := httputil.NewSingleHostReverseProxy(url)
	reverse.Director = nil
	reverse.Rewrite = func(req *httputil.ProxyRequest) {
		req.Out = req.In.Clone(req.In.Context())

		// Override the host to be gcr.io
		req.Out.URL.Scheme = url.Scheme
		req.Out.URL.Host = url.Host
		req.Out.Host = url.Host
	}

	return &proxy{
		tokenSource: oauth2.ReuseTokenSource(nil, tok),
		proxy:       reverse,
	}, nil
}

type proxy struct {
	tokenSource oauth2.TokenSource
	proxy       *httputil.ReverseProxy
}

func (p *proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	token, err := p.tokenSource.Token()
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	token.SetAuthHeader(r)
	p.proxy.ServeHTTP(w, r)
}