diff --git a/test/images/agnhost/README.md b/test/images/agnhost/README.md index 9224f80614a..5a26071bab0 100644 --- a/test/images/agnhost/README.md +++ b/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. diff --git a/test/images/agnhost/VERSION b/test/images/agnhost/VERSION index e3d0696453e..6d28a11dd0e 100644 --- a/test/images/agnhost/VERSION +++ b/test/images/agnhost/VERSION @@ -1 +1 @@ -2.15 +2.16 diff --git a/test/images/agnhost/agnhost.go b/test/images/agnhost/agnhost.go index 1bdedeb6382..bbad09a51b0 100644 --- a/test/images/agnhost/agnhost.go +++ b/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) diff --git a/test/images/agnhost/netexec/netexec.go b/test/images/agnhost/netexec/netexec.go index d36764fd4ff..0a2219d99c7 100644 --- a/test/images/agnhost/netexec/netexec.go +++ b/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) {