Add launch policy for cmd and env vars

cloudbuild.yaml

@@ -31,8 +31,8 @@ steps:
     args: ['finish-image-build',
            '-oem-size=500M',
            '-disk-size-gb=11',
-           '-image-name=attest-cos-dev-${_OUTPUT_IMAGE}',
-           '-image-family=attest-cos-dev',
+           '-image-name=confidential-space-${_OUTPUT_IMAGE}',
+           '-image-family=confidential-space-dev',
            '-image-project=${PROJECT_ID}',
            '-zone=us-central1-a',
            '-project=${PROJECT_ID}']

launcher/container-runner.service

@@ -6,7 +6,8 @@ After=network-online.target gcr-online.target containerd.service
 [Service]
 ExecStart=/var/lib/google/cc_container_launcher --addr=${ATTEST_ENDPOINT}
 # Shutdown the host after the launcher exits
-ExecStopPost=systemctl poweroff
+ExecStopPost=/bin/sleep 60
+ExecStopPost=/usr/bin/systemctl poweroff
 Restart=no
 # RestartSec=90
 StandardOutput=journal+console

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)
 	}
+
+	logger.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 {
@@ -337,10 +345,8 @@ func (r *ContainerRunner) Run(ctx context.Context) error {
 		return fmt.Errorf("failed to fetch and write OIDC token: %v", err)
 	}
 
-	loggerAndStdout := io.MultiWriter(os.Stdout, r.logger.Writer())
-	loggerAndStderr := io.MultiWriter(os.Stderr, r.logger.Writer())
 	for {
-		task, err := r.container.NewTask(ctx, cio.NewCreator(cio.WithStreams(nil, loggerAndStdout, loggerAndStderr)))
+		task, err := r.container.NewTask(ctx, cio.NewCreator(cio.WithStreams(nil, r.logger.Writer(), r.logger.Writer())))
 		if err != nil {
 			return err
 		}
@@ -403,14 +409,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/main.go

@@ -4,6 +4,7 @@ package main
 import (
 	"context"
 	"flag"
+	"io"
 	"log"
 	"os"
 
@@ -21,6 +22,10 @@ var (
 	serverAddr    = flag.String("addr", "", "The server address in the format of host:port")
 )
 
+const (
+	logName = "confidential-space-launcher"
+)
+
 func main() {
 	flag.Parse()
 	os.Exit(run())
@@ -39,11 +44,13 @@ func run() int {
 	}
 	logClient, err := logging.NewClient(context.Background(), projectID)
 	if err != nil {
-		log.Printf("cannot setup Cloud Logging, using the default logger %v", err)
+		logger.Printf("cannot setup Cloud Logging, using the default stdout logger %v", err)
 	} else {
 		defer logClient.Close()
-		logger.Println("logs will publish to Cloud Logging")
-		logger = logClient.Logger("confidential-space-launcher").StandardLogger(logging.Info)
+		logger.Printf("logs will be published to Cloud Logging under the log name %s\n", logName)
+		logger = logClient.Logger(logName).StandardLogger(logging.Info)
+		loggerAndStdout := io.MultiWriter(os.Stdout, logger.Writer()) // for now also print log to stdout
+		logger.SetOutput(loggerAndStdout)
 	}
 
 	spec, err := spec.GetLauncherSpec(mdsClient)

launcher/spec/launch_policy.go

@@ -0,0 +1,67 @@
+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)
+		}
+	}
+
+	return launchPolicy, nil
+}
+
+// Verify will use the LaunchPolicy to verify the given LauncherSpec. If the verification passed, will return nil.
+// If there are multiple violations, the function will return the first error.
+func (p LaunchPolicy) Verify(ls LauncherSpec) error {
+	for _, e := range ls.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(ls.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 TestLaunchPolicy(t *testing.T) {
+	testCases := []struct {
+		testName       string
+		imageLables    map[string]string
+		expectedPolicy LaunchPolicy
+	}{
+		{
+			"single ENV override, CMD override",
+			map[string]string{
+				envOverride: "foo",
+				cmdOverride: "true",
+			},
+			LaunchPolicy{
+				AllowedEnvOverride: []string{"foo"},
+				AllowedCmdOverride: true,
+			},
+		},
+		{
+			"multiple ENV override, no CMD override",
+			map[string]string{
+				envOverride: "foo,bar",
+			},
+			LaunchPolicy{
+				AllowedEnvOverride: []string{"foo", "bar"},
+				AllowedCmdOverride: false,
+			},
+		},
+		{
+			"no ENV override, no CMD override",
+			nil,
+			LaunchPolicy{
+				AllowedEnvOverride: nil,
+				AllowedCmdOverride: false,
+			},
+		},
+		{
+			"empty string in ENV override",
+			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
 	}