Merge pull request #89547 from smarterclayton/graceful

netexec: Allow graceful shutdown testing from netexec

test/images/agnhost/README.md

@@ -392,8 +392,14 @@ Starts a HTTP(S) server on given port with the following endpoints:
     Acceptable values: `http`, `udp`, `sctp`.
   - `tries`: The number of times the request will be performed. Default value: `1`.
 - `/echo`: Returns the given `msg` (`/echo?msg=echoed_msg`)
-- `/exit`: Closes the server with the given code (`/exit?code=some-code`). The `code`
-  is expected to be an integer [0-127] or empty; if it is not, it will return an error message.
+- `/exit`: Closes the server with the given code and graceful shutdown. The endpoint's parameters
+	are:
+	- `code`: The exit code for the process. Default value: 0. Allows an integer [0-127].
+	- `timeout`: The amount of time to wait for connections to close before shutting down.
+		Acceptable values are golang durations. If 0 the process will exit immediately without
+		shutdown.
+	- `wait`: The amount of time to wait before starting shutdown. Acceptable values are
+	  golang durations. If 0 the process will start shutdown immediately.
 - `/healthz`: Returns `200 OK` if the server is ready, `412 Status Precondition Failed`
   otherwise. The server is considered not ready if the UDP server did not start yet or
   it exited.

test/images/agnhost/VERSION

@@ -1 +1 @@
-2.15
+2.16

test/images/agnhost/agnhost.go

@@ -49,7 +49,7 @@ import (
 )
 
 func main() {
-	rootCmd := &cobra.Command{Use: "app", Version: "2.15"}
+	rootCmd := &cobra.Command{Use: "app", Version: "2.16"}
 
 	rootCmd.AddCommand(auditproxy.CmdAuditProxy)
 	rootCmd.AddCommand(connect.CmdConnect)

test/images/agnhost/netexec/netexec.go

@@ -17,6 +17,7 @@ limitations under the License.
 package netexec
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -69,8 +70,14 @@ var CmdNetexec = &cobra.Command{
     Acceptable values: "http", "udp", "sctp".
   - "tries": The number of times the request will be performed. Default value: "1".
 - "/echo": Returns the given "msg" ("/echo?msg=echoed_msg")
-- "/exit": Closes the server with the given code ("/exit?code=some-code"). The "code"
-  is expected to be an integer [0-127] or empty; if it is not, it will return an error message.
+- "/exit": Closes the server with the given code and graceful shutdown. The endpoint's parameters
+	are:
+	- "code": The exit code for the process. Default value: 0. Allows an integer [0-127].
+	- "timeout": The amount of time to wait for connections to close before shutting down.
+		Acceptable values are golang durations. If 0 the process will exit immediately without
+		shutdown.
+	- "wait": The amount of time to wait before starting shutdown. Acceptable values are
+	  golang durations. If 0 the process will start shutdown immediately.
 - "/healthz": Returns "200 OK" if the server is ready, "412 Status Precondition Failed"
   otherwise. The server is considered not ready if the UDP server did not start yet or
   it exited.
@@ -127,25 +134,27 @@ func (a *atomicBool) get() bool {
 }
 
 func main(cmd *cobra.Command, args []string) {
+	exitCh := make(chan shutdownRequest)
+	addRoutes(exitCh)
+
 	go startUDPServer(udpPort)
 	if sctpPort != -1 {
 		go startSCTPServer(sctpPort)
 	}
 
-	addRoutes()
+	server := &http.Server{Addr: fmt.Sprintf(":%d", httpPort)}
 	if len(certFile) > 0 {
-		// only start HTTPS server if a cert is provided
-		startHTTPSServer(httpPort, certFile, privKeyFile)
+		startServer(server, exitCh, func() error { return server.ListenAndServeTLS(certFile, privKeyFile) })
 	} else {
-		startHTTPServer(httpPort)
+		startServer(server, exitCh, server.ListenAndServe)
 	}
 }
 
-func addRoutes() {
+func addRoutes(exitCh chan shutdownRequest) {
 	http.HandleFunc("/", rootHandler)
 	http.HandleFunc("/clientip", clientIPHandler)
 	http.HandleFunc("/echo", echoHandler)
-	http.HandleFunc("/exit", exitHandler)
+	http.HandleFunc("/exit", func(w http.ResponseWriter, req *http.Request) { exitHandler(w, req, exitCh) })
 	http.HandleFunc("/hostname", hostnameHandler)
 	http.HandleFunc("/shell", shellHandler)
 	http.HandleFunc("/upload", uploadHandler)
@@ -156,12 +165,23 @@ func addRoutes() {
 	http.HandleFunc("/shutdown", shutdownHandler)
 }
 
-func startHTTPSServer(httpsPort int, certFile, privKeyFile string) {
-	log.Fatal(http.ListenAndServeTLS(fmt.Sprintf(":%d", httpPort), certFile, privKeyFile, nil))
-}
+func startServer(server *http.Server, exitCh chan shutdownRequest, fn func() error) {
+	go func() {
+		re := <-exitCh
+		ctx, cancelFn := context.WithTimeout(context.Background(), re.timeout)
+		defer cancelFn()
+		err := server.Shutdown(ctx)
+		log.Printf("Graceful shutdown completed with: %v", err)
+		os.Exit(re.code)
+	}()
 
-func startHTTPServer(httpPort int) {
-	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", httpPort), nil))
+	if err := fn(); err != nil {
+		if err == http.ErrServerClosed {
+			// wait until the goroutine calls os.Exit()
+			select {}
+		}
+		log.Fatal(err)
+	}
 }
 
 func rootHandler(w http.ResponseWriter, r *http.Request) {
@@ -179,13 +199,37 @@ func clientIPHandler(w http.ResponseWriter, r *http.Request) {
 	fmt.Fprintf(w, r.RemoteAddr)
 }
 
-func exitHandler(w http.ResponseWriter, r *http.Request) {
-	log.Printf("GET /exit?code=%s", r.FormValue("code"))
-	code, err := strconv.Atoi(r.FormValue("code"))
-	if err == nil || r.FormValue("code") == "" {
+type shutdownRequest struct {
+	code    int
+	timeout time.Duration
+}
+
+func exitHandler(w http.ResponseWriter, r *http.Request, exitCh chan<- shutdownRequest) {
+	waitString := r.FormValue("wait")
+	timeoutString := r.FormValue("timeout")
+	codeString := r.FormValue("code")
+	log.Printf("GET /exit?code=%s&timeout=%s&wait=%s", codeString, timeoutString, waitString)
+	timeout, err := time.ParseDuration(timeoutString)
+	if err != nil && timeoutString != "" {
+		fmt.Fprintf(w, "argument 'timeout' must be a valid golang duration or empty, got %q\n", timeoutString)
+		return
+	}
+	wait, err := time.ParseDuration(waitString)
+	if err != nil && waitString != "" {
+		fmt.Fprintf(w, "argument 'wait' must be a valid golang duration or empty, got %q\n", waitString)
+		return
+	}
+	code, err := strconv.Atoi(codeString)
+	if err != nil && codeString != "" {
+		fmt.Fprintf(w, "argument 'code' must be an integer [0-127] or empty, got %q\n", codeString)
+		return
+	}
+	log.Printf("Will begin shutdown in %s, allowing %s for connections to close, then will exit with %d", wait, timeout, code)
+	time.Sleep(wait)
+	if timeout == 0 {
 		os.Exit(code)
 	}
-	fmt.Fprintf(w, "argument 'code' must be an integer [0-127] or empty, got %q", r.FormValue("code"))
+	exitCh <- shutdownRequest{code: code, timeout: timeout}
 }
 
 func hostnameHandler(w http.ResponseWriter, r *http.Request) {