From 769cffd5df9d9fb15876a8fcb4551a431c35d9c1 Mon Sep 17 00:00:00 2001
From: imneov <neolf@qq.com>
Date: Wed, 31 Jul 2019 18:09:24 +0800
Subject: [PATCH] add router mode

---
 README.md       | 27 +++++++++++++++++++++
 main.go         | 26 ++++++++++++++++----
 proxy/router.go | 63 +++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 111 insertions(+), 5 deletions(-)
 create mode 100644 proxy/router.go

diff --git a/README.md b/README.md
index df49bc4..29fc83c 100644
--- a/README.md
+++ b/README.md
@@ -10,9 +10,36 @@ A global proxy for go modules. see: [https://goproxy.io](https://goproxy.io)
     make
 
 ## Started
+
+
+### Proxy mode    
     
     ./bin/goproxy -listen=0.0.0.0:80 -cacheDir=/tmp/test
 
+### Router mode    
+
+Use the -proxy flag switch to "Router mode", which 
+implements route filter to routing private module 
+or public module .
+
+```
+                                         direct
+                      +----------------------------------> private repo
+                      |
+                 match|pattern
+                      |
+                  +---+---+           +----------+
+go get  +-------> |goproxy| +-------> |goproxy.io| +---> golang.org/x/net
+                  +-------+           +----------+
+                 router mode           proxy mode
+```
+
+In Router mode, use the -exclude flag set pattern , direct to the repo which 
+match the module path, pattern are matched to the full path specified, not only 
+to the host component.
+
+    ./bin/goproxy -listen=0.0.0.0:80 -cacheDir=/tmp/test -proxy https://goproxy.io -exclude "git.private.domain/[abc]"
+
 ## Use docker image
 
     docker run -d -p80:8081 goproxy/goproxy
diff --git a/main.go b/main.go
index e39f6d5..0aa91ed 100644
--- a/main.go
+++ b/main.go
@@ -42,8 +42,12 @@ const listExpire = 5 * time.Minute
 
 var listen string
 var cacheDir string
+var proxyHost string
+var excludeHost string
 
 func init() {
+	flag.StringVar(&excludeHost, "exclude", "", "exclude host pattern")
+	flag.StringVar(&proxyHost, "proxy", "", "next hop proxy for go modules")
 	flag.StringVar(&cacheDir, "cacheDir", "", "go modules cache dir")
 	flag.StringVar(&listen, "listen", "0.0.0.0:8081", "service listen address")
 	flag.Parse()
@@ -64,6 +68,10 @@ func main() {
 	var env struct {
 		GOPATH string
 	}
+	if cacheDir != "" {
+		downloadRoot = filepath.Join(cacheDir, "pkg/mod/cache/download")
+		os.Setenv("GOPATH", cacheDir)
+	}
 	if err := goJSON(&env, "go", "env", "-json", "GOPATH"); err != nil {
 		log.Fatal(err)
 	}
@@ -73,12 +81,20 @@ func main() {
 	}
 	downloadRoot = filepath.Join(list[0], "pkg/mod/cache/download")
 
-	if cacheDir != "" {
-		downloadRoot = filepath.Join(cacheDir, "pkg/mod/cache/download")
-		os.Setenv("GOPATH", cacheDir)
+	var handle http.Handler
+	if proxyHost != "" {
+		fmt.Fprintf(os.Stderr, "ProxyHost %s\n", proxyHost)
+		if excludeHost != "" {
+			fmt.Fprintf(os.Stderr, "ExcludeHost %s\n", excludeHost)
+		}
+		handle = &logger{proxy.NewRouter(proxy.NewServer(new(ops)), &proxy.RouterOps{
+			Pattern: excludeHost,
+			Proxy:   proxyHost,
+		})}
+	} else {
+		handle = &logger{proxy.NewServer(new(ops))}
 	}
-
-	log.Fatal(http.ListenAndServe(listen, &logger{proxy.NewServer(new(ops))}))
+	log.Fatal(http.ListenAndServe(listen, handle))
 }
 
 // goJSON runs the go command and parses its JSON output into dst.
diff --git a/proxy/router.go b/proxy/router.go
new file mode 100644
index 0000000..e628fe7
--- /dev/null
+++ b/proxy/router.go
@@ -0,0 +1,63 @@
+package proxy
+
+import (
+	"crypto/tls"
+	"fmt"
+	"net/http"
+	"net/http/httputil"
+	"net/url"
+	"os"
+	"regexp"
+)
+
+// A RouterOps provides the proxy host and the external pattern
+type RouterOps struct {
+	Pattern string
+	Proxy   string
+}
+
+// A Router is the proxy HTTP server,
+// which implements Route Filter to
+// routing private module or public module .
+type Router struct {
+	srv   *Server
+	proxy *httputil.ReverseProxy
+	regex *regexp.Regexp
+}
+
+// NewRouter returns a new Router using the given operations.
+func NewRouter(srv *Server, ops *RouterOps) *Router {
+	rt := &Router{
+		srv: srv,
+	}
+	if ops != nil {
+		if remote, err := url.Parse(ops.Proxy); err == nil {
+			rt.proxy = httputil.NewSingleHostReverseProxy(remote)
+			rt.proxy.Transport = &http.Transport{
+				TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
+			}
+		}
+		if regex, err := regexp.Compile(ops.Pattern); err == nil {
+			rt.regex = regex
+		}
+	}
+	return rt
+}
+
+func (rt *Router) Direct(path string) bool {
+	if rt.regex == nil {
+		return true
+	}
+	return rt.regex.Match([]byte(path))
+}
+
+func (rt *Router) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	if rt.proxy != nil && rt.Direct(r.URL.Path) {
+		fmt.Fprintf(os.Stdout, "[direct] %s\n", r.URL)
+		rt.srv.ServeHTTP(w, r)
+		return
+	}
+	fmt.Fprintf(os.Stdout, "[proxy] %s\n", r.URL)
+	rt.proxy.ServeHTTP(w, r)
+	return
+}