agent: add -exit-after-auth flag

command/agent.go

@@ -61,8 +61,9 @@ type AgentCommand struct {
 
 	startedCh chan (struct{}) // for tests
 
-	flagConfigs  []string
-	flagLogLevel string
+	flagConfigs       []string
+	flagLogLevel      string
+	flagExitAfterAuth bool
 
 	flagTestVerifyOnly bool
 	flagCombineLogs    bool
@@ -115,6 +116,15 @@ func (c *AgentCommand) Flags() *FlagSets {
 			"\"trace\", \"debug\", \"info\", \"warn\", and \"err\".",
 	})
 
+	f.BoolVar(&BoolVar{
+		Name:    "exit-after-auth",
+		Target:  &c.flagExitAfterAuth,
+		Default: false,
+		Usage: "If set to true, the agent will exit with code 0 after a single " +
+			"successful auth, where success means that a token was retrieved and " +
+			"all sinks successfully wrote it",
+	})
+
 	// Internal-only flags to follow.
 	//
 	// Why hello there little source code reader! Welcome to the Vault source
@@ -223,6 +233,13 @@ func (c *AgentCommand) Run(args []string) int {
 		config.Vault = new(agentConfig.Vault)
 	}
 
+	exitAfterAuth := config.ExitAfterAuth
+	f.Visit(func(fl *flag.Flag) {
+		if fl.Name == "exit-after-auth" {
+			exitAfterAuth = c.flagExitAfterAuth
+		}
+	})
+
 	c.setStringFlag(f, config.Vault.Address, &StringVar{
 		Name:    flagNameAddress,
 		Target:  &c.flagAddress,
@@ -524,7 +541,7 @@ func (c *AgentCommand) Run(args []string) int {
 		ss := sink.NewSinkServer(&sink.SinkServerConfig{
 			Logger:        c.logger.Named("sink.server"),
 			Client:        client,
-			ExitAfterAuth: config.ExitAfterAuth,
+			ExitAfterAuth: exitAfterAuth,
 		})
 		ssDoneCh = ss.DoneCh
 
@@ -534,7 +551,7 @@ func (c *AgentCommand) Run(args []string) int {
 			LogWriter:     c.logWriter,
 			VaultConf:     config.Vault,
 			Namespace:     namespace,
-			ExitAfterAuth: config.ExitAfterAuth,
+			ExitAfterAuth: exitAfterAuth,
 		})
 		tsDoneCh = ts.DoneCh
 

command/agent_test.go

@@ -225,6 +225,16 @@ cache {
 */
 
 func TestAgent_ExitAfterAuth(t *testing.T) {
+	t.Run("via_config", func(t *testing.T) {
+		testAgentExitAfterAuth(t, false)
+	})
+
+	t.Run("via_flag", func(t *testing.T) {
+		testAgentExitAfterAuth(t, true)
+	})
+}
+
+func testAgentExitAfterAuth(t *testing.T, viaFlag bool) {
 	logger := logging.NewVaultLogger(hclog.Trace)
 	coreConfig := &vault.CoreConfig{
 		Logger: logger,
@@ -313,8 +323,13 @@ func TestAgent_ExitAfterAuth(t *testing.T) {
 		logger.Trace("wrote test jwt", "path", in)
 	}
 
+	exitAfterAuthTemplText := "exit_after_auth = true"
+	if viaFlag {
+		exitAfterAuthTemplText = ""
+	}
+
 	config := `
-exit_after_auth = true
+%s
 
 auto_auth {
         method {
@@ -340,7 +355,7 @@ auto_auth {
 }
 `
 
-	config = fmt.Sprintf(config, in, sink1, sink2)
+	config = fmt.Sprintf(config, exitAfterAuthTemplText, in, sink1, sink2)
 	if err := ioutil.WriteFile(conf, []byte(config), 0600); err != nil {
 		t.Fatal(err)
 	} else {
@@ -352,7 +367,12 @@ auto_auth {
 	ui, cmd := testAgentCommand(t, logger)
 	cmd.client = client
 
-	code := cmd.Run([]string{"-config", conf})
+	args := []string{"-config", conf}
+	if viaFlag {
+		args = append(args, "-exit-after-auth")
+	}
+
+	code := cmd.Run(args)
 	if code != 0 {
 		t.Errorf("expected %d to be %d", code, 0)
 		t.Logf("output from agent:\n%s", ui.OutputWriter.String())

sdk/go.mod

@@ -30,6 +30,7 @@ require (
 	golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3 // indirect
 	golang.org/x/sys v0.0.0-20190403152447-81d4e9dc473e
 	golang.org/x/text v0.3.1-0.20181227161524-e6919f6577db // indirect
+	google.golang.org/appengine v1.4.0 // indirect
 	google.golang.org/genproto v0.0.0-20190404172233-64821d5d2107 // indirect
 	google.golang.org/grpc v1.22.0
 )