diff options
| -rw-r--r-- | internal/qna/qna.go | 7 | ||||
| -rw-r--r-- | main.go | 74 |
2 files changed, 34 insertions, 47 deletions
diff --git a/internal/qna/qna.go b/internal/qna/qna.go index aa27d07..d883016 100644 --- a/internal/qna/qna.go +++ b/internal/qna/qna.go @@ -128,3 +128,10 @@ func isDeadlineError(err error) bool { urlErr, ok := err.(*url.Error) return ok && errors.Is(urlErr.Err, context.DeadlineExceeded) } + +func TrimWildcard(san string) string { + if len(san) >= 2 && san[:2] == "*." { + return san[2:] + } + return san +} @@ -1,13 +1,5 @@ // Package main provides onion-grab, a tool that visits a list of domains over // HTTPS to see if they have Onion-Location configured. -// -// Install: -// -// $ go install git.cs.kau.se/rasmoste/onion-grab@latest -// -// Usage: -// -// $ onion-grab -h package main import ( @@ -36,7 +28,6 @@ func main() { MaxResponseHeaderBytes: opts.MaxResponse * 1024 * 1024, }, } - fp, err := os.OpenFile(opts.InputFile, os.O_RDONLY, 0644) if err != nil { log.Printf("ERROR: %v", err) @@ -56,43 +47,50 @@ func main() { ctx, cancel := context.WithCancel(context.Background()) defer cancel() - log.Printf("INFO: starting await handler, ctrl+C to exit\n") + start := time.Now().Round(time.Second) + defer func() { + end := time.Now().Round(time.Second) + log.Printf("INFO: measurement duration was %v\n", end.Sub(start)) + }() + + log.Printf("INFO: ctrl+C to exit prematurely\n") go func() { wg.Add(1) defer wg.Done() - awaitHandler(ctx, cancel) + await(ctx, cancel) }() - log.Printf("INFO: starting %d workers with rate-limit %d requests/s\n", opts.NumWorkers, opts.Limit) + log.Printf("INFO: starting %d workers\n", opts.Limit) for i := 0; i < opts.NumWorkers; i++ { go func() { wg.Add(1) defer wg.Done() - workHandler(ctx, opts, cli, questionCh, answerCh) + for { + select { + case <-ctx.Done(): + return + case question := <-questionCh: + work(ctx, cli, opts.Timeout, question, answerCh) + } + } }() } - log.Printf("INFO: starting work aggregator\n") + log.Printf("INFO: starting work receiver\n") go func() { wg.Add(1) defer wg.Done() - workAggregator(ctx, opts, answerCh) + workReceiver(ctx, opts, answerCh) }() - start := time.Now().Round(time.Second) - defer func() { - end := time.Now().Round(time.Second) - log.Printf("INFO: measurement duration was %v\n", end.Sub(start)) - }() - - log.Printf("INFO: generating work\n") + log.Printf("INFO: starting work generator\n") workGenerator(ctx, opts, fp, questionCh) if err := ctx.Err(); err != nil { log.Printf("WARNING: premature shutdown (context cancelled)\n") } } -func awaitHandler(ctx context.Context, cancel context.CancelFunc) { +func await(ctx context.Context, cancel context.CancelFunc) { sigs := make(chan os.Signal, 1) defer close(sigs) @@ -105,19 +103,8 @@ func awaitHandler(ctx context.Context, cancel context.CancelFunc) { cancel() } -func workHandler(ctx context.Context, opts options.Options, cli *http.Client, questionCh chan qna.Question, answerCh chan qna.Answer) { - for { - select { - case <-ctx.Done(): - return - case question := <-questionCh: - work(ctx, &opts, cli, question, answerCh) - } - } -} - -func work(ctx context.Context, opts *options.Options, cli *http.Client, question qna.Question, answerCh chan qna.Answer) { - cctx, cancel := context.WithTimeout(ctx, opts.Timeout) +func work(ctx context.Context, cli *http.Client, timeout time.Duration, question qna.Question, answerCh chan qna.Answer) { + cctx, cancel := context.WithTimeout(ctx, timeout) defer cancel() answer := qna.Answer{Domain: question.Domain} @@ -147,7 +134,7 @@ func work(ctx context.Context, opts *options.Options, cli *http.Client, question answerCh <- answer } -func workAggregator(ctx context.Context, opts options.Options, answerCh chan qna.Answer) { +func workReceiver(ctx context.Context, opts options.Options, answerCh chan qna.Answer) { p := qna.Progress{} handleAnswer := func(a qna.Answer) { p.AddAnswer(a) @@ -167,14 +154,14 @@ func workAggregator(ctx context.Context, opts options.Options, answerCh chan qna case a := <-answerCh: handleAnswer(a) case <-time.After(opts.Timeout + time.Second): - log.Printf("INFO: metrics@aggregator: summary: \n\n%s\n\n", p.String()) + log.Printf("INFO: metrics@receiver: summary: \n\n%s\n\n", p.String()) return } } case a := <-answerCh: handleAnswer(a) case <-metrics.C: - log.Printf("INFO: metrics@aggregator: \n\n%s\n\n", p.String()) + log.Printf("INFO: metrics@receiver: \n\n%s\n\n", p.String()) } } } @@ -218,7 +205,7 @@ func workGenerator(ctx context.Context, opts options.Options, fp *os.File, quest select { case <-ctx.Done(): return - case questionCh <- qna.Question{Domain: trimWildcard(scanner.Text())}: + case questionCh <- qna.Question{Domain: qna.TrimWildcard(scanner.Text())}: nextLine.Inc() numRequests++ } @@ -254,10 +241,3 @@ func generatorMetrics(ctx context.Context, opts options.Options, nextLine *line. } } } - -func trimWildcard(san string) string { - if len(san) >= 2 && san[:2] == "*." { - return san[2:] - } - return san -} |