Added the fix

Author: lobsterjerusalem

c2/cli/basic.go

@@ -7,16 +7,17 @@ import (
 	"strings"
 	"sync"
 	"time"
+	"sync/atomic"
 
 	"github.com/vulncheck-oss/go-exploit/output"
 	"github.com/vulncheck-oss/go-exploit/protocol"
 )
 
 // A very basic reverse/bind shell handler.
 func Basic(conn net.Conn) {
+	var shutdown int32 = 0
 	// Create channels for communication between goroutines.
 	responseCh := make(chan string)
-	quit := make(chan struct{})
 
 	// Use a WaitGroup to wait for goroutines to finish.
 	var wg sync.WaitGroup
@@ -25,11 +26,17 @@ func Basic(conn net.Conn) {
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
+		defer func (){
+			// Signals for both routines to stop, this should get triggered when socket is closed
+			// and causes it to fail the read
+			atomic.StoreInt32(&shutdown, 1)
+		}()
 		responseBuffer := make([]byte, 1024)
 		for {
-			select {
-			case <-quit:
+			if atomic.LoadInt32(&shutdown) == 1 {
 				return
+			}
+			select {
 			default:
 				_ = conn.SetReadDeadline(time.Now().Add(1 * time.Second))
 				bytesRead, err := conn.Read(responseBuffer)
@@ -53,32 +60,38 @@ func Basic(conn net.Conn) {
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		for response := range responseCh {
-			select {
-			case <-quit:
+		for {
+			if atomic.LoadInt32(&shutdown) == 1 {
 				return
-			default:
+			}
+			select {
+			case response := <-responseCh:
 				output.PrintShell(response)
+			default:
 			}
 		}
 	}()
 
-	for {
-		// read user input until they type 'exit\n' or the socket breaks
-		// note that ReadString is blocking, so they won't know the socket
-		// is broken until they attempt to write something
-		reader := bufio.NewReader(os.Stdin)
-		command, _ := reader.ReadString('\n')
-		ok := protocol.TCPWrite(conn, []byte(command))
-		if !ok || command == "exit\n" {
-			break
+	go func(){ 
+		// no waitgroup for this one because blocking IO, but this should not matter
+		// since we are intentionally not trying to be a multi-implant C2 framework.
+		// There still remains the issue that you would need to hit enter to find out
+		// that the socket is dead but at least we can stop Basic() regardless of this fact.
+		// This issue of unblocking stdin is discussed at length here https://github.com/golang/go/issues/24842
+		for { 
+			reader := bufio.NewReader(os.Stdin)
+			command,_ := reader.ReadString('\n')
+			if atomic.LoadInt32(&shutdown) == 1 {
+				break
+			}
+			ok := protocol.TCPWrite(conn, []byte(command))
+			if !ok || command == "exit\n" {
+				break
+			}
 		}
-	}
-
-	// signal for everyone to shutdown
-	quit <- struct{}{}
-	close(responseCh)
+	}()
 
 	// wait until the go routines are clean up
 	wg.Wait()
+	close(responseCh)
 }

c2/simpleshell/simpleshellserver.go

@@ -18,10 +18,18 @@ import (
 // The server can accept multiple connections, but the user has no way of swapping between them unless
 // the terminate the connection.
 type Server struct {
+	// The socket the server is listening on
 	Listener net.Listener
+	// Allows for us to track this from payloads, tells us if we have a shell
+	Success bool
+	// Lets us know if the server has completed Run(), you can combine this with
+	// Success to signal cleanup operations
+	Finished bool
 }
 
 var serverSingleton *Server
+var clientChan = make(chan net.Conn, 100)
+var wg sync.WaitGroup
 
 // A basic singleton interface for the c2.
 func GetServerInstance() *Server {
@@ -57,18 +65,41 @@ func (shellServer *Server) Init(channel channel.Channel) bool {
 	return true
 }
 
+func (shellServer *Server) KillServer() bool {
+	output.PrintFrameworkStatus("Received shutdown, killing server and client sockets")
+	shellServer.Listener.Close()
+	for {
+		select {
+		case conn := <-clientChan:
+			conn.Close()
+			output.PrintfFrameworkStatus("Force closed socket for: %s", conn.RemoteAddr())
+		default:
+			if len(clientChan) == 0 {
+				output.PrintFrameworkDebug("No more clients to kill:")
+
+				return true
+			}
+		}
+	}
+}
+
+
 // Listen for incoming.
 func (shellServer *Server) Run(timeout int) {
+	defer func (){
+		shellServer.Finished = true
+	}()
+
 	// mutex for user input
 	var cliLock sync.Mutex
 
 	// track if we got a shell or not
-	success := false
+	shellServer.Success = false
 
 	// terminate the server if no shells come in within timeout seconds
 	go func() {
 		time.Sleep(time.Duration(timeout) * time.Second)
-		if !success {
+		if !shellServer.Success {
 			output.PrintFrameworkError("Timeout met. Shutting down shell listener.")
 			shellServer.Listener.Close()
 		}
@@ -83,15 +114,18 @@ func (shellServer *Server) Run(timeout int) {
 				output.PrintFrameworkError(err.Error())
 			}
 
-			return
+			break
 		}
-		success = true
-		output.PrintfFrameworkSuccess("Caught new shell from %v", client.RemoteAddr())
+		shellServer.Success = true
+		wg.Add(1)
+		clientChan <- client
 		go handleSimpleConn(client, &cliLock, client.RemoteAddr())
 	}
+	wg.Wait()
 }
 
 func handleSimpleConn(conn net.Conn, cliLock *sync.Mutex, remoteAddr net.Addr) {
+	defer wg.Done()
 	// connections will stack up here. Currently that will mean a race
 	// to the next connection but we can add in attacker handling of
 	// connections latter

c2/sslshell/sslshellserver.go

@@ -38,8 +38,16 @@ type Server struct {
 	PrivateKeyFile string
 	// The file path to the user provided certificate (if provided)
 	CertificateFile string
+	// Allows for us to track this from payloads, tells us if we have a shell
+	Success bool
+	// Lets us know if the server has completed Run(), you can combine this with
+	// Success to signal cleanup operations
+	Finished bool
 }
 
+
+var clientChan = make(chan net.Conn, 100)
+var wg sync.WaitGroup
 var singleton *Server
 
 // Get a singleton instance of the sslserver c2.
@@ -61,6 +69,24 @@ func (shellServer *Server) CreateFlags() {
 	}
 }
 
+func (shellServer *Server) KillServer() bool {
+       output.PrintFrameworkStatus("Received shutdown, killing server and client sockets")
+       shellServer.Listener.Close()
+       for {
+               select {
+               case conn := <-clientChan:
+                       conn.Close()
+                       output.PrintfFrameworkStatus("Force closed socket for: %s", conn.RemoteAddr())
+               default:
+                       if len(clientChan) == 0 {
+                               output.PrintFrameworkDebug("No more clients to kill:")
+
+                               return true
+                       }
+               }
+       }
+}
+
 // Parses the user provided files or generates the certificate files and starts
 // the TLS listener on the user provided IP/port.
 func (shellServer *Server) Init(channel channel.Channel) bool {
@@ -106,16 +132,20 @@ func (shellServer *Server) Init(channel channel.Channel) bool {
 
 // Listens for incoming SSL/TLS connections spawns a reverse shell handler for each new connection.
 func (shellServer *Server) Run(timeout int) {
+	shellServer.Finished = false
+	defer func (){
+		shellServer.Finished = true
+	}()
 	// mutex for user input
 	var cliLock sync.Mutex
 
 	// track if we got a shell or not
-	success := false
+	shellServer.Success = false
 
 	// terminate the server if no shells come in within timeout seconds
 	go func() {
 		time.Sleep(time.Duration(timeout) * time.Second)
-		if !success {
+		if !shellServer.Success {
 			output.PrintFrameworkError("Timeout met. Shutting down shell listener.")
 			shellServer.Listener.Close()
 		}
@@ -130,15 +160,19 @@ func (shellServer *Server) Run(timeout int) {
 				output.PrintFrameworkError(err.Error())
 			}
 
-			return
+			break
 		}
-		success = true
+		shellServer.Success = true
 		output.PrintfFrameworkSuccess("Caught new shell from %v", client.RemoteAddr())
+		clientChan <- client
+		wg.Add(1)
 		go handleSimpleConn(client, &cliLock, client.RemoteAddr())
 	}
+	wg.Wait()
 }
 
 func handleSimpleConn(conn net.Conn, cliLock *sync.Mutex, remoteAddr net.Addr) {
+	defer wg.Done()
 	// connections will stack up here. Currently that will mean a race
 	// to the next connection but we can add in attacker handling of
 	// connections latter