From ecca93a9a27ba965ef592fd496d4b504da8628b0 Mon Sep 17 00:00:00 2001
From: Rasmus Dahlberg <rasmus@rgdd.se>
Date: Sun, 2 Apr 2023 15:01:47 +0200
Subject: Minor edits and renaming

---
 main.go | 74 ++++++++++++++++++++++++-----------------------------------------
 1 file changed, 27 insertions(+), 47 deletions(-)

(limited to 'main.go')

diff --git a/main.go b/main.go
index 79dde96..726da75 100644
--- a/main.go
+++ b/main.go
@@ -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
-}
-- 
cgit v1.2.3