cmd/pkgsite: support proxy

If the -proxy flag is provided, the pkgsite command will fetch modules
from the proxy if they can't be found locally.

For golang/go#47780

Change-Id: I02b9fddac9013d9d3859b78e6c8887282469c9a8
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/345271
Trust: Jonathan Amsterdam <jba@google.com>
Run-TryBot: Jonathan Amsterdam <jba@google.com>
Reviewed-by: Jamal Carvalho <jamal@golang.org>
Reviewed-by: Julie Qiu <julie@golang.org>

2 files changed, 46 insertions, 12 deletions

diff --git a/cmd/pkgsite/main.go b/cmd/pkgsite/main.go
index 95895fa0..31f11020 100644
--- a/cmd/pkgsite/main.go
+++ b/cmd/pkgsite/main.go
@@ -37,6 +37,7 @@ import (
 	"golang.org/x/pkgsite/internal/frontend"
 	"golang.org/x/pkgsite/internal/log"
 	"golang.org/x/pkgsite/internal/middleware"
+	"golang.org/x/pkgsite/internal/proxy"
 	"golang.org/x/pkgsite/internal/source"
 )
 
@@ -46,6 +47,7 @@ var (
 	_          = flag.String("static", "static", "path to folder containing static files served")
 	gopathMode = flag.Bool("gopath_mode", false, "assume that local modules' paths are relative to GOPATH/src")
 	httpAddr   = flag.String("http", defaultAddr, "HTTP service address to listen for incoming requests on")
+	useProxy   = flag.Bool("proxy", false, "fetch from GOPROXY if not found locally")
 )
 
 func main() {
@@ -62,15 +64,33 @@ func main() {
 		paths = []string{"."}
 	}
 
-	server, err := newServer(ctx, paths, *gopathMode)
+	var prox *proxy.Client
+	if *useProxy {
+		fmt.Fprintf(os.Stderr, "BYPASSING LICENSE CHECKING: MAY DISPLAY NON-REDISTRIBUTABLE INFORMATION\n")
+		url := os.Getenv("GOPROXY")
+		if url == "" {
+			die("GOPROXY environment variable is not set")
+		}
+		var err error
+		prox, err = proxy.New(url)
+		if err != nil {
+			die("connecting to proxy: %s", err)
+		}
+	}
+	server, err := newServer(ctx, paths, *gopathMode, prox)
 	if err != nil {
-		log.Fatalf(ctx, "newServer: %v", err)
+		die("%s", err)
 	}
 	router := dcensus.NewRouter(frontend.TagRoute)
 	server.Install(router.Handle, nil, nil)
 	mw := middleware.Timeout(54 * time.Second)
 	log.Infof(ctx, "Listening on addr %s", *httpAddr)
-	log.Fatal(ctx, http.ListenAndServe(*httpAddr, mw(router)))
+	die("%v", http.ListenAndServe(*httpAddr, mw(router)))
+}
+
+func die(format string, args ...interface{}) {
+	fmt.Fprintf(os.Stderr, format, args...)
+	os.Exit(1)
 }
 
 func collectPaths(args []string) []string {
@@ -81,12 +101,16 @@ func collectPaths(args []string) []string {
 	return paths
 }
 
-func newServer(ctx context.Context, paths []string, gopathMode bool) (*frontend.Server, error) {
+func newServer(ctx context.Context, paths []string, gopathMode bool, prox *proxy.Client) (*frontend.Server, error) {
 	getters := buildGetters(ctx, paths, gopathMode)
+	if prox != nil {
+		getters = append(getters, fetch.NewProxyModuleGetter(prox))
+	}
 	lds := fetchdatasource.Options{
-		Getters:            getters,
-		SourceClient:       source.NewClient(time.Second),
-		BypassLicenseCheck: true,
+		Getters:              getters,
+		SourceClient:         source.NewClient(time.Second),
+		ProxyClientForLatest: prox,
+		BypassLicenseCheck:   true,
 	}.New()
 	server, err := frontend.NewServer(frontend.ServerConfig{
 		DataSourceGetter: func(context.Context) internal.DataSource { return lds },
diff --git a/cmd/pkgsite/main_test.go b/cmd/pkgsite/main_test.go
index b7490764..3a90036e 100644
--- a/cmd/pkgsite/main_test.go
+++ b/cmd/pkgsite/main_test.go
@@ -9,14 +9,22 @@ import (
 	"flag"
 	"net/http"
 	"net/http/httptest"
+	"path/filepath"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
+	"golang.org/x/pkgsite/internal/proxy/proxytest"
 )
 
 func Test(t *testing.T) {
-	flag.Set("static", "../../static")
-	server, err := newServer(context.Background(), []string{"../../internal/fetch/testdata/has_go_mod"}, false)
+	repoPath := func(fn string) string { return filepath.Join("..", "..", fn) }
+	localModule := repoPath("internal/fetch/testdata/has_go_mod")
+	flag.Set("static", repoPath("static"))
+	testModules := proxytest.LoadTestModules(repoPath("internal/proxy/testdata"))
+	prox, teardown := proxytest.SetupTestClient(t, testModules)
+	defer teardown()
+
+	server, err := newServer(context.Background(), []string{localModule}, false, prox)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -24,9 +32,11 @@ func Test(t *testing.T) {
 	server.Install(mux.Handle, nil, nil)
 	w := httptest.NewRecorder()
 
-	mux.ServeHTTP(w, httptest.NewRequest("GET", "/example.com/testmod", nil))
-	if w.Code != http.StatusOK {
-		t.Errorf("%q: got status code = %d, want %d", "/testmod", w.Code, http.StatusOK)
+	for _, url := range []string{"/example.com/testmod", "/example.com/single/pkg"} {
+		mux.ServeHTTP(w, httptest.NewRequest("GET", url, nil))
+		if w.Code != http.StatusOK {
+			t.Errorf("%q: got status code = %d, want %d", url, w.Code, http.StatusOK)
+		}
 	}
 }