Skip to content

Commit

Permalink
i am sorry for this one...fix stdin providers
Browse files Browse the repository at this point in the history
  • Loading branch information
lc committed May 14, 2024
1 parent bdf392f commit 67e9445
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 16 deletions.
19 changes: 11 additions & 8 deletions cmd/gau/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package main

import (
"bufio"
"context"
"io"
"os"
"sync"
Expand Down Expand Up @@ -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 {
Expand All @@ -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)

Expand Down
13 changes: 5 additions & 8 deletions runner/runner.go
Original file line number Diff line number Diff line change
Expand Up @@ -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":
Expand All @@ -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)
}()
}
}
Expand Down

0 comments on commit 67e9445

Please sign in to comment.