Refactored afew places

Makefile

@@ -1,20 +1,51 @@
-tests:
-	go test . ./... -v -race -covermode atomic -coverprofile coverage.out && go tool cover -html coverage.out -o coverage.html
+# Variables for coverage and profiling files
+COVERAGE_FILE := coverage.out
+COVERAGE_HTML := coverage.html
+CPU_PROFILE := cpu.prof
+MEM_PROFILE := mem.prof
 
-tests_without_race:
-	go test . ./... -v -covermode atomic -coverprofile coverage.out && go tool cover -html coverage.out -o coverage.html
+# Default target runs tests with race conditions
+.PHONY: all
+all: test
 
+# Test with race detection, coverage report generation
+.PHONY: test
+test:
+	go test ./... -v -race -covermode=atomic -coverprofile=$(COVERAGE_FILE)
+	go tool cover -html=$(COVERAGE_FILE) -o $(COVERAGE_HTML)
+
+# Test without race detection, coverage report generation
+.PHONY: test-no-race
+test-no-race:
+	go test ./... -v -covermode=atomic -coverprofile=$(COVERAGE_FILE)
+	go tool cover -html=$(COVERAGE_FILE) -o $(COVERAGE_HTML)
+
+# Format the code
+.PHONY: format
 format:
-	go fmt . ./...
+	go fmt ./...
 
+# Run benchmarks with memory allocation statistics
+.PHONY: bench
 bench:
-	go test . ./... -bench . -benchmem -cpu 1
+	go test ./... -bench . -benchmem -cpu=1
+
+# Generate CPU and memory profiles while running benchmarks
+.PHONY: profile-bench
+profile-bench:
+	go test ./... -bench . -cpuprofile=$(CPU_PROFILE) -memprofile=$(MEM_PROFILE) -cpu=1
 
-report_bench:
-	go test . ./... -cpuprofile cpu.prof -memprofile mem.prof -bench . -cpu 1
+# Generate CPU profiling report
+.PHONY: cpu-report
+cpu-report:
+	go tool pprof $(CPU_PROFILE)
 
-cpu_report:
-	go tool pprof cpu.prof
+# Generate memory profiling report
+.PHONY: mem-report
+mem-report:
+	go tool pprof $(MEM_PROFILE)
 
-mem_report:
-	go tool pprof mem.prof
+# Clean up profiling and coverage files
+.PHONY: clean
+clean:
+	rm -f $(COVERAGE_FILE) $(COVERAGE_HTML) $(CPU_PROFILE) $(MEM_PROFILE)

cli/cli.go

@@ -2,8 +2,8 @@ package cli
 
 import (
 	"bufio"
+	"context"
 	"errors"
-	"flag"
 	"fmt"
 	"log"
 	"os"
@@ -16,122 +16,112 @@ import (
 	"github.com/kenjoe41/goSubsWordlist/output"
 )
 
-// Cli accepts a list of URLs, one URL per line, from stdin and generates a wordlist from all subdomains found in the list.
 func Cli(includeRoot, silent bool) error {
-	// Print Header text
 	if !silent {
 		output.Beautify()
 	}
 
-	// This is a CPU-bound task, increasing the threads beyond what's available will just make it slow so removed the flag option.
 	concurrency := runtime.NumCPU()
-
-	// This is divided up in the subroutine for loop, so a value below 2 is BS.
 	if concurrency < 2 {
 		concurrency = 2
-	} else {
-		// We have 2 channels to share the concurrency with, let's reassure them that they'll have equal share.
-		concurrency *= 2
 	}
 
-	// Create channels to use
+	// Channels and WaitGroups
 	domains := make(chan string)
 	subdomains := make(chan string)
-	output := make(chan string)
-
-	// Domain Input worker
-	var domainsWG sync.WaitGroup
-	for i := 0; i < concurrency/2; i++ {
-		domainsWG.Add(1)
+	words := make(chan string)
 
-		go func() {
-			extract, err := fasttld.New(fasttld.SuffixListParams{})
-			if err != nil {
-				log.Fatal(err) // unlikely
-			}
-			for domain := range domains {
-				if domain == "" {
-					// Log something but continue to next domain if available
-					// log.Printf("Failed to get domain from: %s", domain)
-					continue
-				}
-				subdomain := ezutils.ExtractSubdomain(domain, includeRoot, extract)
-
-				if subdomain == "" {
-					// Log something but continue to next domain if available
-					// log.Printf("Failed to get subdomain for domain: %s", domain)
-					continue
-				}
-				subdomains <- subdomain
-			}
-			domainsWG.Done()
-		}()
-	}
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
-	var subdomainsWG sync.WaitGroup
+	var wg sync.WaitGroup
 
+	// Domain input workers
+	wg.Add(concurrency / 2)
 	for i := 0; i < concurrency/2; i++ {
-		subdomainsWG.Add(1)
-
-		go func() {
-			for inSubdomains := range subdomains {
-				// Split the subdomain into separate words by the '.' char.
-				// Returns slice of words.
-				subWords := strings.Split(inSubdomains, ".")
-
-				// Print to console for now
-				for _, subword := range subWords {
-					output <- subword
-				}
-			}
-			subdomainsWG.Done()
-		}()
+		go processDomains(ctx, &wg, domains, subdomains, includeRoot)
 	}
 
-	// Close subdomains channel when done reading from domains chan.
-	go func() {
-		domainsWG.Wait()
-		close(subdomains)
-	}()
+	// Subdomain processing workers
+	wg.Add(concurrency / 2)
+	for i := 0; i < concurrency/2; i++ {
+		go processSubdomains(ctx, &wg, subdomains, words)
+	}
 
-	var outputWG sync.WaitGroup
-	outputWG.Add(1)
+	// Output processor
 	go func() {
-		for word := range output {
+		for word := range words {
 			fmt.Println(word)
 		}
-		outputWG.Done()
 	}()
 
-	// Close the Output Chan after subdomain worker is done.
-	go func() {
-		subdomainsWG.Wait()
-		close(output)
-	}()
+	// Read input from stdin
+	if err := readStdin(domains); err != nil {
+		return err
+	}
+	close(domains)
 
-	// Check for stdin input
-	stat, _ := os.Stdin.Stat()
-	if (stat.Mode() & os.ModeCharDevice) != 0 {
-		flag.Usage()
-		return errors.New("No domains or urls detected. Hint: cat domains.txt | goSubsWordlist")
+	// Wait for workers to complete
+	wg.Wait()
+	close(subdomains)
+	close(words)
+
+	return nil
+}
+
+func processDomains(ctx context.Context, wg *sync.WaitGroup, domains <-chan string, subdomains chan<- string, includeRoot bool) {
+	defer wg.Done()
+	extract, err := fasttld.New(fasttld.SuffixListParams{})
+	if err != nil {
+		log.Fatal(err) // unlikely
+	}
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case domain, ok := <-domains:
+			if !ok {
+				return
+			}
+			if subdomain := ezutils.ExtractSubdomain(domain, includeRoot, extract); subdomain != "" {
+				subdomains <- subdomain
+			}
+		}
 	}
+}
 
-	sc := bufio.NewScanner(os.Stdin)
+func processSubdomains(ctx context.Context, wg *sync.WaitGroup, subdomains <-chan string, words chan<- string) {
+	defer wg.Done()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case subdomain, ok := <-subdomains:
+			if !ok {
+				return
+			}
+			subWords := strings.Split(subdomain, ".")
+			for _, word := range subWords {
+				words <- word
+			}
+		}
+	}
+}
 
-	for sc.Scan() {
-		domains <- sc.Text()
+func readStdin(domains chan<- string) error {
+	stat, _ := os.Stdin.Stat()
+	if (stat.Mode() & os.ModeCharDevice) != 0 {
+		return errors.New("No domains or URLs detected. Hint: cat domains.txt | goSubsWordlist")
 	}
 
-	// Close domains chan
-	close(domains)
+	scanner := bufio.NewScanner(os.Stdin)
+	for scanner.Scan() {
+		domains <- scanner.Text()
+	}
 
-	// check there were no errors reading stdin (unlikely)
-	if err := sc.Err(); err != nil {
+	if err := scanner.Err(); err != nil {
 		return err
 	}
 
-	// Wait until the output waitgroup is done
-	outputWG.Wait()
-
 	return nil
 }

cli/cli_test.go

@@ -5,21 +5,70 @@ import (
 	"testing"
 )
 
-func TestCLI(t *testing.T) {
+func TestCLI_NoStdin(t *testing.T) {
+	// Test when there is no input from Stdin
 	if err := Cli(false, false); err == nil {
 		t.Error("Expected error for no Stdin input")
 	}
+}
+
+func TestCLI_WithStdin(t *testing.T) {
+	// Test with simulated Stdin input
+	userInput := "example.com\nsub.example.com\nanother.sub.example.com\n\n"
+
+	funcDefer, err := mockStdin(t, userInput)
+	if err != nil {
+		t.Fatalf("Error mocking Stdin: %v", err)
+	}
+	defer funcDefer()
+
+	if err := Cli(false, false); err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+}
+
+func TestCLI_WithIncludeRoot(t *testing.T) {
+	// Test including the root domain in the wordlist
+	userInput := "example.com\nsub.example.com\n\n"
+
+	funcDefer, err := mockStdin(t, userInput)
+	if err != nil {
+		t.Fatalf("Error mocking Stdin: %v", err)
+	}
+	defer funcDefer()
+
+	if err := Cli(true, false); err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+}
 
+func TestCLI_WithSilentFlag(t *testing.T) {
+	// Test with the silent flag enabled (no banner output)
 	userInput := "example.com\nsub.example.com\n\n"
 
 	funcDefer, err := mockStdin(t, userInput)
 	if err != nil {
-		t.Errorf("%q", err)
+		t.Fatalf("Error mocking Stdin: %v", err)
+	}
+	defer funcDefer()
+
+	if err := Cli(false, true); err != nil {
+		t.Errorf("Unexpected error: %v", err)
 	}
+}
+
+// Test with invalid input
+func TestCLI_InvalidInput(t *testing.T) {
+	userInput := "invalid-url"
 
+	funcDefer, err := mockStdin(t, userInput)
+	if err != nil {
+		t.Fatalf("Error mocking Stdin: %v", err)
+	}
 	defer funcDefer()
+
 	if err := Cli(false, false); err != nil {
-		t.Errorf("%q", err)
+		t.Errorf("Unexpected error: %v", err)
 	}
 }
 
@@ -30,27 +79,25 @@ func mockStdin(t *testing.T, dummyInput string) (funcDefer func(), err error) {
 
 	oldOsStdin := os.Stdin
 	tmpfile, err := os.CreateTemp(t.TempDir(), t.Name())
-
 	if err != nil {
 		return nil, err
 	}
 
-	content := []byte(dummyInput)
-
-	if _, err := tmpfile.Write(content); err != nil {
+	if _, err := tmpfile.Write([]byte(dummyInput)); err != nil {
 		return nil, err
 	}
 
 	if _, err := tmpfile.Seek(0, 0); err != nil {
 		return nil, err
 	}
 
-	// Set stdin to the temp file
+	// Set os.Stdin to the temp file
 	os.Stdin = tmpfile
 
 	return func() {
-		// clean up
+		// Clean up
 		os.Stdin = oldOsStdin
+		tmpfile.Close()
 		os.Remove(tmpfile.Name())
 	}, nil
 }

ezutils/ezutils.go

@@ -4,14 +4,11 @@ import (
 	"github.com/elliotwutingfeng/go-fasttld"
 )
 
-// ExtractSubdomain extracts the subdomain from a given url.
-// If includeRootPtr is true, the second-level domain will be included
+// ExtractSubdomain extracts the subdomain from a given URL.
+// If includeRootPtr is true, the second-level domain will be included in the result.
 func ExtractSubdomain(url string, includeRootPtr bool, extract *fasttld.FastTLD) string {
 	result, err := extract.Extract(fasttld.URLParams{URL: url})
-	if err != nil {
-		return ""
-	}
-	if result.HostType != fasttld.HostName {
+	if err != nil || result.HostType != fasttld.HostName {
 		return ""
 	}
 	if len(result.SubDomain) > 0 {