Add launch policy for cmd and env vars

Signed-off-by: Jiankun Lu <jiankun@google.com>

launcher/container_runner.go

@@ -102,12 +102,20 @@ func NewRunner(ctx context.Context, cdClient *containerd.Client, token oauth2.To
 	logger.Printf("Operator Override Env Vars : %v\n", envs)
 	logger.Printf("Operator Override Cmd      : %v\n", launchSpec.Cmd)
 
-	imagelabels, err := getImageLabels(ctx, image)
+	imageLabels, err := getImageLabels(ctx, image)
 	if err != nil {
 		logger.Printf("Failed to get image OCI labels %v\n", err)
-	} else {
-		logger.Printf("Image Labels               : %v\n", imagelabels)
 	}
+	log.Printf("Image Labels               : %v\n", imageLabels)
+
+	launchPolicy, err := spec.GetLaunchPolicy(imageLabels)
+	if err != nil {
+		return nil, err
+	}
+	if err := launchPolicy.Verify(launchSpec); err != nil {
+		return nil, err
+	}
+
 	if imageConfig, err := image.Config(ctx); err != nil {
 		logger.Println(err)
 	} else {
@@ -403,14 +411,12 @@ func getImageLabels(ctx context.Context, image containerd.Image) (map[string]str
 	if err != nil {
 		return nil, err
 	}
-
 	switch ic.MediaType {
 	case v1.MediaTypeImageConfig, images.MediaTypeDockerSchema2Config:
 		p, err := content.ReadBlob(ctx, image.ContentStore(), ic)
 		if err != nil {
 			return nil, err
 		}
-
 		var ociimage v1.Image
 		if err := json.Unmarshal(p, &ociimage); err != nil {
 			return nil, err

launcher/spec/launch_policy.go

@@ -0,0 +1,70 @@
+package spec
+
+import (
+	"fmt"
+	"strconv"
+	"strings"
+)
+
+// LaunchPolicy contains policies on starting the container.
+// The policy comes from the labels of the image.
+type LaunchPolicy struct {
+	AllowedEnvOverride []string
+	AllowedCmdOverride bool
+}
+
+const (
+	envOverride = "tee.launch_policy.allow_env_override"
+	cmdOverride = "tee.launch_policy.allow_cmd_override"
+)
+
+// GetLaunchPolicy takes in a map[string] string which should come from image labels,
+// and will try to parse it into a LaunchPolicy. Extra fields will be ignored.
+func GetLaunchPolicy(imageLabels map[string]string) (LaunchPolicy, error) {
+	var err error
+	launchPolicy := LaunchPolicy{}
+	if v, ok := imageLabels[envOverride]; ok {
+		envs := strings.Split(v, ",")
+		for _, env := range envs {
+			// strip out empty env name
+			if env != "" {
+				launchPolicy.AllowedEnvOverride = append(launchPolicy.AllowedEnvOverride, env)
+			}
+		}
+	}
+
+	if v, ok := imageLabels[cmdOverride]; ok {
+		if launchPolicy.AllowedCmdOverride, err = strconv.ParseBool(v); err != nil {
+			return LaunchPolicy{}, fmt.Errorf("value of LABEL %s of the image is not a boolean %s", cmdOverride, v)
+		}
+	} else {
+		// default is false
+		launchPolicy.AllowedCmdOverride = false
+	}
+
+	return launchPolicy, nil
+}
+
+// Verify will use the LaunchPolicy to verify the given LauncherSpec. If the verification passed, will return nil.
+// If there is any issue, the function will return the error based on the first error.
+func (p LaunchPolicy) Verify(lp LauncherSpec) error {
+	for _, e := range lp.Envs {
+		if !contains(p.AllowedEnvOverride, e.Name) {
+			return fmt.Errorf("env var %s is not allowed to be overridden on this image; allowed envs to be overridden: %v", e, p.AllowedEnvOverride)
+		}
+	}
+	if !p.AllowedCmdOverride && len(lp.Cmd) > 0 {
+		return fmt.Errorf("CMD is not allowed to be overridden on this image")
+	}
+
+	return nil
+}
+
+func contains(strs []string, target string) bool {
+	for _, s := range strs {
+		if s == target {
+			return true
+		}
+	}
+	return false
+}

launcher/spec/launch_policy_test.go

@@ -0,0 +1,143 @@
+package spec
+
+import (
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+)
+
+func TestLauncPolicy(t *testing.T) {
+	testCases := []struct {
+		testName       string
+		imageLables    map[string]string
+		expectedPolicy LaunchPolicy
+	}{
+		{
+			"case 1",
+			map[string]string{
+				envOverride: "foo",
+				cmdOverride: "true",
+			},
+			LaunchPolicy{
+				AllowedEnvOverride: []string{"foo"},
+				AllowedCmdOverride: true,
+			},
+		},
+		{
+			"case 2",
+			map[string]string{
+				envOverride: "foo,bar",
+			},
+			LaunchPolicy{
+				AllowedEnvOverride: []string{"foo", "bar"},
+				AllowedCmdOverride: false,
+			},
+		},
+		{
+			"case 3 default case",
+			nil,
+			LaunchPolicy{
+				AllowedEnvOverride: nil,
+				AllowedCmdOverride: false,
+			},
+		},
+		{
+			"case 4 empty string in env",
+			map[string]string{
+				envOverride: ",,,foo",
+				cmdOverride: "false",
+			},
+			LaunchPolicy{
+				AllowedEnvOverride: []string{"foo"},
+				AllowedCmdOverride: false,
+			},
+		},
+	}
+
+	for _, testcase := range testCases {
+		t.Run(testcase.testName, func(t *testing.T) {
+			got, err := GetLaunchPolicy(testcase.imageLables)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			if !cmp.Equal(got, testcase.expectedPolicy) {
+				t.Errorf("Launchspec got %+v, want %+v", got, testcase.expectedPolicy)
+			}
+		})
+	}
+}
+
+func TestVerify(t *testing.T) {
+	testCases := []struct {
+		testName  string
+		policy    LaunchPolicy
+		spec      LauncherSpec
+		expectErr bool
+	}{
+		{
+			"allow everything",
+			LaunchPolicy{
+				AllowedEnvOverride: []string{"foo"},
+				AllowedCmdOverride: true,
+			},
+			LauncherSpec{
+				Envs: []EnvVar{{Name: "foo", Value: "foo"}},
+				Cmd:  []string{"foo"},
+			},
+			false,
+		},
+		{
+			"default case",
+			LaunchPolicy{},
+			LauncherSpec{},
+			false,
+		},
+		{
+			"env override violation",
+			LaunchPolicy{
+				AllowedEnvOverride: []string{"foo"},
+			},
+			LauncherSpec{
+				Envs: []EnvVar{{Name: "bar", Value: ""}},
+			},
+			true,
+		},
+		{
+			"cmd violation",
+			LaunchPolicy{
+				AllowedCmdOverride: false,
+			},
+			LauncherSpec{
+				Cmd: []string{"foo"},
+			},
+			true,
+		},
+		{
+			"allow everything",
+			LaunchPolicy{
+				AllowedEnvOverride: []string{"foo"},
+				AllowedCmdOverride: true,
+			},
+			LauncherSpec{
+				Envs: []EnvVar{{Name: "foo", Value: "foo"}},
+				Cmd:  []string{"foo"},
+			},
+			false,
+		},
+	}
+	for _, testCase := range testCases {
+		t.Run(testCase.testName, func(t *testing.T) {
+			err := testCase.policy.Verify(testCase.spec)
+			if testCase.expectErr {
+				if err == nil {
+					t.Errorf("expected error, but got nil")
+				}
+			} else {
+				if err != nil {
+					t.Errorf("expected no error, but got %v", err)
+				}
+			}
+		})
+	}
+}

launcher/spec/launcher_spec.go

@@ -33,8 +33,8 @@ const (
 	restartPolicyKey           = "tee-restart-policy"
 	cmdKey                     = "tee-cmd"
 	envKeyPrefix               = "tee-env-"
-	instanceAttributes         = "instance/attributes/?recursive=true"
 	impersonateServiceAccounts = "tee-impersonate-service-accounts"
+	instanceAttributesQuery    = "instance/attributes/?recursive=true"
 )
 
 var errImageRefNotSpecified = fmt.Errorf("%s is not specified in the custom metadata", imageRefKey)
@@ -97,6 +97,7 @@ func (s *LauncherSpec) UnmarshalJSON(b []byte) error {
 			s.Envs = append(s.Envs, EnvVar{strings.TrimPrefix(k, envKeyPrefix), v})
 		}
 	}
+
 	return nil
 }
 
@@ -105,7 +106,7 @@ func (s *LauncherSpec) UnmarshalJSON(b []byte) error {
 // ImageRef (tee-image-reference) is required, will return an error if
 // ImageRef is not presented in the metadata.
 func GetLauncherSpec(client *metadata.Client) (LauncherSpec, error) {
-	data, err := client.Get(instanceAttributes)
+	data, err := client.Get(instanceAttributesQuery)
 	if err != nil {
 		return LauncherSpec{}, err
 	}

launcher/spec/launcher_spec_test.go

@@ -18,7 +18,9 @@ func TestLauncherSpecUnmarshalJSONHappyCases(t *testing.T) {
 				"tee-env-foo":"bar",
 				"tee-image-reference":"docker.io/library/hello-world:latest",
 				"tee-restart-policy":"Always",
-				"tee-impersonate-service-accounts":"sv1@developer.gserviceaccount.com,sv2@developer.gserviceaccount.com"
+				"tee-impersonate-service-accounts":"sv1@developer.gserviceaccount.com,sv2@developer.gserviceaccount.com",
+				"tee-stderr-echo":"false",
+				"tee-stdout-echo":"false"
 			}`,
 		},
 		{