|
| 1 | +package cmd |
| 2 | + |
| 3 | +import ( |
| 4 | + "context" |
| 5 | + "fmt" |
| 6 | + "net/http" |
| 7 | + "os" |
| 8 | + "os/signal" |
| 9 | + "strings" |
| 10 | + "syscall" |
| 11 | + |
| 12 | + "github.com/elazarl/goproxy" |
| 13 | + "github.com/linuxsuren/api-testing/extensions/collector/pkg" |
| 14 | + "github.com/linuxsuren/api-testing/extensions/collector/pkg/filter" |
| 15 | + "github.com/spf13/cobra" |
| 16 | +) |
| 17 | + |
| 18 | +type option struct { |
| 19 | + port int |
| 20 | + filterPath string |
| 21 | + output string |
| 22 | +} |
| 23 | + |
| 24 | +// NewRootCmd creates the root command |
| 25 | +func NewRootCmd() (c *cobra.Command) { |
| 26 | + opt := &option{} |
| 27 | + c = &cobra.Command{ |
| 28 | + Use: "atest-collector", |
| 29 | + Short: "A collector for API testing, it will start a HTTP proxy server", |
| 30 | + RunE: opt.runE, |
| 31 | + } |
| 32 | + flags := c.Flags() |
| 33 | + flags.IntVarP(&opt.port, "port", "p", 8080, "The port for the proxy") |
| 34 | + flags.StringVarP(&opt.filterPath, "filter-path", "", "", "The path prefix for filtering") |
| 35 | + flags.StringVarP(&opt.output, "output", "o", "sample.yaml", "The output file") |
| 36 | + |
| 37 | + _ = cobra.MarkFlagRequired(flags, "filter-path") |
| 38 | + return |
| 39 | +} |
| 40 | + |
| 41 | +type responseFilter struct { |
| 42 | + urlFilter *filter.URLPathFilter |
| 43 | + collects *pkg.Collects |
| 44 | +} |
| 45 | + |
| 46 | +func (f *responseFilter) filter(resp *http.Response, ctx *goproxy.ProxyCtx) *http.Response { |
| 47 | + contentType := resp.Header.Get("Content-Type") |
| 48 | + if !strings.Contains(contentType, "application/json") { |
| 49 | + return resp |
| 50 | + } |
| 51 | + |
| 52 | + req := resp.Request |
| 53 | + if f.urlFilter.Filter(req.URL) { |
| 54 | + f.collects.Add(req.Clone(context.TODO())) |
| 55 | + } |
| 56 | + return resp |
| 57 | +} |
| 58 | + |
| 59 | +func (o *option) runE(cmd *cobra.Command, args []string) (err error) { |
| 60 | + urlFilter := &filter.URLPathFilter{PathPrefix: o.filterPath} |
| 61 | + collects := pkg.NewCollects() |
| 62 | + responseFilter := &responseFilter{urlFilter: urlFilter, collects: collects} |
| 63 | + |
| 64 | + proxy := goproxy.NewProxyHttpServer() |
| 65 | + proxy.Verbose = true |
| 66 | + proxy.OnResponse().DoFunc(responseFilter.filter) |
| 67 | + |
| 68 | + exporter := pkg.NewSampleExporter() |
| 69 | + collects.AddEvent(exporter.Add) |
| 70 | + |
| 71 | + srv := &http.Server{ |
| 72 | + Addr: fmt.Sprintf(":%d", o.port), |
| 73 | + Handler: proxy, |
| 74 | + } |
| 75 | + |
| 76 | + sig := make(chan os.Signal, 1) |
| 77 | + signal.Notify(sig, os.Interrupt, syscall.SIGTERM) |
| 78 | + go func() { |
| 79 | + <-sig |
| 80 | + collects.Stop() |
| 81 | + _ = srv.Shutdown(context.Background()) |
| 82 | + }() |
| 83 | + |
| 84 | + cmd.Println("Starting the proxy server with port", o.port) |
| 85 | + _ = srv.ListenAndServe() |
| 86 | + var data string |
| 87 | + if data, err = exporter.Export(); err == nil { |
| 88 | + err = os.WriteFile(o.output, []byte(data), 0644) |
| 89 | + } |
| 90 | + return |
| 91 | +} |
0 commit comments