Merge pull request #4 from octoman90/workerPool

Author: man90es

README.md

@@ -37,10 +37,11 @@ cat addresses.txt | ./proxyshiva
 
 ## Flags
 
-| Flag         | Description                                                  |
-| ------------ | ------------------------------------------------------------ |
-| -json        | Output full data in JSON format                              |
-| -interactive | Don't exit after completing the task and wait for more input |
-| -skipcert    | Skip the TLS certificate verification                        |
-| -skipres     | Skip reserved IP addresses                                   |
-| -timeout=15  | Request timeout in seconds (15 by default)                   |
+| Flag          | Description                                                  |
+| ------------- | ------------------------------------------------------------ |
+| -json         | Output full data in JSON format                              |
+| -interactive  | Don't exit after completing the task and wait for more input |
+| -skipcert     | Skip the TLS certificate verification                        |
+| -skipres      | Skip reserved IP addresses                                   |
+| -parallel=100 | How many requests to make simultaneously (100 by default)    |
+| -timeout=15   | Request timeout in seconds (15 by default)                   |

inputParser/inputParser.go

@@ -1,16 +1,47 @@
 package inputParser
 
 import (
+	"regexp"
 	"strconv"
 	"strings"
 
 	"github.com/octoman90/proxyshiva/proxy"
 	"inet.af/netaddr"
 )
 
+func validateScanned(str string) bool {
+	ipRegex := `((?:(?:25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}(?:25[0-5]|2[0-4]\d|[01]?\d\d?))`
+	protocolRegex := `(http|https|socks4|socks5)`
+	portRegex := `\d+`
+
+	regex := regexp.MustCompile(`` +
+		// Match protocol or protocol list
+		`^(` + protocolRegex + `(,` + protocolRegex + `)*)+` +
+
+		// Match "://"
+		`:\/\/` +
+
+		// Match IP or IP range
+		ipRegex + `(-` + ipRegex + `)?` +
+
+		// Match ":"
+		`:` +
+
+		// Match port or port range
+		`(` + portRegex + `)(-` + portRegex + `)?$`,
+	)
+
+	return regex.MatchString(str)
+}
+
 func RequestGenerator(in string) chan proxy.Proxy {
 	out := make(chan proxy.Proxy)
 
+	if !validateScanned(in) {
+		defer close(out)
+		return out
+	}
+
 	schemeEndIndex := strings.Index(in, "://")
 	addressStartIndex := schemeEndIndex + 3
 	addressEndIndex := strings.LastIndex(in, ":")

main.go

@@ -6,7 +6,6 @@ import (
 	"flag"
 	"fmt"
 	"os"
-	"regexp"
 	"sync"
 
 	"github.com/octoman90/proxyshiva/inputParser"
@@ -15,69 +14,58 @@ import (
 
 var wg sync.WaitGroup
 
-func validateScanned(str string) bool {
-	ipRegex := `((?:(?:25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}(?:25[0-5]|2[0-4]\d|[01]?\d\d?))`
-	protocolRegex := `(http|https|socks4|socks5)`
-	portRegex := `\d+`
-
-	regex := regexp.MustCompile(`` +
-		// Match protocol or protocol list
-		`^(` + protocolRegex + `(,` + protocolRegex + `)*)+` +
-
-		// Match "://"
-		`:\/\/` +
-
-		// Match IP or IP range
-		ipRegex + `(-` + ipRegex + `)?` +
-
-		// Match ":"
-		`:` +
+func worker(taskQueue <-chan proxy.Proxy, resultQueue chan<- *proxy.Proxy, timeout *int, skipCert *bool) {
+	for proxy := range taskQueue {
+		proxy.Check(resultQueue, timeout, skipCert)
+	}
+}
 
-		// Match port or port range
-		`(` + portRegex + `)(-` + portRegex + `)?$`,
-	)
+func printer(resultQueue <-chan *proxy.Proxy, JSON *bool) {
+	for result := range resultQueue {
+		if *JSON { // Print out every result in JSON format
+			jr, _ := json.Marshal(*result)
+			fmt.Println(string(jr))
+		} else if result.Good { // Print out good proxies in short format
+			fmt.Printf("%v://%v:%v\n", result.Scheme, result.Address, result.Port)
+		}
 
-	return regex.MatchString(str)
+		wg.Done()
+	}
 }
 
 func main() {
+	// Parse flags
 	flagJSON := flag.Bool("json", false, "Output full data in JSON format")
 	flagInteractive := flag.Bool("interactive", false, "Don't exit after completing the task and wait for more input")
 	flagSkipCert := flag.Bool("skipcert", false, "Skip the TLS certificate verification")
 	flagTimeout := flag.Int("timeout", 15, "Request timeout in seconds")
 	flagSkipRes := flag.Bool("skipres", false, "Skip reserved IP addresses")
+	flagParallel := flag.Int("parallel", 100, "How many requests to make simultaneously")
 	flag.Parse()
 
-	resultQueue := make(chan *proxy.Proxy)
+	// Create communication queues
+	taskQueue := make(chan proxy.Proxy, 100)
+	resultQueue := make(chan *proxy.Proxy, 100)
+	defer close(taskQueue)
 	defer close(resultQueue)
 
-	// Receive and print out completed checks
-	go func() {
-		for result := range resultQueue {
-			if *flagJSON { // Print out every result in JSON format
-				jr, _ := json.Marshal(*result)
-				fmt.Println(string(jr))
-			} else if result.Good { // Print out good proxies in short format
-				fmt.Printf("%v://%v:%v\n", result.Scheme, result.Address, result.Port)
-			}
+	// Spawn a printer goroutine
+	go printer(resultQueue, flagJSON)
 
-			wg.Done()
-		}
-	}()
+	// Spawn worker goroutines
+	for i := 0; i < *flagParallel; i++ {
+		go worker(taskQueue, resultQueue, flagTimeout, flagSkipCert)
+	}
 
 	// Scan for input
 	scanner := bufio.NewScanner(os.Stdin)
 	for {
 		if scanner.Scan() {
-			scanned := scanner.Text()
-			if valid := validateScanned(scanned); !valid {
-				continue
-			}
-
-			for proxy := range inputParser.RequestGenerator(scanned) {
+			for proxy := range inputParser.RequestGenerator(scanner.Text()) {
+				// Copy tasks from individual input channels to a single task queue
 				if !*flagSkipRes || !proxy.IsReserved() {
 					wg.Add(1)
-					go proxy.Check(resultQueue, flagTimeout, flagSkipCert)
+					taskQueue <- proxy
 				}
 			}
 		}