cos-gpu-installer-v2: Add cos package
Change-Id: I77a029ae87affd29961acb722f06c869b72ac70d
diff --git a/src/pkg/cos/artifacts.go b/src/pkg/cos/artifacts.go
new file mode 100644
index 0000000..48b3594
--- /dev/null
+++ b/src/pkg/cos/artifacts.go
@@ -0,0 +1,143 @@
+package cos
+
+import (
+ "fmt"
+ "io/ioutil"
+ "os"
+ "path/filepath"
+
+ log "github.com/golang/glog"
+ "github.com/pkg/errors"
+
+ "pkg/utils"
+)
+
+const (
+ // TODO(mikewu): consider making GCS buckets as flags.
+ cosToolsGCS = "cos-tools"
+ internalGCS = "container-vm-image-staging"
+ chromiumOSSDKGCS = "chromiumos-sdk"
+ kernelInfo = "kernel_info"
+ kernelSrcArchive = "kernel-src.tar.gz"
+ kernelHeaders = "kernel-headers.tgz"
+ toolchainURL = "toolchain_url"
+ toolchainArchive = "toolchain.tar.xz"
+ toolchainEnv = "toolchain_env"
+ crosKernelRepo = "https://chromium.googlesource.com/chromiumos/third_party/kernel"
+)
+
+// ArtifactsDownloader defines the interface to download COS artifacts.
+type ArtifactsDownloader interface {
+ DownloadKernelSrc(destDir string) error
+ DownloadToolchainEnv(destDir string) error
+ DownloadToolchain(destDir string) error
+ DownloadKernelHeaders(destDir string) error
+ DownloadArtifact(destDir, artifact string) error
+ GetArtifact(artifact string) ([]byte, error)
+}
+
+// GCSDownloader is the struct downloading COS artifacts from GCS bucket.
+type GCSDownloader struct {
+ envReader *EnvReader
+ Internal bool
+}
+
+// DownloadKernelSrc downloads COS kernel sources to destination directory.
+func (d *GCSDownloader) DownloadKernelSrc(destDir string) error {
+ return d.DownloadArtifact(destDir, kernelSrcArchive)
+}
+
+// DownloadToolchainEnv downloads toolchain compilation environment variables to destination directory.
+func (d *GCSDownloader) DownloadToolchainEnv(destDir string) error {
+ return d.DownloadArtifact(destDir, toolchainEnv)
+}
+
+// DownloadToolchain downloads toolchain package to destination directory.
+func (d *GCSDownloader) DownloadToolchain(destDir string) error {
+ downloadURL, err := d.getToolchainURL()
+ if err != nil {
+ return errors.Wrap(err, "failed to download toolchain")
+ }
+ outputPath := filepath.Join(destDir, toolchainArchive)
+ if err := utils.DownloadContentFromURL(downloadURL, outputPath, toolchainArchive); err != nil {
+ return errors.Wrap(err, "failed to download toolchain")
+ }
+ return nil
+}
+
+// DownloadKernelHeaders downloads COS kernel headers to destination directory.
+func (d *GCSDownloader) DownloadKernelHeaders(destDir string) error {
+ return d.DownloadArtifact(destDir, kernelHeaders)
+}
+
+// GetArtifact gets an artifact from GCS buckets and returns its content.
+func (d *GCSDownloader) GetArtifact(artifactPath string) ([]byte, error) {
+ tmpDir, err := ioutil.TempDir("", "tmp")
+ if err != nil {
+ return nil, errors.Wrap(err, "failed to create temp dir")
+ }
+ defer os.RemoveAll(tmpDir)
+
+ if err = d.DownloadArtifact(tmpDir, artifactPath); err != nil {
+ return nil, errors.Wrapf(err, "failed to download artifact %s", artifactPath)
+ }
+
+ content, err := ioutil.ReadFile(filepath.Join(tmpDir, filepath.Base(artifactPath)))
+ if err != nil {
+ return nil, errors.Wrapf(err, "failed to read file %s", filepath.Join(tmpDir, artifactPath))
+ }
+
+ return content, nil
+}
+
+// DownloadArtifact downloads an artifact from GCS buckets, including public bucket and internal bucket.
+// TODO(mikewu): consider allow users to pass in GCS directories in arguments.
+func (d *GCSDownloader) DownloadArtifact(destDir, artifactPath string) error {
+ var err error
+
+ if err = utils.DownloadFromGCS(destDir, cosToolsGCS, d.artifactPublicPath(artifactPath)); err == nil {
+ return nil
+ }
+ log.Errorf("Failed to download %s from public GCS: %v", artifactPath, err)
+
+ if d.Internal {
+ if err = utils.DownloadFromGCS(destDir, internalGCS, d.artifactInternalPath(artifactPath)); err == nil {
+ return nil
+ }
+ log.Errorf("Failed to download %s from internal GCS: %v", artifactPath, err)
+ }
+
+ return errors.Errorf("failed to download %s", artifactPath)
+}
+
+func (d *GCSDownloader) artifactPublicPath(artifactPath string) string {
+ return fmt.Sprintf("%s/%s", d.envReader.BuildNumber(), artifactPath)
+}
+
+func (d *GCSDownloader) artifactInternalPath(artifactPath string) string {
+ return fmt.Sprintf("lakitu-release/R%s-%s/%s", d.envReader.Milestone(), d.envReader.BuildNumber(), artifactPath)
+}
+
+func (d *GCSDownloader) getToolchainURL() (string, error) {
+ // First, check if the toolchain path is available locally
+ tcPath := d.envReader.ToolchainPath()
+ if tcPath != "" {
+ log.V(2).Info("Found toolchain path file locally")
+ return fmt.Sprintf("https://storage.googleapis.com/%s/%s", chromiumOSSDKGCS, tcPath), nil
+ }
+
+ // Next, check if the toolchain path is available in GCS.
+ tmpDir, err := ioutil.TempDir("", "temp")
+ if err != nil {
+ return "", errors.Wrap(err, "failed to create tmp dir")
+ }
+ defer os.RemoveAll(tmpDir)
+ if err := d.DownloadArtifact(tmpDir, toolchainURL); err != nil {
+ return "", err
+ }
+ toolchainURLContent, err := ioutil.ReadFile(filepath.Join(tmpDir, toolchainURL))
+ if err != nil {
+ return "", errors.Wrap(err, "failed to read toolchain URL file")
+ }
+ return string(toolchainURLContent), nil
+}
diff --git a/src/pkg/cos/cos.go b/src/pkg/cos/cos.go
new file mode 100644
index 0000000..549dd57
--- /dev/null
+++ b/src/pkg/cos/cos.go
@@ -0,0 +1,280 @@
+// Package cos provides functionality to read and configure system configs that are specific to COS images.
+package cos
+
+import (
+ "fmt"
+ "io/ioutil"
+ "os"
+ "os/exec"
+ "path/filepath"
+ "strings"
+ "syscall"
+
+ "pkg/utils"
+
+ log "github.com/golang/glog"
+ "github.com/pkg/errors"
+)
+
+const (
+ espPartition = "/dev/sda12"
+ utsFilepath = "include/generated/utsrelease.h"
+)
+
+var (
+ execCommand = exec.Command
+)
+
+// DisableKernelModuleLocking disables kernel modules signing enforcement and loadpin so that unsigned kernel modules
+// can be loaded to COS kernel.
+func DisableKernelModuleLocking() error {
+ log.Info("Checking if third party kernel modules can be installed")
+
+ mountDir, err := ioutil.TempDir("", "mountdir")
+ if err != nil {
+ return errors.Wrap(err, "failed to create mount dir")
+ }
+
+ if err := syscall.Mount(espPartition, mountDir, "vfat", 0, ""); err != nil {
+ return errors.Wrap(err, "failed to mount path")
+ }
+
+ grubCfgPath := filepath.Join(mountDir, "esp/efi/boot/grub.cfg")
+ grubCfg, err := ioutil.ReadFile(grubCfgPath)
+ if err != nil {
+ return errors.Wrapf(err, "failed to read grub config from %s", grubCfgPath)
+ }
+
+ grubCfgStr := string(grubCfg)
+ needReboot := false
+ for _, kernelOption := range []string{
+ "module.sig_enforce",
+ "loadpin.enforce",
+ "loadpin.enabled",
+ } {
+ if newGrubCfgStr, needRebootOption := disableKernelOptionFromGrubCfg(kernelOption, grubCfgStr); needRebootOption {
+ needReboot = true
+ grubCfgStr = newGrubCfgStr
+ }
+ }
+
+ if needReboot {
+ log.Info("Modifying grub config to disable module locking.")
+ if err := os.Rename(grubCfgPath, grubCfgPath+".orig"); err != nil {
+ return errors.Wrapf(err, "failed to rename file %s", grubCfgPath)
+ }
+ if err := ioutil.WriteFile(grubCfgPath, []byte(grubCfgStr), 0644); err != nil {
+ return errors.Wrapf(err, "failed to write to file %s", grubCfgPath)
+ }
+ } else {
+ log.Info("Module locking has been disabled.")
+ }
+
+ syscall.Sync()
+ if err := syscall.Unmount(mountDir, 0); err != nil {
+ return err
+ }
+
+ if needReboot {
+ log.Warning("Rebooting")
+ if err := syscall.Reboot(syscall.LINUX_REBOOT_CMD_RESTART); err != nil {
+ return errors.Wrap(err, "failed to reboot")
+ }
+ }
+ return nil
+}
+
+// SetCompilationEnv sets compilation environment variables (e.g. CC, CXX) for third-party kernel module compilation.
+// TODO(mikewu): pass environment variables to the *exec.Cmd that runs the installer.
+func SetCompilationEnv(downloader ArtifactsDownloader) error {
+ log.Info("Downloading compilation environment variables")
+
+ compilationEnvs := make(map[string]string)
+
+ if err := downloader.DownloadToolchainEnv(os.TempDir()); err != nil {
+ // Required to support COS builds not having toolchain_env file
+ log.Info("Using default compilation environment variables")
+ compilationEnvs["CC"] = "x86_64-cros-linux-gnu-gcc"
+ compilationEnvs["CXX"] = "x86_64-cros-linux-gnu-g++"
+ } else {
+ if compilationEnvs, err = utils.LoadEnvFromFile(os.TempDir(), toolchainEnv); err != nil {
+ return errors.Wrap(err, "failed to parse toolchain_env file")
+ }
+ }
+
+ log.Info("Setting compilation environment variables")
+ for key, value := range compilationEnvs {
+ log.Infof("%s=%s", key, value)
+ os.Setenv(key, value)
+ }
+ return nil
+}
+
+// InstallCrossToolchain installs COS toolchain to destination directory.
+func InstallCrossToolchain(downloader ArtifactsDownloader, destDir string) error {
+ log.Info("Installing the toolchain")
+
+ if err := os.MkdirAll(destDir, 0755); err != nil {
+ return errors.Wrapf(err, "failed to create dir %s", destDir)
+ }
+ if empty, _ := utils.IsDirEmpty(destDir); !empty {
+ log.Info("Found existing toolchain. Skipping download and installation")
+ return nil
+ }
+
+ if err := downloader.DownloadToolchain(destDir); err != nil {
+ return errors.Wrap(err, "failed to download toolchain")
+ }
+
+ if err := exec.Command("tar", "xf", filepath.Join(destDir, toolchainArchive), "-C", destDir).Run(); err != nil {
+ return errors.Wrap(err, "failed to extract toolchain archive tarball")
+ }
+
+ log.Info("Configuring environment variables for cross-compilation")
+ os.Setenv("PATH", fmt.Sprintf("%s/bin:%s", destDir, os.Getenv("PATH")))
+ os.Setenv("SYSROOT", filepath.Join(destDir, "usr/x86_64-cros-linux-gnu"))
+ return nil
+}
+
+// InstallKernelSrcPkg installs COS kernel source package to destination directory.
+func InstallKernelSrcPkg(downloader ArtifactsDownloader, destDir string) error {
+ log.Info("Installing the kernel source package")
+
+ if err := downloadKernelSrc(downloader, destDir); err != nil {
+ return errors.Wrap(err, "failed to download kernel source")
+ }
+
+ if err := configureKernel(destDir); err != nil {
+ return errors.Wrap(err, "failed to configure kernel source")
+ }
+
+ if err := correctKernelMagicVersionIfNeeded(destDir); err != nil {
+ return errors.Wrap(err, "failed to run correctKernelMagicVersionIfNeeded")
+ }
+
+ return nil
+}
+
+// InstallKernelHeaderPkg installs kernel header package to destination directory.
+func InstallKernelHeaderPkg(downloader ArtifactsDownloader, destDir string) error {
+ log.Info("Installing the kernel header package")
+
+ if err := os.MkdirAll(destDir, 0755); err != nil {
+ return errors.Wrapf(err, "failed to create dir %s", destDir)
+ }
+ if empty, _ := utils.IsDirEmpty(destDir); !empty {
+ return nil
+ }
+
+ log.Info("Kernel headers not found locally, downloading")
+ if err := downloader.DownloadKernelHeaders(destDir); err != nil {
+ return errors.Wrap(err, "failed to download kernel headers")
+ }
+ if err := exec.Command("tar", "xf", filepath.Join(destDir, kernelHeaders), "-C", destDir).Run(); err != nil {
+ return errors.Wrap(err, "failed to extract kernel header tarball")
+ }
+
+ return nil
+}
+
+// ConfigureModuleSymvers copys Module.symvers file from kernel header dir to kernel source dir.
+func ConfigureModuleSymvers(kernelHeaderDir, kernelSrcDir string) error {
+ log.Info("Configuring Module.symvers file")
+ if err := utils.CopyFile(filepath.Join(kernelHeaderDir, "Module.symvers"),
+ filepath.Join(kernelSrcDir, "Module.symvers")); err != nil {
+ return errors.Wrap(err, "failed to copy Module.symvers file")
+ }
+ return nil
+}
+
+func disableKernelOptionFromGrubCfg(kernelOption, grubCfg string) (newGrubCfg string, needReboot bool) {
+ newGrubCfg = grubCfg
+ needReboot = false
+ kernelOptionEnabled := fmt.Sprintf("%v=1", kernelOption)
+ kernelOptionDisabled := fmt.Sprintf("%v=0", kernelOption)
+
+ if strings.Contains(grubCfg, kernelOption) {
+ if strings.Contains(grubCfg, kernelOptionEnabled) {
+ newGrubCfg = strings.ReplaceAll(grubCfg, kernelOptionEnabled, kernelOptionDisabled)
+ needReboot = true
+ }
+ } else {
+ newGrubCfg = strings.ReplaceAll(grubCfg, "cros_efi", fmt.Sprintf("cros_efi %v", kernelOptionDisabled))
+ needReboot = true
+ }
+ return newGrubCfg, needReboot
+}
+
+func downloadKernelSrc(downloader ArtifactsDownloader, destDir string) error {
+ if err := os.MkdirAll(destDir, 0755); err != nil {
+ return errors.Wrapf(err, "failed to create dir %s", destDir)
+ }
+
+ if empty, _ := utils.IsDirEmpty(destDir); !empty {
+ return nil
+ }
+
+ log.Info("Kernel sources not found locally, downloading")
+ if err := downloader.DownloadKernelSrc(destDir); err != nil {
+ return errors.Wrap(err, "failed to download kernel sources")
+ }
+ if err := exec.Command("tar", "xf", filepath.Join(destDir, kernelSrcArchive), "-C", destDir).Run(); err != nil {
+ return errors.Wrap(err, "failed to extract kernel source tarball")
+ }
+
+ return nil
+}
+
+func configureKernel(kernelSrcDir string) error {
+ log.Info("Configuring kernel")
+ // TODO(mikewu): consider getting kernel configs from kernel headers.
+ kConfig, err := exec.Command("zcat", "/proc/config.gz").Output()
+ if err != nil {
+ return errors.Wrap(err, "failed to read kernel config")
+ }
+ if err := ioutil.WriteFile(filepath.Join(kernelSrcDir, ".config"), kConfig, 0644); err != nil {
+ return errors.Wrap(err, "failed to write kernel config file")
+ }
+ cmd := exec.Command("make", "olddefconfig")
+ cmd.Stdout = os.Stdout
+ cmd.Stderr = os.Stderr
+ cmd.Dir = kernelSrcDir
+ if err := cmd.Run(); err != nil {
+ return errors.Wrap(err, "failed to run `make olddefconfig`")
+ }
+
+ cmd = exec.Command("make", "modules_prepare")
+ cmd.Dir = kernelSrcDir
+ if err := cmd.Run(); err != nil {
+ return errors.Wrap(err, "failed to run `make modules_prepare`")
+ }
+
+ // COS doesn't enable module versioning, disable Module.symvers file check.
+ os.Setenv("IGNORE_MISSING_MODULE_SYMVERS", "1")
+ return nil
+}
+
+func correctKernelMagicVersionIfNeeded(kernelSrcDir string) error {
+ // Normally COS kernel release version has a "+" in the end, e.g. "4.19.102+". But
+ // the utsrelease file generated here doesn't have it, e.g. "4.19.102". Thus we need
+ // to correct the utsrelease file to make it match the real COS kernel release version.
+ utsCmd, err := execCommand("uname", "-r").Output()
+ if err != nil {
+ return errors.Wrap(err, "failed to run `uname -r`")
+ }
+ kernelVersionCmd := strings.TrimSpace(string(utsCmd))
+ utsFile, err := ioutil.ReadFile(filepath.Join(kernelSrcDir, utsFilepath))
+ if err != nil {
+ return errors.Wrap(err, "failed to read utsrelease file")
+ }
+
+ kernelVersionFile := strings.Trim(strings.Fields(string(utsFile))[2], `"`)
+ if kernelVersionCmd != kernelVersionFile {
+ newUtsFile := strings.ReplaceAll(string(utsFile), kernelVersionFile, kernelVersionCmd)
+ log.Info("Modifying kernel release version magic string in source files")
+ if err := ioutil.WriteFile(filepath.Join(kernelSrcDir, utsFilepath), []byte(newUtsFile), 0644); err != nil {
+ return errors.Wrap(err, "failed to write to utsrelease file")
+ }
+ }
+ return nil
+}
diff --git a/src/pkg/cos/cos_test.go b/src/pkg/cos/cos_test.go
new file mode 100644
index 0000000..c6eda3e
--- /dev/null
+++ b/src/pkg/cos/cos_test.go
@@ -0,0 +1,307 @@
+package cos
+
+import (
+ "fmt"
+ "io/ioutil"
+ "os"
+ "os/exec"
+ "path/filepath"
+ "strconv"
+ "strings"
+ "testing"
+
+ "pkg/utils"
+
+ log "github.com/golang/glog"
+)
+
+var (
+ mockCmdStdout string
+ mockCmdExitStatus = 0
+)
+
+func fakeExecCommand(command string, args ...string) *exec.Cmd {
+ cs := []string{"-test.run=TestHelperProcess", "--", command}
+ cs = append(cs, args...)
+ cmd := exec.Command(os.Args[0], cs...)
+ es := strconv.Itoa(mockCmdExitStatus)
+ cmd.Env = []string{"GO_WANT_HELPER_PROCESS=1",
+ "STDOUT=" + mockCmdStdout,
+ "EXIT_STATUS=" + es}
+ return cmd
+}
+
+// TestHelperProcess is not a real test. It is a helper process for faking exec.Command.
+func TestHelperProcess(t *testing.T) {
+ if os.Getenv("GO_WANT_HELPER_PROCESS") != "1" {
+ return
+ }
+ fmt.Fprintf(os.Stdout, os.Getenv("STDOUT"))
+ es, err := strconv.Atoi(os.Getenv("EXIT_STATUS"))
+ if err != nil {
+ t.Fatalf("Failed to convert EXIT_STATUS to int: %v", err)
+ }
+ os.Exit(es)
+}
+
+func TestCorrectKernelMagicVersionIfNeeded(t *testing.T) {
+ execCommand = fakeExecCommand
+ defer func() {
+ execCommand = exec.Command
+ mockCmdExitStatus = 0
+ }()
+ for _, tc := range []struct {
+ testName string
+ kernelVersionUname string
+ utsRelease string
+ expectedNewUTSRelease string
+ }{
+ {
+ "NeedHack",
+ "4.19.101+",
+ `#define UTS_RELEASE "4.19.100+"`,
+ `#define UTS_RELEASE "4.19.101+"`,
+ },
+ {
+ "NoNeedHack",
+ "4.19.101+",
+ `#define UTS_RELEASE "4.19.101+"`,
+ `#define UTS_RELEASE "4.19.101+"`,
+ },
+ } {
+
+ tmpDir, err := ioutil.TempDir("", "testing")
+ if err != nil {
+ t.Fatalf("Failed to create temp dir: %v", err)
+ }
+ defer os.RemoveAll(tmpDir)
+ utsFile := filepath.Join(tmpDir, utsFilepath)
+ if err := os.MkdirAll(filepath.Dir(utsFile), 0755); err != nil {
+ t.Fatalf("Failed to create dir: %v", err)
+ }
+ if err := ioutil.WriteFile(utsFile, []byte(tc.utsRelease), 0644); err != nil {
+ t.Fatalf("Failed to write to utsfile: %v", err)
+ }
+ mockCmdStdout = tc.kernelVersionUname
+
+ if err := correctKernelMagicVersionIfNeeded(tmpDir); err != nil {
+ t.Fatalf("Failed to run correctKernelMagicVersionIfNeeded: %v", err)
+ }
+
+ gotUTSRelease, err := ioutil.ReadFile(utsFile)
+ if err != nil {
+ t.Fatalf("Failed to read utsfile: %v", err)
+ }
+ if string(gotUTSRelease) != tc.expectedNewUTSRelease {
+ t.Errorf("%s: Unexpected newUtsRelease, want: %s, got: %s", tc.testName, tc.expectedNewUTSRelease, gotUTSRelease)
+ }
+ }
+}
+
+func TestDownloadKernelSrc(t *testing.T) {
+ tmpDir, err := ioutil.TempDir("", "testing")
+ if err != nil {
+ t.Fatalf("Failed to create temp dir: %v", err)
+ }
+ defer os.RemoveAll(tmpDir)
+
+ downloder := fakeDownloader{}
+ if err := downloadKernelSrc(&downloder, tmpDir); err != nil {
+ t.Fatalf("Failed to run downloadKernelSrc: %v", err)
+ }
+
+ if _, err := os.Stat(filepath.Join(tmpDir, "kernel-source")); err != nil {
+ t.Errorf("Failed to get kernel source file: %v", err)
+ }
+}
+
+func TestInstallKernelHeaderPkg(t *testing.T) {
+ tmpDir, err := ioutil.TempDir("", "testing")
+ if err != nil {
+ t.Fatalf("Failed to create temp dir: %v", err)
+ }
+ defer os.RemoveAll(tmpDir)
+
+ downloder := fakeDownloader{}
+ if err := InstallKernelHeaderPkg(&downloder, tmpDir); err != nil {
+ t.Fatalf("Failed to run InstallKernelHeaderPkg: %v", err)
+ }
+
+ if _, err := os.Stat(filepath.Join(tmpDir, "kernel-header")); err != nil {
+ t.Errorf("Failed to get kernel headers file: %v", err)
+ }
+}
+
+func TestSetCompilationEnv(t *testing.T) {
+ origEnvs := os.Environ()
+ defer func() {
+ os.Clearenv()
+ for _, env := range origEnvs {
+ log.Info(env)
+ fields := strings.SplitN(env, "=", 2)
+ os.Setenv(fields[0], fields[1])
+ }
+ }()
+ tmpDir, err := ioutil.TempDir("", "testing")
+ if err != nil {
+ t.Fatalf("Failed to create temp dir: %v", err)
+ }
+ defer os.RemoveAll(tmpDir)
+ os.Setenv("TMPDIR", tmpDir)
+
+ downloder := fakeDownloader{}
+ if err := SetCompilationEnv(&downloder); err != nil {
+ t.Fatalf("Failed to run SetCompilationEnv: %v", err)
+ }
+
+ for _, tc := range []struct {
+ envKey string
+ expectedEnvValue string
+ }{
+ {"CC", "x86_64-cros-linux-gnu-clang"},
+ {"CXX", "x86_64-cros-linux-gnu-clang++"},
+ } {
+ if os.Getenv(tc.envKey) != tc.expectedEnvValue {
+ t.Errorf("Unexpected env %s value: want: %s, got: %s", tc.envKey, tc.expectedEnvValue, os.Getenv(tc.envKey))
+ }
+ }
+}
+
+func TestInstallCrossToolchain(t *testing.T) {
+ origEnvs := os.Environ()
+ defer func() {
+ os.Clearenv()
+ for _, env := range origEnvs {
+ log.Info(env)
+ fields := strings.SplitN(env, "=", 2)
+ os.Setenv(fields[0], fields[1])
+ }
+ }()
+ tmpDir, err := ioutil.TempDir("", "testing")
+ if err != nil {
+ t.Fatalf("Failed to create temp dir: %v", err)
+ }
+ defer os.RemoveAll(tmpDir)
+ origPath := os.Getenv("PATH")
+
+ downloder := fakeDownloader{}
+ if err := InstallCrossToolchain(&downloder, tmpDir); err != nil {
+ t.Fatalf("Failed to run InstallCrossToolchain: %v", err)
+ }
+
+ if _, err := os.Stat(filepath.Join(tmpDir, "x86_64-cros-linux-gnu-clang")); err != nil {
+ t.Errorf("Failed to check file in toolchain: %v", err)
+ }
+ for _, tc := range []struct {
+ envKey string
+ expectedEnvValue string
+ }{
+ {"PATH", tmpDir + "/bin:" + origPath},
+ {"SYSROOT", filepath.Join(tmpDir, "usr/x86_64-cros-linux-gnu")},
+ } {
+ if os.Getenv(tc.envKey) != tc.expectedEnvValue {
+ t.Errorf("Unexpected env %s value: want: %s, got: %s", tc.envKey, tc.expectedEnvValue, os.Getenv(tc.envKey))
+ }
+ }
+}
+
+func TestDisableKernelOptionFromGrubCfg(t *testing.T) {
+ for _, tc := range []struct {
+ testName string
+ kernelOption string
+ grubCfg string
+ expectedNewGrubCfg string
+ expectedNeedReboot bool
+ }{
+ {
+ "LoadPin",
+ "loadpin.enabled",
+
+ `BOOT_IMAGE=/syslinux/vmlinuz.A init=/usr/lib/systemd/systemd boot=local rootwait ro noresume noswap ` +
+ `loglevel=7 noinitrd console=ttyS0 security=apparmor virtio_net.napi_tx=1 ` +
+ `systemd.unified_cgroup_hierarchy=false systemd.legacy_systemd_cgroup_controller=false csm.disabled=1 ` +
+ `dm_verity.error_behavior=3 dm_verity.max_bios=-1 dm_verity.dev_wait=1 i915.modeset=1 cros_efi root=/dev/dm-0 ` +
+ `"dm=1 vroot none ro 1,0 2539520 verity payload=PARTUUID=36547742-9356-EF4E-B9AD-F8DED2F6D087 ` +
+ `hashtree=PARTUUID=36547742-9356-EF4E-B9AD-F8DED2F6D087 hashstart=2539520 alg=sha256 ` +
+ `root_hexdigest=0ff80250bd97ad47a65e7cd330ab70bcf5013d7a86817dca59fcac77f0ba1a8f ` +
+ `salt=414038a6ed9b1f528c327aff4eac16ad5ca4a6699d142ae096e90374af907c34`,
+
+ `BOOT_IMAGE=/syslinux/vmlinuz.A init=/usr/lib/systemd/systemd boot=local rootwait ro noresume noswap ` +
+ `loglevel=7 noinitrd console=ttyS0 security=apparmor virtio_net.napi_tx=1 ` +
+ `systemd.unified_cgroup_hierarchy=false systemd.legacy_systemd_cgroup_controller=false csm.disabled=1 ` +
+ `dm_verity.error_behavior=3 dm_verity.max_bios=-1 dm_verity.dev_wait=1 i915.modeset=1 cros_efi loadpin.enabled=0 root=/dev/dm-0 ` +
+ `"dm=1 vroot none ro 1,0 2539520 verity payload=PARTUUID=36547742-9356-EF4E-B9AD-F8DED2F6D087 ` +
+ `hashtree=PARTUUID=36547742-9356-EF4E-B9AD-F8DED2F6D087 hashstart=2539520 alg=sha256 ` +
+ `root_hexdigest=0ff80250bd97ad47a65e7cd330ab70bcf5013d7a86817dca59fcac77f0ba1a8f ` +
+ `salt=414038a6ed9b1f528c327aff4eac16ad5ca4a6699d142ae096e90374af907c34`,
+ true,
+ },
+ {
+ "LoadPinEnabled",
+ "loadpin.enabled",
+ "cros_efi loadpin.enabled=1",
+ "cros_efi loadpin.enabled=0",
+ true,
+ },
+ {
+ "LoadPinDisabled",
+ "loadpin.enabled",
+ "cros_efi loadpin.enabled=0",
+ "cros_efi loadpin.enabled=0",
+ false,
+ },
+ } {
+ newGrubCfg, needReboot := disableKernelOptionFromGrubCfg(tc.kernelOption, tc.grubCfg)
+ if newGrubCfg != tc.expectedNewGrubCfg || needReboot != tc.expectedNeedReboot {
+ t.Errorf("%v: Unexpected output:\nexpect grubcfg: %v\ngot grubcfg: %v\nexpect needReboot: %v, got needReboot: %v",
+ tc.testName, tc.expectedNewGrubCfg, newGrubCfg, tc.expectedNeedReboot, needReboot)
+ }
+ }
+}
+
+type fakeDownloader struct {
+}
+
+func (*fakeDownloader) DownloadKernelSrc(destDir string) error {
+ var archive = map[string][]byte{
+ "kernel-source": []byte("foo"),
+ }
+ if err := utils.CreateTarFile(filepath.Join(destDir, kernelSrcArchive), archive); err != nil {
+ return fmt.Errorf("Failed to download kernel source: %v", err)
+ }
+ return nil
+}
+
+func (*fakeDownloader) DownloadToolchainEnv(destDir string) error {
+ toolchainEnvStr := `CC=x86_64-cros-linux-gnu-clang
+CXX=x86_64-cros-linux-gnu-clang++
+`
+ if err := ioutil.WriteFile(filepath.Join(destDir, toolchainEnv), []byte(toolchainEnvStr), 0644); err != nil {
+ return fmt.Errorf("Failed to download toolchain env file: %v", err)
+ }
+ return nil
+}
+
+func (*fakeDownloader) DownloadToolchain(destDir string) error {
+ var archive = map[string][]byte{
+ "x86_64-cros-linux-gnu-clang": []byte("foo"),
+ }
+ if err := utils.CreateTarFile(filepath.Join(destDir, toolchainArchive), archive); err != nil {
+ return fmt.Errorf("Failed to download toolchain archive: %v", err)
+ }
+ return nil
+}
+
+func (*fakeDownloader) DownloadKernelHeaders(destDir string) error {
+ var archive = map[string][]byte{
+ "kernel-header": []byte("bar"),
+ }
+ if err := utils.CreateTarFile(filepath.Join(destDir, kernelHeaders), archive); err != nil {
+ return fmt.Errorf("Failed to download kernel headers: %v", err)
+ }
+ return nil
+}
+
+func (*fakeDownloader) DownloadArtifact(string, string) error { return nil }
+
+func (*fakeDownloader) GetArtifact(string) ([]byte, error) { return nil, nil }
diff --git a/src/pkg/cos/env_reader.go b/src/pkg/cos/env_reader.go
new file mode 100644
index 0000000..340b516
--- /dev/null
+++ b/src/pkg/cos/env_reader.go
@@ -0,0 +1,77 @@
+package cos
+
+import (
+ "io/ioutil"
+ "path/filepath"
+ "syscall"
+
+ "pkg/utils"
+
+ "github.com/pkg/errors"
+)
+
+const (
+ osReleasePath = "/etc/os-release"
+ toolchainPathFile = "/etc/toolchain-path"
+
+ buildID = "BUILD_ID"
+ version = "VERSION"
+ kernelCommitID = "KERNEL_COMMIT_ID"
+)
+
+// EnvReader is to read system configurations of COS.
+// TODO(mikewu): rename EnvReader to a better name.
+type EnvReader struct {
+ osRelease map[string]string
+ toolchainPath string
+ uname syscall.Utsname
+}
+
+// NewEnvReader returns an instance of EnvReader.
+func NewEnvReader(hostRootPath string) (reader *EnvReader, err error) {
+ reader = &EnvReader{}
+ reader.osRelease, err = utils.LoadEnvFromFile(hostRootPath, osReleasePath)
+ if err != nil {
+ return nil, errors.Wrapf(err, "failed to read OsRelease file from %s", osReleasePath)
+ }
+
+ if toolchainPath, err := ioutil.ReadFile(filepath.Join(hostRootPath, toolchainPathFile)); err == nil {
+ reader.toolchainPath = string(toolchainPath)
+ }
+
+ if err := syscall.Uname(&reader.uname); err != nil {
+ return nil, errors.Wrap(err, "failed to get uname")
+ }
+ return reader, nil
+}
+
+// OsRelease returns configs of /etc/os-release as a map.
+func (c *EnvReader) OsRelease() map[string]string { return c.osRelease }
+
+// BuildNumber returns COS build number.
+func (c *EnvReader) BuildNumber() string { return c.osRelease[buildID] }
+
+// Milestone returns COS milestone.
+func (c *EnvReader) Milestone() string { return c.osRelease[version] }
+
+// KernelCommit returns commit hash of the COS kernel.
+func (c *EnvReader) KernelCommit() string { return c.osRelease[kernelCommitID] }
+
+// ToolchainPath returns the toolchain path of the COS version.
+// It may return an empty string if the COS version doesn't support the feature.
+func (c *EnvReader) ToolchainPath() string { return c.toolchainPath }
+
+// KernelRelease return COS kernel release, i.e. `uname -r`
+func (c *EnvReader) KernelRelease() string { return charsToString(c.uname.Release[:]) }
+
+// charsToString converts a c-style byte array (null-terminated string) to string.
+func charsToString(chars []int8) string {
+ s := make([]byte, 0, len(chars))
+ for _, ch := range chars {
+ if ch == 0 {
+ break
+ }
+ s = append(s, byte(ch))
+ }
+ return string(s)
+}
diff --git a/src/pkg/cos/env_reader_test.go b/src/pkg/cos/env_reader_test.go
new file mode 100644
index 0000000..a6ceaf3
--- /dev/null
+++ b/src/pkg/cos/env_reader_test.go
@@ -0,0 +1,94 @@
+package cos
+
+import (
+ "fmt"
+ "io/ioutil"
+ "os"
+ "path/filepath"
+ "reflect"
+ "testing"
+)
+
+func TestEnvReader(t *testing.T) {
+ testDir, err := ioutil.TempDir("", "testing")
+ if err != nil {
+ t.Fatalf("Failed to create tempdir: %v", err)
+ }
+ defer os.RemoveAll(testDir)
+
+ osReleaseString := `BUILD_ID=12688.0.0
+NAME="Container-Optimized OS"
+KERNEL_COMMIT_ID=5d8615d1e135275cbfdf9522517a3b198e7199ee
+GOOGLE_CRASH_ID=Lakitu
+VERSION_ID=80
+BUG_REPORT_URL="https://cloud.google.com/container-optimized-os/docs/resources/support-policy#contact_us"
+PRETTY_NAME="Container-Optimized OS from Google"
+VERSION=80
+GOOGLE_METRICS_PRODUCT_ID=26
+HOME_URL="https://cloud.google.com/container-optimized-os/docs"
+ID=cos`
+ if err := createConfigFile(osReleaseString, osReleasePath, testDir); err != nil {
+ t.Fatalf("Failed to create osRelease file: %v", err)
+ }
+ toolchainPathString := `2019/11/x86_64-cros-linux-gnu-2019.11.16.041937.tar.xz`
+ if err := createConfigFile(toolchainPathString, toolchainPathFile, testDir); err != nil {
+ t.Fatalf("Failed to create toolchain path file: %v", err)
+ }
+
+ envReader, err := NewEnvReader(testDir)
+ if err != nil {
+ t.Fatalf("Failed to create EnvReader: %v", err)
+ }
+
+ for _, tc := range []struct {
+ testName string
+ got interface{}
+ expect interface{}
+ }{
+ {
+ "OsRelease",
+ envReader.OsRelease(),
+ map[string]string{
+ "BUILD_ID": "12688.0.0",
+ "NAME": "Container-Optimized OS",
+ "KERNEL_COMMIT_ID": "5d8615d1e135275cbfdf9522517a3b198e7199ee",
+ "GOOGLE_CRASH_ID": "Lakitu",
+ "VERSION_ID": "80",
+ "BUG_REPORT_URL": "https://cloud.google.com/container-optimized-os/docs/resources/support-policy#contact_us",
+ "PRETTY_NAME": "Container-Optimized OS from Google",
+ "VERSION": "80",
+ "GOOGLE_METRICS_PRODUCT_ID": "26",
+ "HOME_URL": "https://cloud.google.com/container-optimized-os/docs",
+ "ID": "cos",
+ },
+ },
+ {"BuildNumber", envReader.BuildNumber(), "12688.0.0"},
+ {"Milestone", envReader.Milestone(), "80"},
+ {"Milestone", envReader.KernelCommit(), "5d8615d1e135275cbfdf9522517a3b198e7199ee"},
+ {"ToolchainPath", envReader.ToolchainPath(), "2019/11/x86_64-cros-linux-gnu-2019.11.16.041937.tar.xz"},
+ } {
+ if !reflect.DeepEqual(tc.expect, tc.got) {
+ t.Errorf("Unexpected %s,\nwant: %v\n got: %v", tc.testName, tc.testName, tc.expect)
+ }
+ }
+}
+
+func createConfigFile(configStr, configFileName, testDir string) error {
+ path := filepath.Join(testDir, configFileName)
+ if err := os.MkdirAll(filepath.Dir(path), 0744); err != nil {
+ return fmt.Errorf("Failed to create dir: %v", err)
+ }
+ configFile, err := os.Create(path)
+ if err != nil {
+ return fmt.Errorf("Failed to create file: %v", err)
+ }
+ defer configFile.Close()
+
+ if _, err = configFile.WriteString(configStr); err != nil {
+ return fmt.Errorf("Failed to write to file %s: %v", configFile.Name(), err)
+ }
+ if err = configFile.Close(); err != nil {
+ return fmt.Errorf("Failed to close file %s: %v", configFile.Name(), err)
+ }
+ return nil
+}
diff --git a/src/pkg/cos/extensions.go b/src/pkg/cos/extensions.go
new file mode 100644
index 0000000..4d77ad8
--- /dev/null
+++ b/src/pkg/cos/extensions.go
@@ -0,0 +1,92 @@
+package cos
+
+import (
+ "fmt"
+ "path/filepath"
+ "regexp"
+
+ "pkg/utils"
+
+ log "github.com/golang/glog"
+ "github.com/pkg/errors"
+)
+
+const (
+ // GPUExtension is the name of GPU extension.
+ GPUExtension = "gpu"
+)
+
+// ExtensionsDownloader is the struct downloading COS extensions from GCS bucket.
+type ExtensionsDownloader interface {
+ ListExtensions() ([]string, error)
+ ListExtensionArtifacts(extension string) ([]string, error)
+ DownloadExtensionArtifact(destDir, extension, artifact string) error
+ GetExtensionArtifact(extension, artifact string) ([]byte, error)
+}
+
+// ListExtensions lists all supported extensions.
+func (d *GCSDownloader) ListExtensions() ([]string, error) {
+ var objects []string
+ var err error
+ if objects, err = utils.ListGCSBucket(cosToolsGCS, d.artifactPublicPath("extensions")); err != nil || len(objects) == 0 {
+ log.Errorf("Failed to list extensions from public GCS: %v", err)
+ if d.Internal {
+ if objects, err = utils.ListGCSBucket(internalGCS, d.artifactInternalPath("extensions")); err != nil {
+ log.Errorf("Failed to list extensions from internal GCS: %v", err)
+ }
+ }
+ }
+ if err != nil {
+ return nil, errors.Wrap(err, "failed to list extensions")
+ }
+
+ var extensions []string
+ re := regexp.MustCompile(`extensions/(\w+)$`)
+ for _, object := range objects {
+ if match := re.FindStringSubmatch(object); match != nil {
+ extensions = append(extensions, match[1])
+ }
+ }
+ return extensions, nil
+}
+
+// ListExtensionArtifacts lists all artifacts of a given extension.
+// TODO(mikewu): make this extension specific.
+func (d *GCSDownloader) ListExtensionArtifacts(extension string) ([]string, error) {
+ var objects []string
+ var err error
+ extensionPath := filepath.Join("extensions", extension)
+ if objects, err = utils.ListGCSBucket(cosToolsGCS, d.artifactPublicPath(extensionPath)); err != nil || len(objects) == 0 {
+ log.Errorf("Failed to list extension artifacts from public GCS: %v", err)
+ // TODO(mikewu): use flags to specify GCS directories.
+ if d.Internal {
+ if objects, err = utils.ListGCSBucket(internalGCS, d.artifactInternalPath(extensionPath)); err != nil {
+ log.Errorf("Failed to list extension artifacts from internal GCS: %v", err)
+ }
+ }
+ }
+ if err != nil {
+ return nil, errors.Wrap(err, "failed to list extensions")
+ }
+
+ var artifacts []string
+ re := regexp.MustCompile(fmt.Sprintf(`extensions/%s/(.+)$`, extension))
+ for _, object := range objects {
+ if match := re.FindStringSubmatch(object); match != nil {
+ artifacts = append(artifacts, match[1])
+ }
+ }
+ return artifacts, nil
+}
+
+// DownloadExtensionArtifact downloads an artifact of the given extension.
+func (d *GCSDownloader) DownloadExtensionArtifact(destDir, extension, artifact string) error {
+ artifactPath := filepath.Join("extensions", extension, artifact)
+ return d.DownloadArtifact(destDir, artifactPath)
+}
+
+// GetExtensionArtifact reads the content of an artifact of the given extension.
+func (d *GCSDownloader) GetExtensionArtifact(extension, artifact string) ([]byte, error) {
+ artifactPath := filepath.Join("extensions", extension, artifact)
+ return d.GetArtifact(artifactPath)
+}