rkt: add --set-env to rkt run and rkt prepare

Also renamed --inherit-environment to --inherit-env, the former is too
verbose for consistent use with --set-env considering --set-env is to be
used repeatedly for multiple variables, abbreviate both uniformly.

This incarnation of --set-env applies set variables globally to all apps,
it seems desirable to be able to specify target apps for the variables.
Limiting inheritance to specific apps may also be useful.

Documentation/commands.md

@@ -166,19 +166,25 @@ For example:
 $ sudo rkt run example.com/worker -- --loglevel verbose --- example.com/syncer -- --interval 30s
 ```
 
-#### Inheriting Environment Variables
+#### Influencing Environment Variables
 
-To inherit environment variables from the parent use the `--inherit-environment` flag.
+To inherit all environment variables from the parent use the `--inherit-env` flag.
 
-The inheritance precedence is as follows with the last item replacing previous environment entries:
+To explicitly set individual environment variables use the `--set-env` flag.
+
+The precedence is as follows with the last item replacing previous environment entries:
 
 - Parent environment
 - App image environment
+- Explicitly set environment
 
 ```
 $ export EXAMPLE_ENV=hello
-$ sudo rkt run --inherit-environment example.com/env-printer
+$ export EXAMPLE_OVERRIDE=under
+$ sudo rkt run --inherit-env --set-env FOO=bar --set-env EXAMPLE_OVERRIDE=over example.com/env-printer
 EXAMPLE_ENV=hello
+FOO=bar
+EXAMPLE_OVERRIDE=over
 ```
 
 _TODO: Exit codes_

rkt/prepare.go

@@ -47,7 +47,8 @@ func init() {
 	cmdPrepare.Flags.StringVar(&flagStage1Image, "stage1-image", defaultStage1Image, `image to use as stage1. Local paths and http/https URLs are supported. If empty, Rocket will look for a file called "stage1.aci" in the same directory as rkt itself`)
 	cmdPrepare.Flags.Var(&flagVolumes, "volume", "volumes to mount into the shared container environment")
 	cmdPrepare.Flags.BoolVar(&flagQuiet, "quiet", false, "suppress superfluous output on stdout, print only the UUID on success")
-	cmdPrepare.Flags.BoolVar(&flagInheritEnv, "inherit-environment", false, "inherit all environment variables not set by apps")
+	cmdPrepare.Flags.BoolVar(&flagInheritEnv, "inherit-env", false, "inherit all environment variables not set by apps")
+	cmdPrepare.Flags.Var(&flagExplicitEnv, "set-env", "an environment variable to set for apps in the form name=value")
 }
 
 func runPrepare(args []string) (exit int) {
@@ -119,6 +120,7 @@ func runPrepare(args []string) (exit int) {
 		ExecAppends: appArgs,
 		Volumes:     []types.Volume(flagVolumes),
 		InheritEnv:  flagInheritEnv,
+		ExplicitEnv: flagExplicitEnv.Strings(),
 	}
 
 	if err = stage0.Prepare(pcfg, c.path(), c.uuid); err != nil {

rkt/run.go

@@ -38,6 +38,7 @@ var (
 	flagPrivateNet           bool
 	flagSpawnMetadataService bool
 	flagInheritEnv           bool
+	flagExplicitEnv          envMap
 	cmdRun                   = &Command{
 		Name:    "run",
 		Summary: "Run image(s) in an application container in rocket",
@@ -67,7 +68,8 @@ func init() {
 	cmdRun.Flags.Var(&flagVolumes, "volume", "volumes to mount into the shared container environment")
 	cmdRun.Flags.BoolVar(&flagPrivateNet, "private-net", false, "give container a private network")
 	cmdRun.Flags.BoolVar(&flagSpawnMetadataService, "spawn-metadata-svc", false, "launch metadata svc if not running")
-	cmdRun.Flags.BoolVar(&flagInheritEnv, "inherit-environment", false, "inherit all environment variables not set by apps")
+	cmdRun.Flags.BoolVar(&flagInheritEnv, "inherit-env", false, "inherit all environment variables not set by apps")
+	cmdRun.Flags.Var(&flagExplicitEnv, "set-env", "an environment variable to set for apps in the form name=value")
 	flagVolumes = volumeList{}
 }
 
@@ -235,6 +237,7 @@ func runRun(args []string) (exit int) {
 		ExecAppends:  appArgs,
 		Volumes:      []types.Volume(flagVolumes),
 		InheritEnv:   flagInheritEnv,
+		ExplicitEnv:  flagExplicitEnv.Strings(),
 	}
 	err = stage0.Prepare(pcfg, c.path(), c.uuid)
 	if err != nil {
@@ -287,3 +290,35 @@ func (vl *volumeList) String() string {
 	}
 	return strings.Join(vs, " ")
 }
+
+// envMap implements the flag.Value interface to contain a set of name=value mappings
+type envMap struct {
+	mapping map[string]string
+}
+
+func (e *envMap) Set(s string) error {
+	if e.mapping == nil {
+		e.mapping = make(map[string]string)
+	}
+	pair := strings.SplitN(s, "=", 2)
+	if len(pair) != 2 {
+		return fmt.Errorf("environment variable must be specified as name=value")
+	}
+	if _, exists := e.mapping[pair[0]]; exists {
+		return fmt.Errorf("environment variable %q already set", pair[0])
+	}
+	e.mapping[pair[0]] = pair[1]
+	return nil
+}
+
+func (e *envMap) String() string {
+	return strings.Join(e.Strings(), "\n")
+}
+
+func (e *envMap) Strings() []string {
+	var env []string
+	for n, v := range e.mapping {
+		env = append(env, n+"="+v)
+	}
+	return env
+}

stage0/run.go

@@ -59,6 +59,7 @@ type PrepareConfig struct {
 	ExecAppends [][]string     // appendages to each image's app.exec lines (empty when none, length should match length of Images)
 	Volumes     []types.Volume // list of volumes that rocket can provide to applications
 	InheritEnv  bool           // inherit parent environment into apps
+	ExplicitEnv []string       // always set these environment variables for all the apps
 }
 
 // configuration parameters needed by Run
@@ -80,18 +81,20 @@ func init() {
 	log.SetOutput(ioutil.Discard)
 }
 
-// MergeEnvs amends appEnv setting variables in setEnv before setting anything new from inheritEnv
-// inheritEnv and setEnv are expected to be in the os.Environ() key=value format
-func MergeEnvs(appEnv *types.Environment, inheritEnv, setEnv []string) {
+// MergeEnvs amends appEnv setting variables in setEnv before setting anything new from os.Environ if inheritEnv = true
+// setEnv is expected to be in the os.Environ() key=value format
+func MergeEnvs(appEnv *types.Environment, inheritEnv bool, setEnv []string) {
 	for _, ev := range setEnv {
 		pair := strings.SplitN(ev, "=", 2)
 		appEnv.Set(pair[0], pair[1])
 	}
 
-	for _, ev := range inheritEnv {
-		pair := strings.SplitN(ev, "=", 2)
-		if _, exists := appEnv.Get(pair[0]); !exists {
-			appEnv.Set(pair[0], pair[1])
+	if inheritEnv {
+		for _, ev := range os.Environ() {
+			pair := strings.SplitN(ev, "=", 2)
+			if _, exists := appEnv.Get(pair[0]); !exists {
+				appEnv.Set(pair[0], pair[1])
+			}
 		}
 	}
 }
@@ -146,12 +149,11 @@ func Prepare(cfg PrepareConfig, dir string, uuid *types.UUID) error {
 			a.App.Exec = append(a.App.Exec, cfg.ExecAppends[i]...)
 		}
 
-		if cfg.InheritEnv {
+		if cfg.InheritEnv || len(cfg.ExplicitEnv) > 0 {
 			if a.App == nil {
 				a.App = am.App
 			}
-			// TODO(vc): use last parameter to propagate explicit override values
-			MergeEnvs(&a.App.Environment, os.Environ(), nil)
+			MergeEnvs(&a.App.Environment, cfg.InheritEnv, cfg.ExplicitEnv)
 		}
 		cm.Apps = append(cm.Apps, a)
 	}