From 738051a2781b8ec39d39274c4e59e7961646cb7d Mon Sep 17 00:00:00 2001
From: terrorbyte <terrorbyte@users.noreply.github.com>
Date: Fri, 8 Nov 2024 16:37:40 -0700
Subject: [PATCH] Add clean shutdown and signal handlers to HTTPServeFile

---
 c2/httpservefile/httpservefile.go | 59 +++++++++++++++++++++++++------
 1 file changed, 48 insertions(+), 11 deletions(-)

diff --git a/c2/httpservefile/httpservefile.go b/c2/httpservefile/httpservefile.go
index a0b16c5..2fe4bf8 100644
--- a/c2/httpservefile/httpservefile.go
+++ b/c2/httpservefile/httpservefile.go
@@ -28,6 +28,7 @@ import (
 	"fmt"
 	"net/http"
 	"os"
+	"os/signal"
 	"path"
 	"strings"
 	"time"
@@ -66,6 +67,8 @@ type Server struct {
 	HostedFiles map[string]HostedFile // RealName -> struct
 	// A comma delimited list of all the files to serve
 	FilesToServe string
+
+	done chan bool
 }
 
 var singleton *Server
@@ -180,6 +183,13 @@ func (httpServer *Server) AddFile(realName string, randomName string, data []byt
 	httpServer.HostedFiles[realName] = hostMe
 }
 
+// Shuts down the server cleanly. Any currently connected clients will still have to complete their
+// connections or timeouts.
+func (httpServer *Server) Shutdown() {
+	output.PrintfFrameworkDebug("HTTP server manually invoked shutdown")
+	httpServer.done <- true
+}
+
 // start the HTTP server and listen for incoming requests for `httpServer.FileName`.
 func (httpServer *Server) Run(timeout int) {
 	if len(httpServer.HostedFiles) == 0 {
@@ -210,6 +220,42 @@ func (httpServer *Server) Run(timeout int) {
 	}
 
 	connectionString := fmt.Sprintf("%s:%d", httpServer.HTTPAddr, httpServer.HTTPPort)
+
+	server := http.Server{}
+	httpServer.done = make(chan bool, 1)
+	sigint := make(chan os.Signal, 1)
+	signal.Notify(sigint, os.Interrupt)
+	// Handle the signal interrupt channel. If the signal is triggered, then trigger the done
+	// channel which will clean up the server and close cleanly.
+	go func(sigint <-chan os.Signal, done chan<- bool) {
+		<-sigint
+		done <- true
+		output.PrintfFrameworkStatus("Interrupt signal received")
+		close(done)
+	}(sigint, httpServer.done)
+
+	// Utilize the timeout config option to check to see if the timeout is met or if the
+	// done channel is triggered.
+	go func(donerecv <-chan bool, donesend chan<- bool) {
+		maxTime := time.Now().Add(time.Duration(timeout) * time.Second)
+		for tick := range time.Tick(1 * time.Second) {
+			select {
+			// If the server is marked as done, exit before the timeout
+			case <-donerecv:
+				output.PrintfFrameworkStatus("HTTP server shutting down...")
+				server.Close()
+
+				return
+			default:
+				if tick.After(maxTime) {
+					output.PrintFrameworkStatus("Timeout reached")
+					donesend <- true
+				}
+			}
+		}
+	}(httpServer.done, httpServer.done)
+
+	defer server.Close()
 	go func() {
 		if httpServer.TLS {
 			output.PrintfFrameworkStatus("Starting an HTTPS server on %s", connectionString)
@@ -218,23 +264,14 @@ func (httpServer *Server) Run(timeout int) {
 				// We have no control over the SSL versions supported on the remote target. Be permissive for more targets.
 				MinVersion: tls.VersionSSL30,
 			}
-			server := http.Server{
-				Addr:      connectionString,
-				TLSConfig: tlsConfig,
-			}
-			defer server.Close()
+			server.Addr = connectionString
+			server.TLSConfig = tlsConfig
 			_ = server.ListenAndServeTLS("", "")
 		} else {
 			output.PrintfFrameworkStatus("Starting an HTTP server on %s", connectionString)
 			_ = http.ListenAndServe(connectionString, nil)
 		}
 	}()
-
-	// let the server run for timeout seconds
-	time.Sleep(time.Duration(timeout) * time.Second)
-
-	// We don't actually clean up anything, but exiting c2 will eventually terminate the program
-	output.PrintFrameworkStatus("Shutting down the HTTP Server")
 }
 
 // Returns the random name of the provided filename. If filename is empty, return the first entry.