diff --git a/cmd/gau/main.go b/cmd/gau/main.go index 44fad57..8b4e8a5 100644 --- a/cmd/gau/main.go +++ b/cmd/gau/main.go @@ -2,6 +2,7 @@ package main import ( "bufio" + "context" "io" "os" "sync" @@ -51,9 +52,10 @@ func main() { log.Fatalf("error writing results: %v\n", err) } }(out, config.JSON) - + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() workChan := make(chan runner.Work) - gau.Start(workChan, results) + gau.Start(ctx, workChan, results) domains := flags.Args() if len(domains) > 0 { for _, provider := range gau.Providers { @@ -63,14 +65,15 @@ func main() { } } else { sc := bufio.NewScanner(os.Stdin) - for _, provider := range gau.Providers { - for sc.Scan() { - workChan <- runner.NewWork(sc.Text(), provider) - } - if err := sc.Err(); err != nil { - log.Fatal(err) + for sc.Scan() { + domain := sc.Text() + for _, provider := range gau.Providers { + workChan <- runner.NewWork(domain, provider) } } + if err := sc.Err(); err != nil { + log.Fatal(err) + } } close(workChan) diff --git a/runner/runner.go b/runner/runner.go index 730aaba..28133ae 100644 --- a/runner/runner.go +++ b/runner/runner.go @@ -16,17 +16,14 @@ import ( type Runner struct { sync.WaitGroup - Providers []providers.Provider - threads uint - ctx context.Context - cancelFunc context.CancelFunc + Providers []providers.Provider + threads uint + ctx context.Context } // Init initializes the runner func (r *Runner) Init(c *providers.Config, providers []string, filters providers.Filters) error { r.threads = c.Threads - r.ctx, r.cancelFunc = context.WithCancel(context.Background()) - for _, name := range providers { switch name { case "urlscan": @@ -48,12 +45,12 @@ func (r *Runner) Init(c *providers.Config, providers []string, filters providers } // Starts starts the worker -func (r *Runner) Start(workChan chan Work, results chan string) { +func (r *Runner) Start(ctx context.Context, workChan chan Work, results chan string) { for i := uint(0); i < r.threads; i++ { r.Add(1) go func() { defer r.Done() - r.worker(r.ctx, workChan, results) + r.worker(ctx, workChan, results) }() } }