2 files changed, 13 insertions, 2 deletions
diff --git a/classic.go b/classic.go
index deb14fd..14bb888 100644
--- a/classic.go
+++ b/classic.go
@@ -40,7 +40,13 @@ loop:
 					break loop
 				}
 
-				if p.depth == maxDepth {
+				// Track maxDepth here. A maxDepth
+				// greater than zero means a finite
+				// maxDepth value. If the packets
+				// we're seeing reach that depth,
+				// don't use their URLs to spawn new
+				// fetches.
+				if maxDepth > 0 && p.depth == maxDepth {
 					continue
 				}
 
diff --git a/main.go b/main.go
index 5673441..311cece 100644
--- a/main.go
+++ b/main.go
@@ -16,6 +16,7 @@ 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 (omitted or 0 means no limit)")
+	maxDepth := flag.Int("depth", 0, "Maximum URL depth (omitted or 0 means no limit)")
 
 	flag.Parse()
 
@@ -35,13 +36,17 @@ func main() {
 		log.Fatalf("Invalid -max argument: %d", *maxURLs)
 	}
 
+	if *maxDepth < 0 {
+		log.Fatalf("Invalid -depth argument: %d", *maxDepth)
+	}
+
 	startURL, err := url.Parse(*startRawURL)
 	if err != nil {
 		log.Fatal(err)
 	}
 
 	getLeakProfile(func() {
-		classic(*startURL, *maxConcurrency, *maxURLs, 1)
+		classic(*startURL, *maxConcurrency, *maxURLs, *maxDepth)
 	})
 }