change sumdb proxy impl

sumdb/handler.go

@@ -6,78 +6,101 @@
 package sumdb
 
 import (
+	"context"
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
 	"net/url"
 	"strings"
+	"time"
 )
 
-var enableGoogleSumDB bool
-var supportedSumDB = []string{
-	"sum.golang.org",
-	"gosum.io",
+var supportedSumDB = map[string][]string{
+	"sum.golang.org":       {"https://sum.golang.org/", "https://sum.golang.google.cn/"},
+	"sum.golang.google.cn": {"https://sum.golang.org/", "https://sum.golang.google.cn/"}, // db-name `sum.golang.google.cn` will be replaced in go
+	"gosum.io":             {"https://gosum.io/"},
 }
 
-func init() {
-	go func() {
-		p := "https://sum.golang.org"
-		_, err := http.Get(p)
-		if err == nil {
-			enableGoogleSumDB = true
-		}
-	}()
-}
+var (
+	errSumPathInvalid = errors.New("sumdb request path invalid")
+)
 
-//Handler handles sumdb request
+// Handler handles sumdb request
+// goproxy.io not impl a complete sumdb, just proxy to upstream.
 func Handler(w http.ResponseWriter, r *http.Request) {
-	if strings.HasSuffix(r.URL.Path, "/supported") {
-		for _, supported := range supportedSumDB {
-			uri := fmt.Sprintf("/sumdb/%s/supported", supported)
-			if r.URL.Path == uri {
-				w.WriteHeader(http.StatusOK)
-				return
-			}
-		}
-
+	whichDB, realPath, err := parsePath(r.URL.Path)
+	_, supported := supportedSumDB[whichDB]
+	if err != nil || !supported {
+		// if not check the target db,
+		// curl https://goproxy.io/sumdb/www.google.com will succ
 		w.WriteHeader(http.StatusGone)
+		fmt.Fprint(w, "unsupported db")
 		return
 	}
 
-	if !enableGoogleSumDB {
-		sumViaGoproxy(w, r)
+	// $GOROOT/src/cmd/go/internal/modfetch/sumdb.go@initBase
+	// > Before accessing any checksum database URL using a proxy, the proxy
+	// > client should first fetch <proxyURL>/sumdb/<sumdb-name>/supported.
+	if realPath == "supported" {
+		w.WriteHeader(http.StatusOK)
 		return
 	}
 
-	p := "https://" + strings.TrimPrefix(r.URL.Path, "/sumdb/")
-	proxySumdb(p, w, r)
+	result := make(chan *http.Response)
+	ctx, cancel := context.WithTimeout(r.Context(), 2*time.Second) // 2sec maybe enough
+	defer cancel()
+
+	for _, host := range supportedSumDB[whichDB] {
+		go proxySumdb(ctx, host, realPath, result)
+	}
+
+	select {
+	case resp := <-result:
+		{
+			defer resp.Body.Close()
+			w.WriteHeader(resp.StatusCode)
+			if _, err := io.Copy(w, resp.Body); err != nil {
+				fmt.Fprint(w, err.Error())
+			}
+		}
+	case <-ctx.Done():
+		w.WriteHeader(http.StatusGone)
+		fmt.Fprint(w, ctx.Err().Error())
+		return
+	}
 }
 
-func sumViaGoproxy(w http.ResponseWriter, r *http.Request) {
-	p := "https://goproxy.io" + r.URL.Path
-	proxySumdb(p, w, r)
+func parsePath(rawPath string) (whichDB, path string, err error) {
+	parts := strings.SplitN(rawPath, "/", 4)
+	if len(parts) < 4 {
+		return "", "", errSumPathInvalid
+	}
+	whichDB = parts[2]
+	path = parts[3]
+	return
 }
 
-func proxySumdb(p string, w http.ResponseWriter, r *http.Request) {
-	_, err := url.Parse(p)
+func proxySumdb(ctx context.Context, host, path string, respChan chan<- *http.Response) {
+	urlPath, err := url.Parse(host)
 	if err != nil {
-		w.WriteHeader(http.StatusGone)
-		fmt.Fprintf(w, err.Error())
 		return
 	}
+	urlPath.Path = path
 
-	resp, err := http.Get(p)
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, urlPath.String(), nil)
 	if err != nil {
-		w.WriteHeader(http.StatusGone)
-		fmt.Fprintf(w, err.Error())
 		return
 	}
-	defer resp.Body.Close()
-	w.WriteHeader(resp.StatusCode)
-	if _, err := io.Copy(w, resp.Body); err != nil {
-		fmt.Fprintf(w, err.Error())
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
 		return
 	}
-	return
+
+	select {
+	case <-ctx.Done():
+		resp.Body.Close()
+	case respChan <- resp:
+	}
 
 }

sumdb/handler_test.go

@@ -0,0 +1,98 @@
+// Copyright 2019 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Package sumdb implements sumdb handler proxy.
+package sumdb
+
+import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+)
+
+func TestHandler(t *testing.T) {
+	if ret := t.Run("supported", testSupported); !ret {
+		t.Error("supported test failed, stop test")
+		t.FailNow()
+	}
+	t.Run("proxy", testProxy)
+}
+
+func testSupported(t *testing.T) {
+	type TestCase struct {
+		name          string
+		db            string
+		wantSupported bool
+	}
+
+	tests := []TestCase{
+		{
+			name:          "sum.golang.org",
+			db:            "sum.golang.org",
+			wantSupported: true,
+		},
+		{
+			name:          "gosum.io",
+			db:            "gosum.io",
+			wantSupported: true,
+		},
+		{
+			name:          "sum.golang.google.cn",
+			db:            "sum.golang.google.cn",
+			wantSupported: true,
+		},
+		{
+			name:          "other",
+			db:            "other",
+			wantSupported: false,
+		},
+	}
+
+	for _, testcase := range tests {
+		t.Run(testcase.name, func(t *testing.T) {
+			recoder := httptest.NewRecorder()
+			req := httptest.NewRequest(http.MethodGet, fmt.Sprintf("https://goproxy.io/sumdb/%s/supported", testcase.db), nil)
+			Handler(recoder, req)
+
+			resp := recoder.Result()
+			if support := (resp.StatusCode == http.StatusOK); support != testcase.wantSupported {
+				t.Errorf("db %s: want %v got %v", testcase.db, testcase.wantSupported, support)
+			}
+			resp.Body.Close()
+		})
+	}
+}
+
+func testProxy(t *testing.T) {
+	type TestCase struct {
+		name       string
+		db         string
+		path       string
+		expectSucc bool
+	}
+	tests := []TestCase{
+		{
+			name:       "lookup",
+			db:         "sum.golang.google.cn",
+			path:       "lookup/github.com/goproxyio/goproxy@v1.0.0", // this is a fake testcase
+			expectSucc: true,
+		},
+	}
+
+	for _, testcase := range tests {
+		t.Run(testcase.name, func(t *testing.T) {
+
+			recoder := httptest.NewRecorder()
+			req := httptest.NewRequest(http.MethodGet, fmt.Sprintf("https://goproxy.io/sumdb/%s/%s", testcase.db, testcase.path), nil)
+			Handler(recoder, req)
+
+			resp := recoder.Result()
+			if succ := (resp.StatusCode == http.StatusOK); succ != testcase.expectSucc {
+				t.Errorf("FETCH from db %s/%s got unexpect http status %d", testcase.db, testcase.path, resp.StatusCode)
+				return
+			}
+		})
+	}
+}