151 lines
3.8 KiB
Go
151 lines
3.8 KiB
Go
package main
|
|
|
|
// This is a test application that is used by TestExecServer_Run to verify
|
|
// the behavior of vault agent running as a process supervisor.
|
|
//
|
|
// The app will automatically exit after 1 minute or the --stop-after interval,
|
|
// whichever comes first. It also can serve its loaded environment variables on
|
|
// the given --port. This app will also return the given --exit-code and
|
|
// terminate on SIGTERM unless --use-sigusr1 is specified.
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"flag"
|
|
"fmt"
|
|
"log"
|
|
"net/http"
|
|
"os"
|
|
"os/signal"
|
|
"strings"
|
|
"syscall"
|
|
"time"
|
|
)
|
|
|
|
var (
|
|
port uint
|
|
ignoreStopSignal bool
|
|
sleepAfterStopSignal time.Duration
|
|
useSigusr1StopSignal bool
|
|
stopAfter time.Duration
|
|
exitCode int
|
|
)
|
|
|
|
func init() {
|
|
flag.UintVar(&port, "port", 34000, "port to run the test app on")
|
|
flag.DurationVar(&sleepAfterStopSignal, "sleep-after-stop-signal", 1*time.Second, "time to sleep after getting the signal before exiting")
|
|
flag.BoolVar(&useSigusr1StopSignal, "use-sigusr1", false, "use SIGUSR1 as the stop signal, instead of the default SIGTERM")
|
|
flag.DurationVar(&stopAfter, "stop-after", 0, "stop the process after duration (overrides all other flags if set)")
|
|
flag.IntVar(&exitCode, "exit-code", 0, "exit code to return when this script exits")
|
|
}
|
|
|
|
type Response struct {
|
|
EnvironmentVariables map[string]string `json:"environment_variables"`
|
|
ProcessID int `json:"process_id"`
|
|
}
|
|
|
|
func newResponse() Response {
|
|
respEnv := make(map[string]string, len(os.Environ()))
|
|
for _, envVar := range os.Environ() {
|
|
tokens := strings.Split(envVar, "=")
|
|
respEnv[tokens[0]] = tokens[1]
|
|
}
|
|
|
|
return Response{
|
|
EnvironmentVariables: respEnv,
|
|
ProcessID: os.Getpid(),
|
|
}
|
|
}
|
|
|
|
func handler(w http.ResponseWriter, r *http.Request) {
|
|
var buf bytes.Buffer
|
|
encoder := json.NewEncoder(&buf)
|
|
if r.URL.Query().Get("pretty") == "1" {
|
|
encoder.SetIndent("", " ")
|
|
}
|
|
if err := encoder.Encode(newResponse()); err != nil {
|
|
http.Error(w, err.Error(), http.StatusInternalServerError)
|
|
return
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusOK)
|
|
_, _ = w.Write(buf.Bytes())
|
|
}
|
|
|
|
func main() {
|
|
logger := log.New(os.Stderr, "test-app: ", log.LstdFlags)
|
|
|
|
if err := run(logger); err != nil {
|
|
log.Fatalf("error: %v\n", err)
|
|
}
|
|
|
|
logger.Printf("exit code: %d\n", exitCode)
|
|
|
|
os.Exit(exitCode)
|
|
}
|
|
|
|
func run(logger *log.Logger) error {
|
|
/* */ logger.Println("run: started")
|
|
defer logger.Println("run: done")
|
|
|
|
ctx, cancelContextFunc := context.WithTimeout(context.Background(), 60*time.Second)
|
|
defer cancelContextFunc()
|
|
|
|
flag.Parse()
|
|
|
|
server := http.Server{
|
|
Addr: fmt.Sprintf(":%d", port),
|
|
Handler: http.HandlerFunc(handler),
|
|
ReadTimeout: 20 * time.Second,
|
|
WriteTimeout: 20 * time.Second,
|
|
IdleTimeout: 20 * time.Second,
|
|
}
|
|
|
|
doneCh := make(chan struct{})
|
|
|
|
go func() {
|
|
defer close(doneCh)
|
|
|
|
stopSignal := make(chan os.Signal, 1)
|
|
if useSigusr1StopSignal {
|
|
signal.Notify(stopSignal, syscall.SIGUSR1)
|
|
} else {
|
|
signal.Notify(stopSignal, syscall.SIGTERM)
|
|
}
|
|
|
|
select {
|
|
case <-ctx.Done():
|
|
logger.Println("context done: exiting")
|
|
|
|
case s := <-stopSignal:
|
|
logger.Printf("signal %q: received\n", s)
|
|
|
|
if sleepAfterStopSignal > 0 {
|
|
logger.Printf("signal %q: sleeping for %v simulate cleanup\n", s, sleepAfterStopSignal)
|
|
time.Sleep(sleepAfterStopSignal)
|
|
}
|
|
|
|
case <-time.After(stopAfter):
|
|
logger.Printf("stopping after: %v\n", stopAfter)
|
|
}
|
|
|
|
if err := server.Shutdown(context.Background()); err != nil {
|
|
log.Printf("server shutdown error: %v", err)
|
|
}
|
|
}()
|
|
|
|
logger.Printf("server %s: started\n", server.Addr)
|
|
|
|
if err := server.ListenAndServe(); !errors.Is(err, http.ErrServerClosed) {
|
|
return fmt.Errorf("could not start the server: %v", err)
|
|
}
|
|
|
|
logger.Printf("server %s: done\n", server.Addr)
|
|
|
|
<-doneCh
|
|
|
|
return nil
|
|
}
|