From ec32ff386cf1ee0f300bc990aa4c1abfd3d2f342 Mon Sep 17 00:00:00 2001
From: Rick <1450685+LinuxSuRen@users.noreply.github.com>
Date: Thu, 8 Jun 2023 11:14:11 +0800
Subject: [PATCH] feat: support to set upstream proxy address (#86)

---
 extensions/collector/README.md                |  7 +++
 extensions/collector/cmd/collect.go           | 52 ++++++++++++++++---
 extensions/collector/cmd/collect_test.go      | 16 ++++--
 extensions/collector/pkg/collector.go         | 24 ++++++---
 extensions/collector/pkg/collector_test.go    |  5 +-
 extensions/collector/pkg/exporter.go          | 20 ++++---
 extensions/collector/pkg/exporter_test.go     | 12 +++--
 extensions/collector/pkg/filter/url_filter.go | 11 ++--
 .../collector/pkg/filter/url_filter_test.go   |  4 +-
 .../collector/pkg/testdata/sample_suite.yaml  |  3 +-
 10 files changed, 118 insertions(+), 36 deletions(-)

diff --git a/extensions/collector/README.md b/extensions/collector/README.md
index 95c1d83..669a014 100644
--- a/extensions/collector/README.md
+++ b/extensions/collector/README.md
@@ -10,3 +10,10 @@ It will start a HTTP proxy server, and set the server address to your browser pr
 
 `atest-collector` will record all HTTP requests which has prefix `/answer/api/v1`, and 
 save it to file `sample.yaml` once you close the server.
+
+## Features
+
+* Basic authorization
+* Upstream proxy
+* URL path filter
+* Support save response body or not
diff --git a/extensions/collector/cmd/collect.go b/extensions/collector/cmd/collect.go
index b04cd02..d629199 100644
--- a/extensions/collector/cmd/collect.go
+++ b/extensions/collector/cmd/collect.go
@@ -1,24 +1,33 @@
 package cmd
 
 import (
+	"bytes"
 	"context"
 	"fmt"
+	"io"
 	"net/http"
+	"net/url"
 	"os"
 	"os/signal"
 	"strings"
 	"syscall"
 
 	"github.com/elazarl/goproxy"
+	"github.com/elazarl/goproxy/ext/auth"
 	"github.com/linuxsuren/api-testing/extensions/collector/pkg"
 	"github.com/linuxsuren/api-testing/extensions/collector/pkg/filter"
 	"github.com/spf13/cobra"
 )
 
 type option struct {
-	port       int
-	filterPath string
-	output     string
+	port             int
+	filterPath       []string
+	saveResponseBody bool
+	output           string
+	upstreamProxy    string
+	verbose          bool
+	username         string
+	password         string
 }
 
 // NewRootCmd creates the root command
@@ -31,8 +40,13 @@ func NewRootCmd() (c *cobra.Command) {
 	}
 	flags := c.Flags()
 	flags.IntVarP(&opt.port, "port", "p", 8080, "The port for the proxy")
-	flags.StringVarP(&opt.filterPath, "filter-path", "", "", "The path prefix for filtering")
+	flags.StringSliceVarP(&opt.filterPath, "filter-path", "", []string{}, "The path prefix for filtering")
+	flags.BoolVarP(&opt.saveResponseBody, "save-response-body", "", false, "Save the response body")
 	flags.StringVarP(&opt.output, "output", "o", "sample.yaml", "The output file")
+	flags.StringVarP(&opt.upstreamProxy, "upstream-proxy", "", "", "The upstream proxy")
+	flags.StringVarP(&opt.username, "username", "", "", "The username for basic auth")
+	flags.StringVarP(&opt.password, "password", "", "", "The password for basic auth")
+	flags.BoolVarP(&opt.verbose, "verbose", "", false, "Verbose mode")
 
 	_ = cobra.MarkFlagRequired(flags, "filter-path")
 	return
@@ -41,6 +55,7 @@ func NewRootCmd() (c *cobra.Command) {
 type responseFilter struct {
 	urlFilter *filter.URLPathFilter
 	collects  *pkg.Collects
+	ctx       context.Context
 }
 
 func (f *responseFilter) filter(resp *http.Response, ctx *goproxy.ProxyCtx) *http.Response {
@@ -51,7 +66,16 @@ func (f *responseFilter) filter(resp *http.Response, ctx *goproxy.ProxyCtx) *htt
 
 	req := resp.Request
 	if f.urlFilter.Filter(req.URL) {
-		f.collects.Add(req.Clone(context.TODO()))
+		simpleResp := &pkg.SimpleResponse{StatusCode: resp.StatusCode}
+
+		if resp.Body != nil {
+			buf := new(bytes.Buffer)
+			io.Copy(buf, resp.Body)
+			simpleResp.Body = buf.String()
+			resp.Body = io.NopCloser(buf)
+		}
+
+		f.collects.Add(req.Clone(f.ctx), simpleResp)
 	}
 	return resp
 }
@@ -59,13 +83,25 @@ func (f *responseFilter) filter(resp *http.Response, ctx *goproxy.ProxyCtx) *htt
 func (o *option) runE(cmd *cobra.Command, args []string) (err error) {
 	urlFilter := &filter.URLPathFilter{PathPrefix: o.filterPath}
 	collects := pkg.NewCollects()
-	responseFilter := &responseFilter{urlFilter: urlFilter, collects: collects}
+	responseFilter := &responseFilter{urlFilter: urlFilter, collects: collects, ctx: cmd.Context()}
 
 	proxy := goproxy.NewProxyHttpServer()
-	proxy.Verbose = true
+	proxy.Verbose = o.verbose
+	if o.upstreamProxy != "" {
+		proxy.Tr.Proxy = func(r *http.Request) (*url.URL, error) {
+			return url.Parse(o.upstreamProxy)
+		}
+		proxy.ConnectDial = proxy.NewConnectDialToProxy(o.upstreamProxy)
+		cmd.Println("Using upstream proxy", o.upstreamProxy)
+	}
+	if o.username != "" && o.password != "" {
+		auth.ProxyBasic(proxy, "my_realm", func(user, pwd string) bool {
+			return user == o.username && o.password == pwd
+		})
+	}
 	proxy.OnResponse().DoFunc(responseFilter.filter)
 
-	exporter := pkg.NewSampleExporter()
+	exporter := pkg.NewSampleExporter(o.saveResponseBody)
 	collects.AddEvent(exporter.Add)
 
 	srv := &http.Server{
diff --git a/extensions/collector/cmd/collect_test.go b/extensions/collector/cmd/collect_test.go
index 6521ff0..523fb39 100644
--- a/extensions/collector/cmd/collect_test.go
+++ b/extensions/collector/cmd/collect_test.go
@@ -1,6 +1,9 @@
 package cmd
 
 import (
+	"bytes"
+	"context"
+	"io"
 	"net/http"
 	"net/url"
 	"testing"
@@ -17,19 +20,26 @@ func TestNewRootCmd(t *testing.T) {
 }
 
 func TestResponseFilter(t *testing.T) {
+	targetURL, err := url.Parse("http://foo.com/api/v1")
+	assert.NoError(t, err)
+
 	resp := &http.Response{
 		Header: http.Header{
 			"Content-Type": []string{"application/json; charset=utf-8"},
 		},
 		Request: &http.Request{
-			URL: &url.URL{},
+			URL: targetURL,
 		},
+		Body: io.NopCloser(bytes.NewBuffer([]byte("hello"))),
 	}
 	emptyResp := &http.Response{}
 
 	filter := &responseFilter{
-		urlFilter: &filter.URLPathFilter{},
-		collects:  pkg.NewCollects(),
+		urlFilter: &filter.URLPathFilter{
+			PathPrefix: []string{"/api/v1"},
+		},
+		collects: pkg.NewCollects(),
+		ctx:      context.Background(),
 	}
 	filter.filter(emptyResp, nil)
 	filter.filter(resp, nil)
diff --git a/extensions/collector/pkg/collector.go b/extensions/collector/pkg/collector.go
index 7ee9ae8..5d48881 100644
--- a/extensions/collector/pkg/collector.go
+++ b/extensions/collector/pkg/collector.go
@@ -11,33 +11,46 @@ type Collects struct {
 	once       sync.Once
 	signal     chan string
 	stopSignal chan struct{}
-	keys       map[string]*http.Request
+	keys       map[string]*RequestAndResponse
 	requests   []*http.Request
 	events     []EventHandle
 }
 
+type SimpleResponse struct {
+	StatusCode int
+	Body       string
+}
+
+type RequestAndResponse struct {
+	Request  *http.Request
+	Response *SimpleResponse
+}
+
 // NewCollects creates an instance of Collector
 func NewCollects() *Collects {
 	return &Collects{
 		once:       sync.Once{},
 		signal:     make(chan string, 5),
 		stopSignal: make(chan struct{}, 1),
-		keys:       make(map[string]*http.Request),
+		keys:       make(map[string]*RequestAndResponse),
 	}
 }
 
 // Add adds a HTTP request
-func (c *Collects) Add(req *http.Request) {
+func (c *Collects) Add(req *http.Request, resp *SimpleResponse) {
 	key := fmt.Sprintf("%s-%s", req.Method, req.URL.String())
 	if _, ok := c.keys[key]; !ok {
-		c.keys[key] = req
+		c.keys[key] = &RequestAndResponse{
+			Request:  req,
+			Response: resp,
+		}
 		c.requests = append(c.requests, req)
 		c.signal <- key
 	}
 }
 
 // EventHandle is the collect event handle
-type EventHandle func(r *http.Request)
+type EventHandle func(r *RequestAndResponse)
 
 // AddEvent adds new event handle
 func (c *Collects) AddEvent(e EventHandle) {
@@ -60,7 +73,6 @@ func (c *Collects) handleEvents() {
 				case key := <-c.signal:
 					fmt.Println("receive signal", key)
 					for _, e := range c.events {
-						fmt.Println("handle event", key, e)
 						e(c.keys[key])
 					}
 				case <-c.stopSignal:
diff --git a/extensions/collector/pkg/collector_test.go b/extensions/collector/pkg/collector_test.go
index 2581501..d74af7b 100644
--- a/extensions/collector/pkg/collector_test.go
+++ b/extensions/collector/pkg/collector_test.go
@@ -21,11 +21,12 @@ func TestCollector(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			collects := pkg.NewCollects()
-			collects.AddEvent(func(r *http.Request) {
+			collects.AddEvent(func(reqAndResp *pkg.RequestAndResponse) {
+				r := reqAndResp.Request
 				assert.Equal(t, tt.Request, r)
 			})
 			for i := 0; i < 10; i++ {
-				collects.Add(tt.Request)
+				collects.Add(tt.Request, nil)
 			}
 			collects.Stop()
 		})
diff --git a/extensions/collector/pkg/exporter.go b/extensions/collector/pkg/exporter.go
index 7278e69..bb9eaee 100644
--- a/extensions/collector/pkg/exporter.go
+++ b/extensions/collector/pkg/exporter.go
@@ -3,7 +3,6 @@ package pkg
 import (
 	"fmt"
 	"io"
-	"net/http"
 	"strings"
 
 	atestpkg "github.com/linuxsuren/api-testing/pkg/testing"
@@ -12,20 +11,23 @@ import (
 
 // SampleExporter is a sample exporter
 type SampleExporter struct {
-	TestSuite atestpkg.TestSuite
+	TestSuite        atestpkg.TestSuite
+	saveResponseBody bool
 }
 
 // NewSampleExporter creates a new exporter
-func NewSampleExporter() *SampleExporter {
+func NewSampleExporter(saveResponseBody bool) *SampleExporter {
 	return &SampleExporter{
 		TestSuite: atestpkg.TestSuite{
 			Name: "sample",
 		},
+		saveResponseBody: saveResponseBody,
 	}
 }
 
 // Add adds a request to the exporter
-func (e *SampleExporter) Add(r *http.Request) {
+func (e *SampleExporter) Add(reqAndResp *RequestAndResponse) {
+	r, resp := reqAndResp.Request, reqAndResp.Response
 
 	fmt.Println("receive", r.URL.Path)
 	req := atestpkg.Request{
@@ -42,9 +44,13 @@ func (e *SampleExporter) Add(r *http.Request) {
 
 	testCase := atestpkg.TestCase{
 		Request: req,
-		Expect: atestpkg.Response{
-			StatusCode: http.StatusOK,
-		},
+	}
+
+	if resp != nil {
+		testCase.Expect.StatusCode = resp.StatusCode
+		if e.saveResponseBody && resp.Body != "" {
+			testCase.Expect.Body = resp.Body
+		}
 	}
 
 	specs := strings.Split(r.URL.Path, "/")
diff --git a/extensions/collector/pkg/exporter_test.go b/extensions/collector/pkg/exporter_test.go
index 156a3cb..20293a0 100644
--- a/extensions/collector/pkg/exporter_test.go
+++ b/extensions/collector/pkg/exporter_test.go
@@ -12,15 +12,21 @@ import (
 )
 
 func TestSampleExporter(t *testing.T) {
-	exporter := pkg.NewSampleExporter()
+	exporter := pkg.NewSampleExporter(true)
 	assert.Equal(t, "sample", exporter.TestSuite.Name)
 
 	request, err := newRequest()
 	assert.NoError(t, err)
-	exporter.Add(request)
+	exporter.Add(&pkg.RequestAndResponse{Request: request})
 
 	request, err = newRequest()
-	exporter.Add(request)
+	exporter.Add(&pkg.RequestAndResponse{
+		Request: request,
+		Response: &pkg.SimpleResponse{
+			Body:       "hello",
+			StatusCode: http.StatusOK,
+		},
+	})
 
 	var result string
 	result, err = exporter.Export()
diff --git a/extensions/collector/pkg/filter/url_filter.go b/extensions/collector/pkg/filter/url_filter.go
index 67a785d..7ac4d45 100644
--- a/extensions/collector/pkg/filter/url_filter.go
+++ b/extensions/collector/pkg/filter/url_filter.go
@@ -1,7 +1,6 @@
 package filter
 
 import (
-	"fmt"
 	"net/url"
 	"strings"
 )
@@ -13,11 +12,15 @@ type URLFilter interface {
 
 // URLPathFilter filters the URL with path
 type URLPathFilter struct {
-	PathPrefix string
+	PathPrefix []string
 }
 
 // Filter implements the URLFilter
 func (f *URLPathFilter) Filter(targetURL *url.URL) bool {
-	fmt.Println(targetURL.Path, f.PathPrefix)
-	return strings.HasPrefix(targetURL.Path, f.PathPrefix)
+	for _, prefix := range f.PathPrefix {
+		if strings.HasPrefix(targetURL.Path, prefix) {
+			return true
+		}
+	}
+	return false
 }
diff --git a/extensions/collector/pkg/filter/url_filter_test.go b/extensions/collector/pkg/filter/url_filter_test.go
index ffb4078..d86917b 100644
--- a/extensions/collector/pkg/filter/url_filter_test.go
+++ b/extensions/collector/pkg/filter/url_filter_test.go
@@ -9,6 +9,8 @@ import (
 )
 
 func TestURLPathFilter(t *testing.T) {
-	urlFilter := &filter.URLPathFilter{PathPrefix: "/api"}
+	urlFilter := &filter.URLPathFilter{PathPrefix: []string{"/api/v1", "/api/v2"}}
 	assert.True(t, urlFilter.Filter(&url.URL{Path: "/api/v1"}))
+	assert.True(t, urlFilter.Filter(&url.URL{Path: "/api/v2"}))
+	assert.False(t, urlFilter.Filter(&url.URL{Path: "/api/v3"}))
 }
diff --git a/extensions/collector/pkg/testdata/sample_suite.yaml b/extensions/collector/pkg/testdata/sample_suite.yaml
index aa081d2..41705b2 100644
--- a/extensions/collector/pkg/testdata/sample_suite.yaml
+++ b/extensions/collector/pkg/testdata/sample_suite.yaml
@@ -10,8 +10,6 @@ items:
       Authorization: Bearer token
       Content-Type: application/json
     body: hello
-  expect:
-    statusCode: 200
 - name: v1-1
   request:
     api: http://foo/api/v1
@@ -22,3 +20,4 @@ items:
     body: hello
   expect:
     statusCode: 200
+    body: hello