RSDK-9651 - Add cloud metadata and api key to module env vars (viamro…

…botics#4736)

Co-authored-by: Benjamin Rewis <[email protected]>

config/config.go

@@ -965,6 +965,22 @@ func (config *AuthHandlerConfig) Validate(path string) error {
 	return nil
 }
 
+// ParseAPIKeys parses API keys from the handler config. It will return an empty map
+// if the credential type is not [rpc.CredentialsTypeAPIKey].
+func ParseAPIKeys(handler AuthHandlerConfig) map[string]string {
+	apiKeys := map[string]string{}
+	if handler.Type == rpc.CredentialsTypeAPIKey {
+		for k := range handler.Config {
+			// if it is not a legacy api key indicated by "key(s)" key
+			// current api keys will follow format { [keyId]: [key] }
+			if k != "keys" && k != "key" {
+				apiKeys[k] = handler.Config.String(k)
+			}
+		}
+	}
+	return apiKeys
+}
+
 // CreateTLSWithCert creates a tls.Config with the TLS certificate to be returned.
 func CreateTLSWithCert(cfg *Config) (*tls.Config, error) {
 	cert, err := tls.X509KeyPair([]byte(cfg.Cloud.TLSCertificate), []byte(cfg.Cloud.TLSPrivateKey))

config/module.go

@@ -130,6 +130,17 @@ func (m Module) Equals(other Module) bool {
 	return reflect.DeepEqual(m, other)
 }
 
+// MergeEnvVars will merge the provided environment variables with the existing Environment, with the existing Environment
+// taking priority.
+func (m *Module) MergeEnvVars(env map[string]string) {
+	for k, v := range env {
+		if _, ok := m.Environment[k]; ok {
+			continue
+		}
+		m.Environment[k] = v
+	}
+}
+
 var tarballExtensionsRegexp = regexp.MustCompile(`\.(tgz|tar\.gz)$`)
 
 // NeedsSyntheticPackage returns true if this is a local module pointing at a tarball.
@@ -315,6 +326,7 @@ func (m *Module) FirstRun(
 	for key, val := range env {
 		cmd.Env = append(cmd.Env, key+"="+val)
 	}
+	utils.LogViamEnvVariables("Running first run script with following Viam environment variables", env, logger)
 
 	stdOut, err := cmd.StdoutPipe()
 	if err != nil {

config/module_test.go

@@ -397,6 +397,24 @@ func TestFindMetaJSONFile(t *testing.T) {
 	})
 }
 
+func TestMergeEnvVars(t *testing.T) {
+	t.Run("empty", func(t *testing.T) {
+		m := Module{Environment: map[string]string{}}
+		expected := map[string]string{"abc": "def", "hello": "world"}
+		m.MergeEnvVars(expected)
+		test.That(t, m.Environment, test.ShouldResemble, expected)
+	})
+
+	t.Run("existing env priority", func(t *testing.T) {
+		m := Module{Environment: map[string]string{"hello": "world"}}
+		env := map[string]string{"abc": "def", "hello": "friend"}
+
+		expected := map[string]string{"abc": "def", "hello": "world"}
+		m.MergeEnvVars(env)
+		test.That(t, m.Environment, test.ShouldResemble, expected)
+	})
+}
+
 // testWriteJSON is a t.Helper that serializes `value` to `path` as json.
 func testWriteJSON(t *testing.T, path string, value any) {
 	t.Helper()

config/reader.go

@@ -11,6 +11,7 @@ import (
 	"os"
 	"path/filepath"
 	"runtime"
+	"sort"
 	"time"
 
 	"github.com/a8m/envsubst"
@@ -437,6 +438,36 @@ func processConfigLocalConfig(unprocessedConfig *Config, logger logging.Logger)
 	return processConfig(unprocessedConfig, false, logger)
 }
 
+// additionalModuleEnvVars will get additional environment variables for modules using other parts of the config.
+func additionalModuleEnvVars(cloud *Cloud, auth AuthConfig) map[string]string {
+	env := make(map[string]string)
+	if cloud != nil {
+		env[rutils.PrimaryOrgIDEnvVar] = cloud.PrimaryOrgID
+		env[rutils.LocationIDEnvVar] = cloud.LocationID
+		env[rutils.MachineIDEnvVar] = cloud.MachineID
+		env[rutils.MachinePartIDEnvVar] = cloud.ID
+	}
+	for _, handler := range auth.Handlers {
+		if handler.Type != rpc.CredentialsTypeAPIKey {
+			continue
+		}
+		apiKeys := ParseAPIKeys(handler)
+		if len(apiKeys) == 0 {
+			continue
+		}
+		// the keys come in unsorted, so sort the keys so we'll always get the same API key
+		// if there are no changes
+		keyIDs := make([]string, 0, len(apiKeys))
+		for k := range apiKeys {
+			keyIDs = append(keyIDs, k)
+		}
+		sort.Strings(keyIDs)
+		env[rutils.APIKeyIDEnvVar] = keyIDs[0]
+		env[rutils.APIKeyEnvVar] = apiKeys[keyIDs[0]]
+	}
+	return env
+}
+
 // processConfig processes the config passed in. The config can be either JSON or gRPC derived.
 // If any part of this function errors, the function will exit and no part of the new config will be returned
 // until it is corrected.
@@ -596,6 +627,13 @@ func processConfig(unprocessedConfig *Config, fromCloud bool, logger logging.Log
 		}
 	}
 
+	// add additional environment vars to modules
+	// adding them here ensures that if the parsed API key changes, the module will be restarted with the updated environment.
+	env := additionalModuleEnvVars(cfg.Cloud, cfg.Auth)
+	for _, m := range cfg.Modules {
+		m.MergeEnvVars(env)
+	}
+
 	// now that the attribute maps are converted, validate configs and get implicit dependencies for builtin resource models
 	if err := cfg.Ensure(fromCloud, logger); err != nil {
 		return nil, err

config/reader_test.go

@@ -13,9 +13,11 @@ import (
 	"github.com/pkg/errors"
 	pb "go.viam.com/api/app/v1"
 	"go.viam.com/test"
+	"go.viam.com/utils/rpc"
 
 	"go.viam.com/rdk/config/testutils"
 	"go.viam.com/rdk/logging"
+	"go.viam.com/rdk/utils"
 )
 
 func TestFromReader(t *testing.T) {
@@ -394,3 +396,97 @@ func TestReadTLSFromCache(t *testing.T) {
 		test.That(t, err, test.ShouldBeNil)
 	})
 }
+
+func TestAdditionalModuleEnvVars(t *testing.T) {
+	t.Run("empty", func(t *testing.T) {
+		expected := map[string]string{}
+		observed := additionalModuleEnvVars(nil, AuthConfig{})
+		test.That(t, observed, test.ShouldResemble, expected)
+	})
+
+	cloud1 := Cloud{
+		ID:           "test",
+		LocationID:   "the-location",
+		PrimaryOrgID: "the-primary-org",
+		MachineID:    "the-machine",
+	}
+	t.Run("cloud", func(t *testing.T) {
+		expected := map[string]string{
+			utils.MachinePartIDEnvVar: cloud1.ID,
+			utils.MachineIDEnvVar:     cloud1.MachineID,
+			utils.PrimaryOrgIDEnvVar:  cloud1.PrimaryOrgID,
+			utils.LocationIDEnvVar:    cloud1.LocationID,
+		}
+		observed := additionalModuleEnvVars(&cloud1, AuthConfig{})
+		test.That(t, observed, test.ShouldResemble, expected)
+	})
+
+	authWithExternalCreds := AuthConfig{
+		Handlers: []AuthHandlerConfig{{Type: rpc.CredentialsTypeExternal}},
+	}
+
+	t.Run("auth with external creds", func(t *testing.T) {
+		expected := map[string]string{}
+		observed := additionalModuleEnvVars(nil, authWithExternalCreds)
+		test.That(t, observed, test.ShouldResemble, expected)
+	})
+	apiKeyID := "abc"
+	apiKey := "def"
+	authWithAPIKeyCreds := AuthConfig{
+		Handlers: []AuthHandlerConfig{{Type: rpc.CredentialsTypeAPIKey, Config: utils.AttributeMap{
+			apiKeyID: apiKey,
+			"keys":   []string{apiKeyID},
+		}}},
+	}
+
+	t.Run("auth with api key creds", func(t *testing.T) {
+		expected := map[string]string{
+			utils.APIKeyEnvVar:   apiKey,
+			utils.APIKeyIDEnvVar: apiKeyID,
+		}
+		observed := additionalModuleEnvVars(nil, authWithAPIKeyCreds)
+		test.That(t, observed, test.ShouldResemble, expected)
+	})
+
+	apiKeyID2 := "uvw"
+	apiKey2 := "xyz"
+	order1 := AuthConfig{
+		Handlers: []AuthHandlerConfig{{Type: rpc.CredentialsTypeAPIKey, Config: utils.AttributeMap{
+			apiKeyID:  apiKey,
+			apiKeyID2: apiKey2,
+			"keys":    []string{apiKeyID, apiKeyID2},
+		}}},
+	}
+	order2 := AuthConfig{
+		Handlers: []AuthHandlerConfig{{Type: rpc.CredentialsTypeAPIKey, Config: utils.AttributeMap{
+			apiKeyID2: apiKey2,
+			apiKeyID:  apiKey,
+			"keys":    []string{apiKeyID, apiKeyID2},
+		}}},
+	}
+
+	t.Run("auth with keys in different order are stable", func(t *testing.T) {
+		expected := map[string]string{
+			utils.APIKeyEnvVar:   apiKey,
+			utils.APIKeyIDEnvVar: apiKeyID,
+		}
+		observed := additionalModuleEnvVars(nil, order1)
+		test.That(t, observed, test.ShouldResemble, expected)
+
+		observed = additionalModuleEnvVars(nil, order2)
+		test.That(t, observed, test.ShouldResemble, expected)
+	})
+
+	t.Run("full", func(t *testing.T) {
+		expected := map[string]string{
+			utils.MachinePartIDEnvVar: cloud1.ID,
+			utils.MachineIDEnvVar:     cloud1.MachineID,
+			utils.PrimaryOrgIDEnvVar:  cloud1.PrimaryOrgID,
+			utils.LocationIDEnvVar:    cloud1.LocationID,
+			utils.APIKeyEnvVar:        apiKey,
+			utils.APIKeyIDEnvVar:      apiKeyID,
+		}
+		observed := additionalModuleEnvVars(&cloud1, authWithAPIKeyCreds)
+		test.That(t, observed, test.ShouldResemble, expected)
+	})
+}

module/modmanager/manager.go

@@ -1198,6 +1198,8 @@ func (m *module) startProcess(
 	defer checkTicker.Stop()
 
 	m.logger.CInfow(ctx, "Starting up module", "module", m.cfg.Name)
+	rutils.LogViamEnvVariables("Starting module with following Viam environment variables", moduleEnvironment, m.logger)
+
 	ctxTimeout, cancel := context.WithTimeout(ctx, rutils.GetModuleStartupTimeout(m.logger))
 	defer cancel()
 	for {
@@ -1413,6 +1415,7 @@ func getFullEnvironment(
 		environment["VIAM_MODULE_ID"] = cfg.ModuleID
 	}
 	// Overwrite the base environment variables with the module's environment variables (if specified)
+	// VIAM_MODULE_ROOT is filled out by app.viam.com in cloud robots.
 	for key, value := range cfg.Environment {
 		environment[key] = value
 	}

robot/web/web.go

@@ -603,7 +603,7 @@ func (svc *webService) initAuthHandlers(listenerTCPAddr *net.TCPAddr, options we
 		for _, handler := range options.Auth.Handlers {
 			switch handler.Type {
 			case rpc.CredentialsTypeAPIKey:
-				apiKeys := parseAPIKeys(handler)
+				apiKeys := config.ParseAPIKeys(handler)
 
 				if len(apiKeys) == 0 {
 					return nil, errors.Errorf("%q handler requires non-empty API keys", handler.Type)
@@ -640,18 +640,6 @@ func (svc *webService) initAuthHandlers(listenerTCPAddr *net.TCPAddr, options we
 	return rpcOpts, nil
 }
 
-func parseAPIKeys(handler config.AuthHandlerConfig) map[string]string {
-	apiKeys := map[string]string{}
-	for k := range handler.Config {
-		// if it is not a legacy api key indicated by "key(s)" key
-		// current api keys will follow format { [keyId]: [key] }
-		if k != "keys" && k != "key" {
-			apiKeys[k] = handler.Config.String(k)
-		}
-	}
-	return apiKeys
-}
-
 // Register every API resource grpc service here.
 func (svc *webService) initAPIResourceCollections(ctx context.Context, mod bool) error {
 	// TODO (RSDK-144): only register necessary services

utils/env.go

@@ -5,6 +5,7 @@ import (
 	"regexp"
 	"runtime"
 	"slices"
+	"strings"
 	"time"
 
 	"go.viam.com/rdk/logging"
@@ -30,6 +31,31 @@ const (
 
 	// AndroidFilesDir is hardcoded because golang inits before our android code can override HOME var.
 	AndroidFilesDir = "/data/user/0/com.viam.rdk.fgservice/cache"
+
+	// ViamEnvVarPrefix is the prefix for all Viam-related environment variables.
+	ViamEnvVarPrefix = "VIAM_"
+
+	// APIKeyEnvVar is the environment variable which contains an API key that can be used for
+	// communications to app.viam.com.
+	//nolint:gosec
+	APIKeyEnvVar = "VIAM_API_KEY"
+
+	// APIKeyIDEnvVar is the environment variable which contains an API key ID that can be used for
+	// communications to app.viam.com.
+	//nolint:gosec
+	APIKeyIDEnvVar = "VIAM_API_KEY_ID"
+
+	// MachineIDEnvVar is the environment variable that contains the machine ID of the machine.
+	MachineIDEnvVar = "VIAM_MACHINE_ID"
+
+	// MachinePartIDEnvVar is the environment variable that contains the machine part ID of the machine.
+	MachinePartIDEnvVar = "VIAM_MACHINE_PART_ID"
+
+	// LocationIDEnvVar is the environment variable that contains the location ID of the machine.
+	LocationIDEnvVar = "VIAM_LOCATION_ID"
+
+	// PrimaryOrgIDEnvVar is the environment variable that contains the primary org ID of the machine.
+	PrimaryOrgIDEnvVar = "VIAM_PRIMARY_ORG_ID"
 )
 
 // EnvTrueValues contains strings that we interpret as boolean true in env vars.
@@ -95,3 +121,24 @@ func ViamTCPSockets() bool {
 	return runtime.GOOS == "windows" ||
 		slices.Contains(EnvTrueValues, os.Getenv("VIAM_TCP_SOCKETS"))
 }
+
+// LogViamEnvVariables logs the list of viam environment variables in [os.Environ] along with the env passed in.
+func LogViamEnvVariables(msg string, envVars map[string]string, logger logging.Logger) {
+	var env []string
+	for _, v := range os.Environ() {
+		if !strings.HasPrefix(v, ViamEnvVarPrefix) {
+			continue
+		}
+		env = append(env, v)
+	}
+	for key, val := range envVars {
+		// mask the secret
+		if key == APIKeyEnvVar {
+			val = "XXXXXXXXXX"
+		}
+		env = append(env, key+"="+val)
+	}
+	if len(env) != 0 {
+		logger.Infow(msg, "environment", env)
+	}
+}