From 962a5980030f2d773444da7ea7aaf558d55c0142 Mon Sep 17 00:00:00 2001
From: Justin Head <jhead@jxh.io>
Date: Tue, 14 Jul 2020 21:50:06 -0700
Subject: [PATCH] Add workers flag

---
 cmd/phantom.go          |  2 ++
 internal/proxy/proxy.go | 17 ++++++++++++++++-
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/cmd/phantom.go b/cmd/phantom.go
index 6867fb1..87e6450 100644
--- a/cmd/phantom.go
+++ b/cmd/phantom.go
@@ -27,6 +27,7 @@ func main() {
 	debugArg := flag.Bool("debug", false, "Optional: Enables debug logging")
 	ipv6Arg := flag.Bool("6", false, "Optional: Enables IPv6 support on port 19133 (experimental)")
 	removePortsArg := flag.Bool("remove_ports", false, "Optional: Forces ports to be excluded from pong packets (experimental)")
+	workersArg := flag.Uint("workers", 1, "Optional: Number of workers, useful for tweaking performance (experimental)")
 
 	flag.Usage = usage
 	flag.Parse()
@@ -66,6 +67,7 @@ func main() {
 		idleTimeout,
 		*ipv6Arg,
 		*removePortsArg,
+		*workersArg,
 	})
 
 	if err != nil {
diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go
index 9d9e2de..7eee932 100644
--- a/internal/proxy/proxy.go
+++ b/internal/proxy/proxy.go
@@ -39,6 +39,7 @@ type ProxyPrefs struct {
 	IdleTimeout  time.Duration
 	EnableIPv6   bool
 	RemovePorts  bool
+	NumWorkers   uint
 }
 
 var randSource = rand.NewSource(time.Now().UnixNano())
@@ -125,7 +126,7 @@ func (proxy *ProxyServer) Start() error {
 	log.Info().Msgf("Once your console pings phantom, you should see replies below.")
 
 	// Start processing everything else using the proxy listener
-	proxy.readLoop(proxy.server)
+	proxy.startWorkers(proxy.server)
 
 	return nil
 }
@@ -148,9 +149,23 @@ func (proxy *ProxyServer) Close() {
 	proxy.dead.Set()
 }
 
+func (proxy *ProxyServer) startWorkers(listener net.PacketConn) {
+	log.Info().Msgf("Starting %d workers", proxy.prefs.NumWorkers)
+
+	for i := uint(0); i < proxy.prefs.NumWorkers; i++ {
+		if i < proxy.prefs.NumWorkers-1 {
+			go proxy.readLoop(listener)
+		} else {
+			proxy.readLoop(listener)
+		}
+	}
+}
+
 // Continually reads data from the provided listener and passes it to
 // processDataFromClients until the ProxyServer has been closed.
 func (proxy *ProxyServer) readLoop(listener net.PacketConn) {
+	log.Info().Msgf("Listener starting up: %s", listener.LocalAddr())
+
 	packetBuffer := make([]byte, maxMTU)
 
 	for !proxy.dead.IsSet() {