Add experiment support (google#352)

Co-authored-by: Joshua Krstic <[email protected]>

launcher/container_runner.go

@@ -500,6 +500,8 @@ func (r *ContainerRunner) Run(ctx context.Context) error {
 		return fmt.Errorf("failed to fetch and write OIDC token: %v", err)
 	}
 
+	r.logger.Printf("EnableTestFeatureForImage is set to %v\n", r.launchSpec.Experiments.EnableTestFeatureForImage)
+
 	var streamOpt cio.Opt
 	switch r.launchSpec.LogRedirect {
 	case spec.Nowhere:

launcher/image/cloudbuild.yaml

@@ -15,6 +15,13 @@ steps:
       - |
         cd launcher/launcher
         CGO_ENABLED=0 go build -o ../image/launcher
+  - name: 'gcr.io/cloud-builders/gcloud'
+    id: DownloadExpBinary
+    entrypoint: 'gcloud'
+    args: ['storage',
+           'cp',
+           'gs://confidential-space-images_third-party/confidential_space_experiments',
+           './launcher/image/confidential_space_experiments']
   - name: 'gcr.io/cos-cloud/cos-customizer'
     args: ['start-image-build',
            '-build-context=launcher/image',

launcher/image/preload.sh

@@ -2,11 +2,18 @@
 
 readonly OEM_PATH='/usr/share/oem'
 readonly CS_PATH="${OEM_PATH}/confidential_space"
+readonly EXPERIMENTS_BINARY="confidential_space_experiments"
 
 copy_launcher() {
   cp launcher "${CS_PATH}/cs_container_launcher"
 }
 
+copy_experiment_client() {
+  # DownloadExpBinary creates the file at EXPERIMENTS_BINARY.
+  cp $EXPERIMENTS_BINARY "${CS_PATH}/${EXPERIMENTS_BINARY}"
+  chmod +x "${CS_PATH}/${EXPERIMENTS_BINARY}"
+}
+
 setup_launcher_systemd_unit() {
   cp container-runner.service "${CS_PATH}/container-runner.service"
   cp exit_script.sh "${CS_PATH}/exit_script.sh"
@@ -90,6 +97,8 @@ main() {
 
   # Install container launcher entrypoint.
   configure_entrypoint "entrypoint.sh"
+  # Install experiment client.
+  copy_experiment_client
   # Install container launcher.
   copy_launcher
   setup_launcher_systemd_unit

launcher/internal/experiments/experiments.go

@@ -0,0 +1,42 @@
+// Package experiments contains functionalities to retrieve synced experiments
+package experiments
+
+import (
+	"encoding/json"
+	"fmt"
+	"os"
+)
+
+// Experiments contains the experiments flags this version of the launcher expects to receive.
+// Failure to unmarshal the experiment JSON data will result in an empty object being returned
+// to treat experiment flags as their default value. The error should still be checked.
+type Experiments struct {
+	EnableTestFeatureForImage  bool
+	EnableSignedContainerImage bool
+}
+
+// New takes a filepath, opens the file, and calls ReadJsonInput with the contents
+// of the file.
+// If the file cannot be opened, the experiments map is set to an empty map.
+func New(fpath string) (Experiments, error) {
+	f, err := os.ReadFile(fpath)
+	if err != nil {
+		// Return default values on failure.
+		return Experiments{}, err
+	}
+
+	r, err := readJSONInput(f)
+
+	return r, err
+}
+
+// ReadJSONInput  takes a reader and unmarshals the contents into the experiments map.
+// If the unmarsahlling fails, the experiments map is set to an empty map.
+func readJSONInput(b []byte) (Experiments, error) {
+	var experiments Experiments
+	if err := json.Unmarshal(b, &experiments); err != nil {
+		// Return default values on failure.
+		return Experiments{}, fmt.Errorf("failed to unmarshal json: %w", err)
+	}
+	return experiments, nil
+}

launcher/internal/experiments/experiments_test.go

@@ -0,0 +1,52 @@
+package experiments
+
+import (
+	"testing"
+)
+
+func TestExperiments(t *testing.T) {
+	tests := []struct {
+		input string
+	}{
+		{input: "{\"EnableTestFeatureForImage\":true,\"EnableSignedContainerImage\":true}"},
+		{input: "{\"EnableTestFeatureForImage\":true,\"EnableSignedContainerImage\":true,\"FloatFeature\":-5.6,\"OtherTestFeatureForImage\":false}"},
+	}
+
+	for i, test := range tests {
+		e, err := readJSONInput([]byte(test.input))
+
+		if err != nil {
+			t.Errorf("testcase %d: failed to create experiments object: %v", i, err)
+		}
+
+		if e.EnableTestFeatureForImage == false {
+			t.Errorf("testcase %d: expected EnableTestFeatureForImage to be true, got false", i)
+		}
+
+		if e.EnableSignedContainerImage == false {
+			t.Errorf("testcase %d: expected EnableSignedContainerImage to be true, got false", i)
+		}
+	}
+}
+
+func TestExperimentsBadJson(t *testing.T) {
+	tests := []struct {
+		input string
+	}{
+		{input: "{\"EnableTestFeatureForImage\":true,\"EnableSignedContainerImage\":true"},
+		{input: "{}"},
+		{input: ""},
+	}
+
+	for i, test := range tests {
+		e, _ := readJSONInput([]byte(test.input))
+
+		if e.EnableTestFeatureForImage == true {
+			t.Errorf("testcase %d: expected EnableTestFeatureForImage to be false, got true", i)
+		}
+
+		if e.EnableSignedContainerImage == true {
+			t.Errorf("testcase %d: expected EnableSignedContainerImage to be false, got true", i)
+		}
+	}
+}

launcher/launcher/main.go

@@ -9,6 +9,7 @@ import (
 	"log"
 	"os"
 	"os/exec"
+	"path"
 	"regexp"
 	"strings"
 
@@ -18,6 +19,8 @@ import (
 	"github.com/containerd/containerd/namespaces"
 	"github.com/google/go-tpm-tools/client"
 	"github.com/google/go-tpm-tools/launcher"
+	"github.com/google/go-tpm-tools/launcher/internal/experiments"
+	"github.com/google/go-tpm-tools/launcher/launcherfile"
 	"github.com/google/go-tpm-tools/launcher/spec"
 	"github.com/google/go-tpm/legacy/tpm2"
 )
@@ -28,6 +31,10 @@ const (
 	// panic() returns 2
 	rebootRC = 3 // reboot
 	holdRC   = 4 // hold
+	// experimentDataFile defines where the experiment sync output data is expected to be.
+	experimentDataFile = "experiment_data"
+	// binaryPath contains the path to the experiments binary.
+	binaryPath = "/usr/share/oem/confidential_space/confidential_space_experiments"
 )
 
 var rcMessage = map[int]string{
@@ -72,6 +79,21 @@ func main() {
 		return
 	}
 
+	experimentsFile := path.Join(launcherfile.HostTmpPath, experimentDataFile)
+
+	args := fmt.Sprintf("-output=%s", experimentsFile)
+	err = exec.Command(binaryPath, args).Run()
+	if err != nil {
+		logger.Printf("failure during experiment sync: %v", err)
+	}
+
+	e, err := experiments.New(experimentsFile)
+	if err != nil {
+		logger.Printf("failed to read experiment file %v\n", err)
+		// do not fail if experiment retrieval fails
+	}
+	launchSpec.Experiments = e
+
 	defer func() {
 		// Catch panic to attempt to output to Cloud Logging.
 		if r := recover(); r != nil {

launcher/spec/launch_spec.go

@@ -9,6 +9,7 @@ import (
 	"strings"
 
 	"cloud.google.com/go/compute/metadata"
+	"github.com/google/go-tpm-tools/launcher/internal/experiments"
 )
 
 // RestartPolicy is the enum for the container restart policy.
@@ -92,6 +93,7 @@ type LaunchSpec struct {
 	Region                     string
 	Hardened                   bool
 	LogRedirect                LogRedirectLocation
+	Experiments                experiments.Experiments
 }
 
 // UnmarshalJSON unmarshals an instance attributes list in JSON format from the metadata