feat: add maxURLs CLI flag

I've also added some more input sanitization.

1 files changed, 11 insertions, 1 deletions

diff --git a/main.go b/main.go
index d23b4cf..bf4450d 100644
--- a/main.go
+++ b/main.go
@@ -11,16 +11,26 @@ import (
 func main() {
 	maxConcurrency := flag.Int("c", 0, "Maximum number of concurrent queue pushes")
 	startRawURL := flag.String("url", "", "Entry-point URL")
+	maxURLs := flag.Int("max", 0, "Maximum number of URLs to collect (0 means no limit)")
+
 	flag.Parse()
 
 	if *maxConcurrency == 0 {
 		log.Fatal("Missing -c argument")
 	}
 
+	if *maxConcurrency < 1 {
+		log.Fatalf("Invalid -c argument: %d", *maxConcurrency)
+	}
+
 	if *startRawURL == "" {
 		log.Fatal("Missing -url argument")
 	}
 
+	if *maxURLs < 0 {
+		log.Fatalf("Invalid -max argument: %d", *maxURLs)
+	}
+
 	startURL, err := url.Parse(*startRawURL)
 	if err != nil {
 		log.Fatal(err)
@@ -42,7 +52,7 @@ func main() {
 loop:
 	for list := range worklist {
 		for _, u := range list {
-			if len(seen) == 30 {
+			if len(seen) == *maxURLs {
 				cancel()
 				break loop
 			}